class_docstr: The generic environment runner for the single agent case.
class_docstr_tok: ["The","generic","environment","runner","for","the","single","agent","case","."]
docstr: Helper method to sample n timesteps.
docstr_tok: ["Helper","method","to","sample","n","timesteps","."]
code: def _sample_timesteps(
      self,
      num_timesteps: int,
      explore: bool = True,
      random_actions: bool = False,
      force_reset: bool = False,
  ) -> List[SingleAgentEpisode]:
      """Helper method to sample n timesteps."""
  
      done_episodes_to_return: List[SingleAgentEpisode] = []
  
      # Have to reset the env (on all vector sub_envs).
      if force_reset or self._needs_initial_reset:
          # Create n new episodes and make the `on_episode_created` callbacks.
          self._episodes = []
          for env_index in range(self.num_envs):
              self._episodes.append(SingleAgentEpisode())
              self._make_on_episode_callback(
                  "on_episode_created", env_index
              )
  
          obs, infos = self.env.reset()
  
          # Call `on_episode_start()` callbacks.
          for env_index in range(self.num_envs):
              self._make_on_episode_callback(
                  "on_episode_start", env_index
              )
  
          # We just reset the env. Don't have to force this again in the next
          # call to `self._sample_timesteps()`.
          self._needs_initial_reset = False
  
          self._episodes = [
              SingleAgentEpisode(
                  observation_space=self.env.single_observation_space,
                  action_space=self.env.single_action_space,
              )
              for _ in range(self.num_envs)
          ]
  
          # Set initial obs and states in the episodes.
          for env_index in range(self.num_envs):
              # TODO (sven): Maybe move this into connector pipeline
              # (even if automated).
              self._episodes[env_index].add_env_reset(
                  observation=obs[env_index],
                  infos=infos[env_index],
              )
  
      # Loop through env in enumerate.(self._episodes):
      ts = 0
  
      while ts < num_timesteps:
          # Act randomly.
          if random_actions:
              to_env = {
                  SampleBatch.ACTIONS: self.env.action_space.sample(),
              }
          # Compute an action using the RLModule.
          else:
              to_module = self._env_to_module(
                  rl_module=self.module,
                  episodes=self._episodes,
                  explore=explore,
                  # persistent_data=None, #TODO
              )
              # Explore or not.
              if explore:
                  to_env = self.module.forward_exploration(to_module)
              else:
                  to_env = self.module.forward_inference(to_module)
  
              to_env = self._module_to_env(
                  rl_module=self.module,
                  data=to_env,
                  episodes=self._episodes,
                  explore=explore,
                  # persistent_data=None, #TODO
              )
  
          actions = to_env.pop(SampleBatch.ACTIONS)
  
          obs, rewards, terminateds, truncateds, infos = self.env.step(
              actions
          )
  
          ts += self.num_envs
  
          for env_index in range(self.num_envs):
              # The last entry in self.observations[i] is already the reset
              # obs of the new episode.
              # TODO (simon): This might be unfortunate if a user needs to set a
              #  certain env parameter during different episodes (for example for
              #  benchmarking).
              extra_model_output = tree.map_structure(
                  lambda s: s[env_index], to_env
              )
  
              # In inference we have only the action logits.
              if terminateds[env_index] or truncateds[env_index]:
                  # Finish the episode with the actual terminal observation stored in
                  # the info dict.
                  self._episodes[env_index].add_env_step(
                      # Gym vector env provides the `"final_observation"`.
                      # Pop these out of the infos dict so this information doesn't
                      # appear in the next episode as well (at index=0).
                      infos[env_index].pop("final_observation"),
                      actions[env_index],
                      rewards[env_index],
                      infos=infos[env_index].pop("final_info"),
                      terminated=terminateds[env_index],
                      truncated=truncateds[env_index],
                      extra_model_outputs=extra_model_output,
                  )
  
                  # Make the `on_episode_step` callback (before finalizing the
                  # episode object).
                  self._make_on_episode_callback(
                      "on_episode_step", env_index
                  )
                  done_episodes_to_return.append(
                      self._episodes[env_index].finalize()
                  )
  
                  # Make the `on_episode_env` callback (after having finalized the
                  # episode object).
                  self._make_on_episode_callback(
                      "on_episode_end", env_index
                  )
  
                  # Create a new episode object with already the reset data in it.
                  self._episodes[env_index] = SingleAgentEpisode(
                      observations=[obs[env_index]],
                      infos=[infos[env_index]],
                      observation_space=self.env.single_observation_space,
                      action_space=self.env.single_action_space,
                  )
  
                  # Make the `on_episode_start` callback.
                  self._make_on_episode_callback(
                      "on_episode_start", env_index
                  )
  
              else:
                  self._episodes[env_index].add_env_step(
                      obs[env_index],
                      actions[env_index],
                      rewards[env_index],
                      infos=infos[env_index],
                      extra_model_outputs=extra_model_output,
                  )
  
                  # Make the `on_episode_step` callback.
                  self._make_on_episode_callback(
                      "on_episode_step", env_index
                  )
  
      # Return done episodes ...
      self._done_episodes_for_metrics.extend(done_episodes_to_return)
      # ... and all ongoing episode chunks.
  
      # Also, make sure we start new episode chunks (continuing the ongoing episodes
      # from the to-be-returned chunks).
      ongoing_episodes_continuations = [
          eps.cut(
              len_lookback_buffer=self.config.episode_lookback_horizon
          )
          for eps in self._episodes
      ]
  
      ongoing_episodes_to_return = []
      for eps in self._episodes:
          # Just started Episodes do not have to be returned. There is no data
          # in them anyway.
          if eps.t == 0:
              continue
          eps.validate()
          self._ongoing_episodes_for_metrics[eps.id_].append(eps)
          # Return finalized (numpy'ized) Episodes.
          ongoing_episodes_to_return.append(eps.finalize())
  
      # Continue collecting into the cut Episode chunks.
      self._episodes = ongoing_episodes_continuations
  
      # Record last metrics collection.
      self._ts_since_last_metrics += ts
  
      return done_episodes_to_return + ongoing_episodes_to_return
  
imports: import gymnasium
  import numpy
  import tree
  from collections import defaultdict
  from functools import partial
  from typing import Dict, List, Optional
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
  from ray.rllib.algorithms.callbacks import DefaultCallbacks
  from ray.rllib.core.rl_module.rl_module import RLModule, SingleAgentRLModuleSpec
  from ray.rllib.env.env_runner import EnvRunner
  from ray.rllib.env.single_agent_episode import SingleAgentEpisode
  from ray.rllib.env.utils import _gym_env_creator
  from ray.rllib.evaluation.metrics import RolloutMetrics
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
  from ray.rllib.utils.annotations import ExperimentalAPI, override
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
  from ray.rllib.utils.typing import TensorType
  from ray.tune.registry import ENV_CREATOR, _global_registry
label_desc: Use image node_id 3 for calling the SingleAgentEnvRunner obj's underlying member method code with example usage: obj._sample_timesteps(num_timesteps, explore, random_actions, force_reset) and returns: unknown