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 run n episodes.
  
  See docstring of `self.sample()` for more details.
docstr_tok: ["Helper","method","to","run","n","episodes",".","See","docstring","of","`","self.sample","(",")","`","for","more","details","."]
code: def _sample_episodes(
      self,
      num_episodes: int,
      explore: bool = True,
      random_actions: bool = False,
      with_render_data: bool = False,
  ) -> List["SingleAgentEpisode"]:
      """Helper method to run n episodes.
  
      See docstring of `self.sample()` for more details.
      """
      # If user calls sample(num_timesteps=..) after this, we must reset again
      # at the beginning.
      self._needs_initial_reset = True
  
      done_episodes_to_return: List["SingleAgentEpisode"] = []
  
      obs, infos = self.env.reset()
      episodes = []
      for env_index in range(self.num_envs):
          episodes.append(
              SingleAgentEpisode(
                  observation_space=self.env.single_observation_space,
                  action_space=self.env.single_action_space,
              )
          )
          self._make_on_episode_callback(
              "on_episode_created", env_index, episodes
          )
  
      render_images = [None] * self.num_envs
      if with_render_data:
          render_images = [e.render() for e in self.env.envs]
  
      for env_index in range(self.num_envs):
          episodes[env_index].add_env_reset(
              observation=obs[env_index],
              infos=infos[env_index],
              render_image=render_images[env_index],
          )
          self._make_on_episode_callback(
              "on_episode_start", env_index, episodes
          )
  
      eps = 0
      while eps < num_episodes:
          if random_actions:
              to_env = {
                  SampleBatch.ACTIONS: self.env.action_space.sample(),
              }
          else:
              to_module = self._env_to_module(
                  rl_module=self.module,
                  episodes=episodes,
                  explore=explore,
              )
              # 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=episodes,
                  explore=explore,
              )
  
          actions = to_env.pop(SampleBatch.ACTIONS)
          obs, rewards, terminateds, truncateds, infos = self.env.step(
              actions
          )
          if with_render_data:
              render_images = [e.render() for e in self.env.envs]
  
          for env_index in range(self.num_envs):
              # Extract info and state for vector sub_env.
              # info = {k: v[i] for k, v in infos.items()}
              # The last entry in self.observations[i] is already the reset
              # obs of the new episode.
              extra_model_output = tree.map_structure(
                  lambda s: s[env_index], to_env
              )
  
              if terminateds[env_index] or truncateds[env_index]:
                  eps += 1
  
                  episodes[env_index].add_env_step(
                      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 `on_episode_step` callback before finalizing the episode.
                  self._make_on_episode_callback(
                      "on_episode_step", env_index, episodes
                  )
                  done_episodes_to_return.append(
                      episodes[env_index].finalize()
                  )
  
                  # Make `on_episode_end` callback after finalizing the episode.
                  self._make_on_episode_callback(
                      "on_episode_end", env_index, episodes
                  )
  
                  # Also early-out if we reach the number of episodes within this
                  # for-loop.
                  if eps == num_episodes:
                      break
  
                  # Create a new episode object.
                  episodes[env_index] = SingleAgentEpisode(
                      observations=[obs[env_index]],
                      infos=[infos[env_index]],
                      render_images=None
                      if render_images[env_index] is None
                      else [render_images[env_index]],
                      observation_space=self.env.single_observation_space,
                      action_space=self.env.single_action_space,
                  )
                  # Make `on_episode_start` callback.
                  self._make_on_episode_callback(
                      "on_episode_start", env_index, episodes
                  )
              else:
                  episodes[env_index].add_env_step(
                      obs[env_index],
                      actions[env_index],
                      rewards[env_index],
                      infos=infos[env_index],
                      render_image=render_images[env_index],
                      extra_model_outputs=extra_model_output,
                  )
                  # Make `on_episode_step` callback.
                  self._make_on_episode_callback(
                      "on_episode_step", env_index, episodes
                  )
  
      self._done_episodes_for_metrics.extend(done_episodes_to_return)
      self._ts_since_last_metrics += sum(
          len(eps) for eps in done_episodes_to_return
      )
  
      # Initialized episodes have to be removed as they lack `extra_model_outputs`.
      samples = [
          episode
          for episode in done_episodes_to_return
          if episode.t > 0
      ]
  
      return samples
  
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 4 for calling the SingleAgentEnvRunner obj's underlying member method code with example usage: obj._sample_episodes(num_episodes, explore, random_actions, with_render_data) and returns: samples