def METHOD_NAME(self):
    # Ensure custom scalar reward with batch_size greater than 1 raises
    # ValueError
    obs_spec = BoundedTensorSpec((2, 3), torch.int32, -10, 10)
    action_spec = BoundedTensorSpec((1, ), torch.int64)
    env = RandomAlfEnvironment(
        obs_spec,
        action_spec,
        reward_fn=lambda *_: np.array([1.0]),
        batch_size=5)
    env.reset()
    env._done = False
    action = np.array(0, dtype=np.int64)
    with self.assertRaises(ValueError):
        env.step(action)