def play(env, policy, dm_suite=False, max_steps=None, last_observation=None):
  """Play an episode in env taking actions according to the given policy.

  Environment is first reset and an from then on, a game proceeds. At each
  step, the policy is asked to choose an action and the environment moves
  forward. A Trajectory is created in that way and returns when the episode
  finished, which is either when env returns `done` or max_steps is reached.

  Args:
    env: the environment to play in, conforming to gym.Env or
      DeepMind suite interfaces.
    policy: a function taking a Trajectory and returning a pair consisting
      of an action (int or float) and the confidence in that action (float,
      defined as the log of the probability of taking that action).
    dm_suite: whether we are using the DeepMind suite or the gym interface
    max_steps: for how many steps to play.
    last_observation: last observation from a previous trajectory slice, used to
      begin a new one. Controls whether we reset the environment at the
      beginning - if `None`, resets the env and starts the slice from the
      observation got from reset().

  Returns:
    a completed trajectory slice that was just played.
  """
  done = False
  cur_step = 0
  if last_observation is None:
    # TODO(pkozakowski): Make a Gym wrapper over DM envs to get rid of branches
    # like that.
    last_observation = env.reset().observation if dm_suite else env.reset()
  cur_trajectory = Trajectory(last_observation)
  while not done and (max_steps is None or cur_step < max_steps):
    action, dist_inputs = policy(cur_trajectory)
    action = np.asarray(action)
    step = env.step(action)
    if dm_suite:
      (observation, reward, done) = (
          step.observation, step.reward, step.step_type.last()
      )
      info = {}
    else:
      (observation, reward, done, info) = step

    # Make an EnvInfo out of the supported keys in the info dict.
    env_info = EnvInfo(**{
        key: value for (key, value) in info.items()
        if key in EnvInfo._fields
    })
    cur_trajectory.extend(
        action=action,
        dist_inputs=dist_inputs,
        reward=reward,
        done=done,
        new_observation=observation,
        env_info=env_info,
    )
    cur_step += 1
  return cur_trajectory
Play an episode in env taking actions according to the given policy.

  Environment is first reset and an from then on, a game proceeds. At each
  step, the policy is asked to choose an action and the environment moves
  forward. A Trajectory is created in that way and returns when the episode
  finished, which is either when env returns `done` or max_steps is reached.

  Args:
    env: the environment to play in, conforming to gym.Env or
      DeepMind suite interfaces.
    policy: a function taking a Trajectory and returning a pair consisting
      of an action (int or float) and the confidence in that action (float,
      defined as the log of the probability of taking that action).
    dm_suite: whether we are using the DeepMind suite or the gym interface
    max_steps: for how many steps to play.
    last_observation: last observation from a previous trajectory slice, used to
      begin a new one. Controls whether we reset the environment at the
      beginning - if `None`, resets the env and starts the slice from the
      observation got from reset().

  Returns:
    a completed trajectory slice that was just played.
  
trax/rl/task.py