def test_action_mapping_1d(self):
    """Testing discretization with a mock environment.

    In the mock call we get access to the argument of the
    SimpleContinuousActionsEnv.step method which we check against
    precomputed values of continuous actions.
    """
    num_actions = 4

    with unittest.mock.patch.object(
        gym_utils.RenderedEnv, "step", autospec=True) as mock_step_method:
      env = gym_utils.RenderedEnv(SimpleContinuousActionsEnv(dimensions=1))
      expected_continuous_actions = np.linspace(
          np.min(env.action_space.low),
          np.min(env.action_space.high),
          num=num_actions).flatten()

      env_discrete = gym_utils.ActionDiscretizeWrapper(env, num_actions)
      for discrete_action in range(num_actions):
        env_discrete.step([discrete_action])
        mock_step_method.assert_called_with(
            unittest.mock.ANY,
            expected_continuous_actions[discrete_action])
Testing discretization with a mock environment.

    In the mock call we get access to the argument of the
    SimpleContinuousActionsEnv.step method which we check against
    precomputed values of continuous actions.
    
tensor2tensor/rl/gym_utils_test.py