class_docstr: Rollouts a given policy in a given environment, recording the trajectory (in the form of raw actions
  and observations).
  
  Works with both standard and vectorized environments.
  
  :param env: Environment to run rollouts in. Will be reset before the first rollout.
  :param record_logits: Whether to record the policy logits.
  :param record_step_stats: Whether to record step statistics.
  :param record_episode_stats: Whether to record episode stats (happens only when an episode is done).
  :param record_next_observations: Whether to record next observation (i.e. observation following the action taken).
  :param terminate_on_done: Whether to end the rollout when the env is done (by default resets the env and continues
                            until the desired number of steps has been recorded). Only applicable in non-vectorized
                            scenarios.
class_docstr_tok: ["Rollouts","a","given","policy","in","a","given","environment",",","recording","the","trajectory","(","in","the","form","of","raw","actions","and","observations",")",".","Works","with","both","standard","and","vectorized","environments",".",":","param","env",":","Environment","to","run","rollouts","in",".","Will","be","reset","before","the","first","rollout",".",":","param","record_logits",":","Whether","to","record","the","policy","logits",".",":","param","record_step_stats",":","Whether","to","record","step","statistics",".",":","param","record_episode_stats",":","Whether","to","record","episode","stats","(","happens","only","when","an","episode","is","done",")",".",":","param","record_next_observations",":","Whether","to","record","next","observation","(","i.e",".","observation","following","the","action","taken",")",".",":","param","terminate_on_done",":","Whether","to","end","the","rollout","when","the","env","is","done","(","by","default","resets","the","env","and","continues","until","the","desired","number","of","steps","has","been","recorded",")",".","Only","applicable","in","non-vectorized","scenarios","."]
docstr: Perform and record a rollout with given policy, for given steps or until done.
  
  Note that the env is only reset on the very first rollout with this generator, the following rollouts
  just pick up where the previous left off. If required, you can avoid the initial reset by assigning
  the last observation (which will be recorded with the first step) into `self.last_observation`.
  
  :param policy: Policy to roll out.
  :param n_steps: How many steps to perform. If None, rollouts are performed until done=True.
  :param trajectory_id: Optionally, the ID of the trajectory that we are recording.
  :return: Recorded trajectory.
docstr_tok: ["Perform","and","record","a","rollout","with","given","policy",",","for","given","steps","or","until","done",".","Note","that","the","env","is","only","reset","on","the","very","first","rollout","with","this","generator",",","the","following","rollouts","just","pick","up","where","the","previous","left","off",".","If","required",",","you","can","avoid","the","initial","reset","by","assigning","the","last","observation","(","which","will","be","recorded","with","the","first","step",")","into","`","self.last_observation","`",".",":","param","policy",":","Policy","to","roll","out",".",":","param","n_steps",":","How","many","steps","to","perform",".","If","None",",","rollouts","are","performed","until","done=True",".",":","param","trajectory_id",":","Optionally",",","the","ID","of","the","trajectory","that","we","are","recording",".",":","return",":","Recorded","trajectory","."]
code: def rollout(
      self,
      policy: Policy,
      n_steps: Optional[int],
      trajectory_id: Optional[Any] = None,
  ) -> SpacesTrajectoryRecord:
      """Perform and record a rollout with given policy, for given steps or until done.
  
      Note that the env is only reset on the very first rollout with this generator, the following rollouts
      just pick up where the previous left off. If required, you can avoid the initial reset by assigning
      the last observation (which will be recorded with the first step) into `self.last_observation`.
  
      :param policy: Policy to roll out.
      :param n_steps: How many steps to perform. If None, rollouts are performed until done=True.
      :param trajectory_id: Optionally, the ID of the trajectory that we are recording.
      :return: Recorded trajectory.
      """
      # Check: Logits can be recorded with torch policy only
      if self.record_logits:
          assert isinstance(
              policy, TorchPolicy
          ), "to collect logits, the policy needs to be a Torch policy"
  
      # Initialize a trajectory record
      trajectory_record = SpacesTrajectoryRecord(
          trajectory_id if trajectory_id else self.rollout_counter
      )
      self.rollout_counter += 1
  
      # Reset the environment during the first rollout only
      if self.last_observation is None:
          self.last_observation = self.env.reset()
  
      # Step the desired number of (flat) steps
      step_count = 0
      while True:
          step_record = StructuredSpacesRecord()
  
          # Step through all sub-steps, i.e., step until the env time changes
          current_env_time = self.env.get_env_time()
          while np.all(current_env_time == self.env.get_env_time()):
              record = self._record_sub_step(policy=policy)
              step_record.append(record)
              # note that this also handles the special case of a done env after the first step
              if np.alltrue(record.done):
                  break
  
          if self.record_step_stats:
              step_record.step_stats = self.env.get_stats(
                  LogStatsLevel.STEP
              ).last_stats
  
          if (
              self.record_episode_stats
              and not self.is_vectorized
              and step_record.is_done()
          ):
              step_record.episode_stats = self.env.get_stats(
                  LogStatsLevel.EPISODE
              ).last_stats
  
          # Redistribute actor rewards, if available
          actor_rewards = self.env.get_actor_rewards()
          if actor_rewards is not None:
              assert len(actor_rewards) == len(
                  step_record.substep_records
              )
              for substep_record, reward in zip(
                  step_record.substep_records, actor_rewards
              ):
                  substep_record.reward = reward
  
          trajectory_record.append(step_record)
  
          # Limit maximum number of steps
          step_count += 1
          if n_steps and step_count >= n_steps:
              break
  
          # End prematurely on env done if desired
          if (
              self.terminate_on_done
              and not self.is_vectorized
              and step_record.is_done()
          ):
              break
  
      return trajectory_record
  
file_path: maze/maze/core/rollout/rollout_generator.py
imports: from copy import deepcopy
  from typing import Union, Optional, Any, Callable
  import numpy
  import torch
  from maze.core.agent.policy import Policy
  from maze.core.agent.torch_policy import TorchPolicy
  from maze.core.env.maze_env import MazeEnv
  from maze.core.log_stats.log_stats import LogStatsLevel, LogStatsAggregator, LogStatsValue
  from maze.core.trajectory_recording.records.spaces_record import SpacesRecord
  from maze.core.trajectory_recording.records.structured_spaces_record import StructuredSpacesRecord
  from maze.core.trajectory_recording.records.trajectory_record import SpacesTrajectoryRecord
  from maze.core.wrappers.log_stats_wrapper import LogStatsWrapper
  from maze.perception.perception_utils import convert_to_numpy
  from maze.train.parallelization.vector_env.structured_vector_env import StructuredVectorEnv
label_desc: Use image node_id 4 for calling the RolloutGenerator obj's underlying member method code with example usage: obj.rollout(policy, n_steps, trajectory_id) and returns: trajectory_record