def policy(self, trajectory, temperature=1):
    """Chooses an action to play after a trajectory."""
    tr_slice = trajectory.suffix(self._max_slice_length)
    trajectory_np = tr_slice.to_np(timestep_to_np=self.task.timestep_to_np)
    # Add batch dimension to trajectory_np and run the model.
    obs = trajectory_np.observation[None, ...]
    values = self._run_value_model(obs, use_eval_model=False)
    # We insisit that values and observations have the shape
    # (batch, length, ...), where the length is the number of subsequent
    # observations on a given trajectory
    assert values.shape[:1] == obs.shape[:1]
    # We select the last element in the batch and the value
    # related to the last (current) observation
    values = values[0, -1, :]
    # temperature == 0 is used in another place in order to trigger eval
    if np.random.random_sample() < self._exploration_rate(self._epoch) and \
        temperature == 1:
      sample = np.array(self.task.action_space.sample())
    else:
      # this is our way of doing the argmax
      sample = jnp.argmax(values)
    result = (sample, values)
    if fastmath.backend_name() == 'jax':
      result = fastmath.nested_map(lambda x: x.copy(), result)
    return result
Chooses an action to play after a trajectory.
trax/rl/training.py