code: def rollout(
      env,
      agent,
      *,
      max_episode_length=np.inf,
      animated=False,
      speedup=1,
      deterministic=False,
  ):
      """Sample a single episode of the agent in the environment.
  
      Args:
          agent (Policy): Agent 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.
          speedup (float): Factor by which to decrease the wait time between
              rendered 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 (T + 1, S^*) (the
                  unflattened state space of the current environment).
              * actions(np.array): Non-flattened array of actions. Should have
                  shape (T, S^*) (the unflattened action space of the current
                  environment).
              * rewards(np.array): Array of rewards of shape (T,) (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.
              * episode_infos(Dict[str, np.array]): Dictionary of stacked,
                  non-flattened `episode_info` arrays.
              * dones(np.array): Array of termination signals.
  
      """
      del speedup
      env_steps = []
      agent_infos = []
      observations = []
      episode_infos = []
      last_obs, episode_info = env.reset()
      agent.reset()
      episode_length = 0
      if animated:
          env.visualize()
      while episode_length < (max_episode_length or np.inf):
          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_infos.append(episode_info)
          episode_length += 1
          if es.last:
              break
          last_obs = es.observation
  
      return dict(
          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]
          ),
          episode_infos=stack_tensor_dict_list(episode_infos),
          dones=np.array([es.terminal for es in env_steps]),
      )
  
label_desc: Use image node_id 1 for calling a global function with example usage: rollout(env, agent) and returns: dict