code: def rollout(
      env: MazeEnv,
      agent: Policy,
      n_max_steps: int,
      render: bool = False,
      log_dir: str = ".",
  ) -> float:
      """
      Rolls out environment and collects total reward for one epoch.
  
      :param env: Environment to roll out.
      :param agent: Agent with .compute_action(obs).
      :param n_max_steps: Max. number of steps to take.
      :param render: Whether to render the environment.
      :param log_dir: Logging directory.
      :return: Cumulative reward.
      """
  
      cumulative_reward = 0
      i = 0
  
      # Suppress output.
      with open(os.devnull, "w") as devnull:
          with contextlib.redirect_stdout(devnull):
              # Initialize MonitoringSetup for rendering support.
              with MonitoringSetup(
                  env, log_dir=log_dir
              ) as monitored_env:
                  obs = monitored_env.reset()
                  done = False
  
                  while not done and i < n_max_steps:
                      action = agent.compute_action(
                          obs,
                          maze_state=monitored_env.get_maze_state(),
                          actor_id=monitored_env.actor_id(),
                      )
                      obs, reward, done, _ = monitored_env.step(action)
                      cumulative_reward += reward
  
                      if render:
                          monitored_env.render()
  
                      i += 1
  
      env.close()
  
      return cumulative_reward
  
file_path: maze/maze/utils/notebooks.py
imports: import contextlib
  import os
  import subprocess
  import sys
  from maze.core.agent.policy import Policy
  from maze.core.env.maze_env import MazeEnv
  from maze.core.trajectory_recording.utils.monitoring_setup import MonitoringSetup
label_desc: Use image node_id 2 for calling a global function with example usage: rollout(env, agent, n_max_steps, render, log_dir) and returns: cumulative_reward