def play(self, policy, max_steps=None, only_eval=False):
    """Play an episode in env taking actions according to the given policy."""
    if max_steps is None:
      max_steps = self._max_steps
    if only_eval:
      cur_trajectory = play(
          self._eval_env, policy, self._dm_suite,
          # Only step up to the time limit.
          max_steps=min(max_steps, self._time_limit),
          # Always reset at the beginning of an eval episode.
          last_observation=None,
      )
    else:
      cur_trajectory = play(
          self._env, policy, self._dm_suite,
          # Only step up to the time limit, used up by all slices of the same
          # trajectory.
          max_steps=min(max_steps, self._n_steps_left),
          # Pass the environmnent state between play() calls, so one episode can
          # span multiple training epochs.
          # NOTE: Cutting episodes between epochs may hurt e.g. with
          # Transformers.
          # TODO(pkozakowski): Join slices together if this becomes a problem.
          last_observation=self._last_observation,
      )
      # Update the number of steps left to reach the time limit.
      # NOTE: This should really be done as an env wrapper.
      # TODO(pkozakowski): Do that once we wrap the DM envs in Gym interface.
      # The initial reset doesn't count.
      self._n_steps_left -= len(cur_trajectory) - 1
      assert self._n_steps_left >= 0
      if self._n_steps_left == 0:
        cur_trajectory.done = True
      # Pass the last observation between trajectory slices.
      if cur_trajectory.done:
        self._last_observation = None
        # Reset the time limit.
        self._n_steps_left = self._time_limit
      else:
        self._last_observation = cur_trajectory.last_observation

    cur_trajectory.calculate_returns(self._gamma)
    return cur_trajectory
Play an episode in env taking actions according to the given policy.
trax/rl/task.py