def compute_optimal_action_with_environment_dynamics(
    observation, environment_dynamics):
  """Computes the optimal action using the environment dynamics.

  Args:
    observation: a (possibly batched) observation.
    environment_dynamics: environment dynamics object (an instance of
      `non_stationary_stochastic_environment.EnvironmentDynamics`)

  Returns:
    The optimal action.
  """
  return environment_dynamics.compute_optimal_action(observation)