def testWrongPolicyInfoType(self): dims = (10, 1) log_probability = tf.fill(dims, value=-0.5) info = policy_step.PolicyInfo(log_probability=log_probability) input_tensor = tf.fill(dims, value=_GREEDY) result = policy_utilities.set_bandit_policy_type(info, input_tensor) self.assertNotIsInstance(result, policy_utilities.PolicyInfo) self.assertAllEqual(info.log_probability, result.log_probability)
def _action(self, time_step, policy_state, seed): observation_and_action_constraint_splitter = ( self.observation_and_action_constraint_splitter) if observation_and_action_constraint_splitter is not None: _, mask = observation_and_action_constraint_splitter( time_step.observation) zero_logits = tf.cast(tf.zeros_like(mask), tf.float32) masked_categorical = masked.MaskedCategorical(zero_logits, mask) #Modified to accomodate scalar action spaces #action_ = tf.cast(masked_categorical.sample() + self.action_spec.minimum, # self.action_spec.dtype) action_ = tf.reshape( tf.cast(masked_categorical.sample() + self.action_spec.minimum, self.action_spec.dtype), [1]) # If the action spec says each action should be shaped (1,), add another # dimension so the final shape is (B, 1) rather than (B,). if self.action_spec.shape.rank == 1: action_ = tf.expand_dims(action_, axis=-1) else: outer_dims = nest_utils.get_outer_shape(time_step, self._time_step_spec) action_ = tensor_spec.sample_spec_nest(self._action_spec, seed=seed, outer_dims=outer_dims) # TODO(b/78181147): Investigate why this control dependency is required. if time_step is not None: with tf.control_dependencies(tf.nest.flatten(time_step)): action_ = tf.nest.map_structure(tf.identity, action_) step = policy_step.PolicyStep(action_, policy_state) if self.emit_log_probability: if observation_and_action_constraint_splitter is not None: log_probability = masked_categorical.log_prob( action_ - self.action_spec.minimum) else: action_probability = tf.nest.map_structure( _uniform_probability, self._action_spec) log_probability = tf.nest.map_structure( tf.math.log, action_probability) info = policy_step.PolicyInfo(log_probability=log_probability) return step._replace(info=info) return step
def _action(self, time_step, policy_state, seed): outer_dims = nest_utils.get_outer_shape(time_step, self._time_step_spec) action_ = tensor_spec.sample_spec_nest( self._action_spec, seed=seed, outer_dims=outer_dims) # TODO(b/78181147): Investigate why this control dependency is required. if time_step is not None: with tf.control_dependencies(tf.nest.flatten(time_step)): action_ = tf.nest.map_structure(tf.identity, action_) step = policy_step.PolicyStep(action_, policy_state) if self.emit_log_probability: action_probability = tf.nest.map_structure(_uniform_probability, self._action_spec) log_probability = tf.nest.map_structure(tf.math.log, action_probability) info = policy_step.PolicyInfo(log_probability=log_probability) return step._replace(info=info) return step