class TestBug(tf.test.TestCase):

    def testBugWrapper(self):
        obs_spec = array_spec.BoundedArraySpec((2, 3), np.int32, -10, 10)
        action_spec = array_spec.BoundedArraySpec((1,), np.int32, -10, 10)
        nested_env = random_py_environment.RandomPyEnvironment(
            obs_spec,
            action_spec,
            reward_fn=lambda *_: np.array([1.0]),
            batch_size=2)
        env = wrappers.ActionDiscretizeWrapper(nested_env, num_actions=100)
        self.assertNotEqual(env.batched, nested_env.batched)  # Don't raise when it should
        self.assertNotEqual(nested_env.batch_size, env.batch_size)  # Don't raise when it should

class PyEnvironmentBaseWrapper(py_environment.PyEnvironment):
  """PyEnvironment wrapper forwards calls to the given environment."""

  def __getattr__(self, name):
    """Forward all other calls to the base environment."""
    return getattr(self._env, name)

  @property
  def batched(self):
    return self._env.bacthed

  @property
  def batch_size(self):
    return self._env.batch_size