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 one sub-step in the environment and return the record of it.
  
  Resets non-vectorised envs when done.
  
  :param policy: The policy to roll out.
  :return: Spaces record with the data recorded during this sub-step.
docstr_tok: ["Perform","one","sub-step","in","the","environment","and","return","the","record","of","it",".","Resets","non-vectorised","envs","when","done",".",":","param","policy",":","The","policy","to","roll","out",".",":","return",":","Spaces","record","with","the","data","recorded","during","this","sub-step","."]
code: def _record_sub_step(
      self, policy: Union[Policy, TorchPolicy]
  ) -> SpacesRecord:
      """Perform one sub-step in the environment and return the record of it.
  
      Resets non-vectorised envs when done.
  
      :param policy: The policy to roll out.
      :return: Spaces record with the data recorded during this sub-step.
      """
      record = SpacesRecord(
          actor_id=self.env.actor_id(),
          batch_shape=[self.env.n_envs] if self.is_vectorized else None,
      )
  
      # Record copy of the observation (as by default, the policy converts and handles it in place)
      record.observation = self.last_observation
  
      # Sample action and record logits if configured
      # Note: Copy the observation (as by default, the policy converts and handles it in place)
      if self.record_logits:
          with torch.no_grad():
              step_policy_output = policy.compute_substep_policy_output(
                  self.last_observation.copy(), actor_id=record.actor_id
              )
              action = convert_to_numpy(
                  step_policy_output.prob_dist.sample(),
                  cast=None,
                  in_place=False,
              )
              record.logits = step_policy_output.action_logits
      else:
          # Inject the MazeEnv state if desired by the policy
          maze_state = (
              self.env.get_maze_state()
              if policy.needs_state()
              else None
          )
          env = self.env if policy.needs_env() else None
          action = policy.compute_action(
              self.last_observation.copy(),
              actor_id=record.actor_id,
              maze_state=maze_state,
              env=env,
              deterministic=False,
          )
      record.action = action
      record.policy_record = policy.write_policy_record()
  
      # Take the step
      (
          self.last_observation,
          record.reward,
          record.done,
          record.info,
      ) = self.env.step(action)
  
      # Record the resulting observation if requested
      if self.record_next_observations:
          record.next_observation = self.last_observation
  
      # Reset the env if done and keep the terminal observation
      if not self.is_vectorized and record.done:
          record.info["terminal_observation"] = deepcopy(
              self.last_observation
          )
          self.last_observation = self.env.reset()
  
      return 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 5 for calling the RolloutGenerator obj's underlying member method code with example usage: obj._record_sub_step(policy) and returns: record