code: def rollout(
      env,
      agent,
      *,
      max_episode_length=np.inf,
      animated=False,
      pause_per_frame=None,
      deterministic=False,
  ):
      """Sample a single episode of the agent in the environment.
  
      Args:
          agent (Policy): Policy used to select actions.
          env (Environment): Environment to perform actions in.
          max_episode_length (int): If the episode reaches this many timesteps,
              it is truncated.
          animated (bool): If true, render the environment after each step.
          pause_per_frame (float): Time to sleep between steps. Only relevant if
              animated == true.
          deterministic (bool): If true, use the mean action returned by the
              stochastic policy instead of sampling from the returned action
              distribution.
  
      Returns:
          dict[str, np.ndarray or dict]: Dictionary, with keys:
              * observations(np.array): Flattened array of observations.
                  There should be one more of these than actions. Note that
                  observations[i] (for i < len(observations) - 1) was used by the
                  agent to choose actions[i]. Should have shape
                  :math:`(T + 1, S^*)`, i.e. the unflattened observation space of
                      the current environment.
              * actions(np.array): Non-flattened array of actions. Should have
                  shape :math:`(T, S^*)`, i.e. the unflattened action space of
                  the current environment.
              * rewards(np.array): Array of rewards of shape :math:`(T,)`, i.e. a
                  1D array of length timesteps.
              * agent_infos(Dict[str, np.array]): Dictionary of stacked,
                  non-flattened `agent_info` arrays.
              * env_infos(Dict[str, np.array]): Dictionary of stacked,
                  non-flattened `env_info` arrays.
              * dones(np.array): Array of termination signals.
  
      """
      env_steps = []
      agent_infos = []
      observations = []
      last_obs, episode_infos = env.reset()
      agent.reset()
      episode_length = 0
      if animated:
          env.visualize()
      while episode_length < (max_episode_length or np.inf):
          if pause_per_frame is not None:
              time.sleep(pause_per_frame)
          a, agent_info = agent.get_action(last_obs)
          if deterministic and "mean" in agent_info:
              a = agent_info["mean"]
          es = env.step(a)
          env_steps.append(es)
          observations.append(last_obs)
          agent_infos.append(agent_info)
          episode_length += 1
          if es.last:
              break
          last_obs = es.observation
  
      return dict(
          episode_infos=episode_infos,
          observations=np.array(observations),
          actions=np.array([es.action for es in env_steps]),
          rewards=np.array([es.reward for es in env_steps]),
          agent_infos=stack_tensor_dict_list(agent_infos),
          env_infos=stack_tensor_dict_list(
              [es.env_info for es in env_steps]
          ),
          dones=np.array([es.terminal for es in env_steps]),
      )
  
file_path: garage/src/garage/_functions.py
imports: from collections import defaultdict
  import time
  import click
  from dowel import tabular
  import numpy
  from garage import EpisodeBatch, StepType
  from garage.np import discount_cumsum, stack_tensor_dict_list
label_desc: Use image node_id 2 for calling a global function with example usage: rollout(env, agent) and returns: dict