class_docstr: Implementation of the model-based DreamerV3 RL algorithm described in [1].
class_docstr_tok: ["Implementation","of","the","model-based","DreamerV3","RL","algorithm","described","in","[","1","]","."]
code: def training_step(self) -> ResultDict:
      results = {}
  
      env_runner = self.workers.local_worker()
  
      # Push enough samples into buffer initially before we start training.
      if self.training_iteration == 0:
          logger.info(
              "Filling replay buffer so it contains at least "
              f"{self.config.batch_size_B * self.config.batch_length_T} timesteps "
              "(required for a single train batch)."
          )
  
      # Have we sampled yet in this `training_step()` call?
      have_sampled = False
      with self._timers[SAMPLE_TIMER]:
          # Continue sampling from the actual environment (and add collected samples
          # to our replay buffer) as long as we:
          while (
              # a) Don't have at least batch_size_B x batch_length_T timesteps stored
              # in the buffer. This is the minimum needed to train.
              self.replay_buffer.get_num_timesteps()
              < (self.config.batch_size_B * self.config.batch_length_T)
              # b) The computed `training_ratio` is >= the configured (desired)
              # training ratio (meaning we should continue sampling).
              or self.training_ratio >= self.config.training_ratio
              # c) we have not sampled at all yet in this `training_step()` call.
              or not have_sampled
          ):
              # Sample using the env runner's module.
              done_episodes, ongoing_episodes = env_runner.sample()
              # Add ongoing and finished episodes into buffer. The buffer will
              # automatically take care of properly concatenating (by episode IDs)
              # the different chunks of the same episodes, even if they come in via
              # separate `add()` calls.
              self.replay_buffer.add(
                  episodes=done_episodes + ongoing_episodes
              )
              have_sampled = True
  
              # We took B x T env steps.
              env_steps_last_regular_sample = sum(
                  len(eps) for eps in done_episodes + ongoing_episodes
              )
              total_sampled = env_steps_last_regular_sample
  
              # If we have never sampled before (just started the algo and not
              # recovered from a checkpoint), sample B random actions first.
              if self._counters[NUM_AGENT_STEPS_SAMPLED] == 0:
                  d_, o_ = env_runner.sample(
                      num_timesteps=(
                          self.config.batch_size_B
                          * self.config.batch_length_T
                      )
                      - env_steps_last_regular_sample,
                      random_actions=True,
                  )
                  self.replay_buffer.add(episodes=d_ + o_)
                  total_sampled += sum(len(eps) for eps in d_ + o_)
  
              self._counters[NUM_AGENT_STEPS_SAMPLED] += total_sampled
              self._counters[NUM_ENV_STEPS_SAMPLED] += total_sampled
  
      # Summarize environment interaction and buffer data.
      results[ALL_MODULES] = report_sampling_and_replay_buffer(
          replay_buffer=self.replay_buffer,
      )
  
      # Continue sampling batch_size_B x batch_length_T sized batches from the buffer
      # and using these to update our models (`LearnerGroup.update()`) until the
      # computed `training_ratio` is larger than the configured one, meaning we should
      # go back and collect more samples again from the actual environment.
      # However, when calculating the `training_ratio` here, we use only the
      # trained steps in this very `training_step()` call over the most recent sample
      # amount (`env_steps_last_regular_sample`), not the global values. This is to
      # avoid a heavy overtraining at the very beginning when we have just pre-filled
      # the buffer with the minimum amount of samples.
      replayed_steps_this_iter = sub_iter = 0
      while (
          replayed_steps_this_iter / env_steps_last_regular_sample
      ) < self.config.training_ratio:
          # Time individual batch updates.
          with self._timers[LEARN_ON_BATCH_TIMER]:
              logger.info(
                  f"\tSub-iteration {self.training_iteration}/{sub_iter})"
              )
  
              # Draw a new sample from the replay buffer.
              sample = self.replay_buffer.sample(
                  batch_size_B=self.config.batch_size_B,
                  batch_length_T=self.config.batch_length_T,
              )
              replayed_steps = (
                  self.config.batch_size_B * self.config.batch_length_T
              )
              replayed_steps_this_iter += replayed_steps
  
              if isinstance(
                  env_runner.env.single_action_space,
                  gym.spaces.Discrete,
              ):
                  sample["actions_ints"] = sample[SampleBatch.ACTIONS]
                  sample[SampleBatch.ACTIONS] = one_hot(
                      sample["actions_ints"],
                      depth=env_runner.env.single_action_space.n,
                  )
  
              # Perform the actual update via our learner group.
              train_results = self.learner_group.update_from_batch(
                  batch=SampleBatch(sample).as_multi_agent(),
                  reduce_fn=self._reduce_results,
              )
              self._counters[NUM_AGENT_STEPS_TRAINED] += replayed_steps
              self._counters[NUM_ENV_STEPS_TRAINED] += replayed_steps
  
              # Perform additional (non-gradient updates), such as the critic EMA-copy
              # update.
              with self._timers["critic_ema_update"]:
                  self.learner_group.additional_update(
                      timestep=self._counters[NUM_ENV_STEPS_SAMPLED],
                      reduce_fn=self._reduce_results,
                  )
  
              if self.config.report_images_and_videos:
                  report_predicted_vs_sampled_obs(
                      # TODO (sven): DreamerV3 is single-agent only.
                      results=train_results[DEFAULT_POLICY_ID],
                      sample=sample,
                      batch_size_B=self.config.batch_size_B,
                      batch_length_T=self.config.batch_length_T,
                      symlog_obs=do_symlog_obs(
                          env_runner.env.single_observation_space,
                          self.config.symlog_obs,
                      ),
                  )
  
              res = train_results[DEFAULT_POLICY_ID]
              logger.info(
                  f"\t\tWORLD_MODEL_L_total={res['WORLD_MODEL_L_total']:.5f} ("
                  f"L_pred={res['WORLD_MODEL_L_prediction']:.5f} ("
                  f"decoder/obs={res['WORLD_MODEL_L_decoder']} "
                  f"L_rew={res['WORLD_MODEL_L_reward']} "
                  f"L_cont={res['WORLD_MODEL_L_continue']}); "
                  f"L_dyn/rep={res['WORLD_MODEL_L_dynamics']:.5f})"
              )
              msg = "\t\t"
              if self.config.train_actor:
                  msg += f"L_actor={res['ACTOR_L_total']:.5f} "
              if self.config.train_critic:
                  msg += f"L_critic={res['CRITIC_L_total']:.5f} "
              logger.info(msg)
  
              sub_iter += 1
              self._counters[NUM_GRAD_UPDATES_LIFETIME] += 1
  
      # Update weights - after learning on the LearnerGroup - on all EnvRunner
      # workers.
      with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]:
          # Only necessary if RLModule is not shared between (local) EnvRunner and
          # (local) Learner.
          if (
              not self.config.share_module_between_env_runner_and_learner
          ):
              self._counters[
                  NUM_TRAINING_STEP_CALLS_SINCE_LAST_SYNCH_WORKER_WEIGHTS
              ] = 0
              self._counters[NUM_SYNCH_WORKER_WEIGHTS] += 1
              self.workers.sync_weights(
                  from_worker_or_learner_group=self.learner_group
              )
  
      # Try trick from https://medium.com/dive-into-ml-ai/dealing-with-memory-leak-
      # issue-in-keras-model-training-e703907a6501
      if self.config.gc_frequency_train_steps and (
          self.training_iteration % self.config.gc_frequency_train_steps
          == 0
      ):
          with self._timers[GARBAGE_COLLECTION_TIMER]:
              gc.collect()
  
      # Add train results and the actual training ratio to stats. The latter should
      # be close to the configured `training_ratio`.
      results.update(train_results)
      results[ALL_MODULES][
          "actual_training_ratio"
      ] = self.training_ratio
  
      # Return all results.
      return results
  
imports: import copy
  import gc
  import logging
  import tree
  from typing import Any, Dict, List, Optional, Union
  import gymnasium
  import numpy
  from ray.rllib.algorithms.algorithm import Algorithm
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
  from ray.rllib.algorithms.dreamerv3.dreamerv3_catalog import DreamerV3Catalog
  from ray.rllib.algorithms.dreamerv3.utils import do_symlog_obs
  from ray.rllib.algorithms.dreamerv3.utils.env_runner import DreamerV3EnvRunner
  from ray.rllib.algorithms.dreamerv3.utils.summaries import report_predicted_vs_sampled_obs, report_sampling_and_replay_buffer
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
  from ray.rllib.models.catalog import MODEL_DEFAULTS
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
  from ray.rllib.utils import deep_update
  from ray.rllib.utils.annotations import override, PublicAPI
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.numpy import one_hot
  from ray.rllib.utils.metrics import ALL_MODULES, GARBAGE_COLLECTION_TIMER, LEARN_ON_BATCH_TIMER, NUM_AGENT_STEPS_SAMPLED, NUM_AGENT_STEPS_TRAINED, NUM_ENV_STEPS_SAMPLED, NUM_ENV_STEPS_TRAINED, NUM_GRAD_UPDATES_LIFETIME, NUM_SYNCH_WORKER_WEIGHTS, NUM_TRAINING_STEP_CALLS_SINCE_LAST_SYNCH_WORKER_WEIGHTS, SAMPLE_TIMER, SYNCH_WORKER_WEIGHTS_TIMER
  from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer
  from ray.rllib.utils.typing import LearningRateOrSchedule, ResultDict
label_desc: Use image node_id 4 for calling the DreamerV3 obj's underlying member method code with example usage: obj.training_step() and returns: results