class_docstr: An environment runner to collect data from vectorized gymnasium environments.
class_docstr_tok: ["An","environment","runner","to","collect","data","from","vectorized","gymnasium","environments","."]
docstr: 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.
      """
      done_episodes_to_return = []
  
      obs, _ = self.env.reset()
      episodes = [SingleAgentEpisode() for _ in range(self.num_envs)]
  
      # Multiply states n times according to our vector env batch size (num_envs).
      states = tree.map_structure(
          lambda s: np.repeat(s, self.num_envs, axis=0),
          self.module.get_initial_state(),
      )
      is_first = np.ones((self.num_envs,))
  
      render_images = [None] * self.num_envs
      if with_render_data:
          render_images = [e.render() for e in self.env.envs]
  
      for i in range(self.num_envs):
          episodes[i].add_env_reset(
              observation=obs[i],
              render_image=render_images[i],
          )
  
      eps = 0
      while eps < num_episodes:
          if random_actions:
              actions = self.env.action_space.sample()
          else:
              batch = {
                  STATE_IN: tree.map_structure(
                      lambda s: tf.convert_to_tensor(s), states
                  ),
                  SampleBatch.OBS: tf.convert_to_tensor(obs),
                  "is_first": tf.convert_to_tensor(is_first),
              }
  
              if explore:
                  outs = self.module.forward_exploration(batch)
              else:
                  outs = self.module.forward_inference(batch)
  
              actions = outs[SampleBatch.ACTIONS].numpy()
              if isinstance(
                  self.env.single_action_space, gym.spaces.Discrete
              ):
                  actions = np.argmax(actions, axis=-1)
              states = tree.map_structure(
                  lambda s: s.numpy(), outs[STATE_OUT]
              )
  
          obs, rewards, terminateds, truncateds, infos = self.env.step(
              actions
          )
          if with_render_data:
              render_images = [e.render() for e in self.env.envs]
  
          for i in range(self.num_envs):
              # The last entry in self.observations[i] is already the reset
              # obs of the new episode.
              if terminateds[i] or truncateds[i]:
                  eps += 1
  
                  episodes[i].add_env_step(
                      observation=infos["final_observation"][i],
                      action=actions[i],
                      reward=rewards[i],
                      terminated=terminateds[i],
                      truncated=truncateds[i],
                  )
                  done_episodes_to_return.append(episodes[i])
  
                  # Also early-out if we reach the number of episodes within this
                  # for-loop.
                  if eps == num_episodes:
                      break
  
                  # Reset h-states to the model's initial ones b/c we are starting a
                  # new episode.
                  for k, v in self.module.get_initial_state().items():
                      states[k][i] = v.numpy()
                  is_first[i] = True
  
                  episodes[i] = SingleAgentEpisode(
                      observations=[obs[i]],
                      render_images=(
                          [render_images[i]]
                          if with_render_data
                          else None
                      ),
                  )
              else:
                  episodes[i].add_env_step(
                      observation=obs[i],
                      action=actions[i],
                      reward=rewards[i],
                      render_image=render_images[i],
                  )
                  is_first[i] = False
  
      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
      )
  
      # If user calls sample(num_timesteps=..) after this, we must reset again
      # at the beginning.
      self._needs_initial_reset = True
  
      return done_episodes_to_return
  
imports: from collections import defaultdict
  from functools import partial
  from typing import List, Tuple
  import gymnasium
  import numpy
  import tree
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
  from ray.rllib.core.models.base import STATE_IN, STATE_OUT
  from ray.rllib.env.env_runner import EnvRunner
  from ray.rllib.env.wrappers.atari_wrappers import NoopResetEnv, MaxAndSkipEnv
  from ray.rllib.env.wrappers.dm_control_wrapper import DMCEnv
  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 override
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.env.single_agent_episode import SingleAgentEpisode
  from ray.rllib.utils.numpy import one_hot
  from ray.tune.registry import ENV_CREATOR, _global_registry
label_desc: Use image node_id 4 for calling the DreamerV3EnvRunner obj's underlying member method code with example usage: obj._sample_episodes(num_episodes, explore, random_actions, with_render_data) and returns: done_episodes_to_return