fe3101f5a5b02e9e7b4ddb260cda15b05cfb6dd7,tf_agents/policies/random_tf_policy_test.py,RandomTFPolicyTest,testInfoSpec,#RandomTFPolicyTest#Any#,188

Before Change

        "PolicyInfo",
        ("log_probability", "predicted_rewards"))
    // Set default empty tuple for all fields.
    PolicyInfo.__new__.__defaults__ = ((),) * len(PolicyInfo._fields)

    action_spec = [
        tensor_spec.BoundedTensorSpec((2, 3), dtype, -10, 10),

After Change

        tensor_spec.BoundedTensorSpec((2, 3), dtype, -10, 10),
        tensor_spec.BoundedTensorSpec((1, 2), dtype, -10, 10)
    ]
    info_spec = [
        tensor_spec.TensorSpec([1], dtype=tf.float32, name="loc"),
        tensor_spec.TensorSpec([1], dtype=tf.float32, name="scale")
    ]
    time_step_spec, time_step = self.create_time_step()
    policy = random_tf_policy.RandomTFPolicy(
        time_step_spec=time_step_spec,
        action_spec=action_spec,
        info_spec=info_spec)

    // Test without batch
    action_step = policy.action(time_step)
    tf.nest.assert_same_structure(action_spec, action_step.action)
    self.assertEqual((2, 3,), action_step.action[0].shape)
    self.assertEqual((1, 2,), action_step.action[1].shape)
    tf.nest.assert_same_structure(info_spec, action_step.info)
    self.assertEqual((1,), action_step.info[0].shape)
    self.assertEqual((1,), action_step.info[1].shape)

    // Test with batch, we should see the additional outer batch dim for both
    // `action` and `info`.
    batch_size = 2
    batched_time_step = self.create_batch(time_step, batch_size)
    batched_action_step = policy.action(batched_time_step)
    tf.nest.assert_same_structure(action_spec, batched_action_step.action)
    self.assertEqual((batch_size, 2, 3,), batched_action_step.action[0].shape)
    self.assertEqual((batch_size, 1, 2,), batched_action_step.action[1].shape)
    tf.nest.assert_same_structure(info_spec, batched_action_step.info)

Instances