code: 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
  
imports: import collections
  import os
  import gin
  import gym
  import numpy
  from trax import fastmath
  from trax.rl import advantages
  from trax.supervised import training
label_desc: Use image node_id 1 for calling a global function with example usage: play(env, policy, dm_suite, max_steps, last_observation) and returns: cur_trajectory