image-match/python/py/Lib/site-packages/ray/rllib/algorithms/dreamerv3/dreamerv3.py

"""
[1] Mastering Diverse Domains through World Models - 2023
D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
https://arxiv.org/pdf/2301.04104v1.pdf

[2] Mastering Atari with Discrete World Models - 2021
D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
https://arxiv.org/pdf/2010.02193.pdf
"""

import logging
from typing import Any, Dict, Optional, Union

import gymnasium as gym
from typing_extensions import Self

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.add_is_firsts_to_batch import (
    AddIsFirstsToBatch,
)
from ray.rllib.algorithms.dreamerv3.utils.summaries import (
    report_dreamed_eval_trajectory_vs_samples,
    report_predicted_vs_sampled_obs,
    report_sampling_and_replay_buffer,
)
from ray.rllib.connectors.common import AddStatesFromEpisodesToBatch
from ray.rllib.core import DEFAULT_MODULE_ID
from ray.rllib.core.columns import Columns
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.env import INPUT_ENV_SINGLE_SPACES
from ray.rllib.execution.rollout_ops import synchronous_parallel_sample
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.utils import deep_update
from ray.rllib.utils.annotations import PublicAPI, override
from ray.rllib.utils.metrics import (
    ENV_RUNNER_RESULTS,
    LEARN_ON_BATCH_TIMER,
    LEARNER_RESULTS,
    NUM_ENV_STEPS_SAMPLED_LIFETIME,
    NUM_ENV_STEPS_TRAINED_LIFETIME,
    NUM_GRAD_UPDATES_LIFETIME,
    NUM_SYNCH_WORKER_WEIGHTS,
    REPLAY_BUFFER_RESULTS,
    SAMPLE_TIMER,
    SYNCH_WORKER_WEIGHTS_TIMER,
    TIMERS,
)
from ray.rllib.utils.numpy import one_hot
from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer
from ray.rllib.utils.typing import LearningRateOrSchedule

logger = logging.getLogger(__name__)


class DreamerV3Config(AlgorithmConfig):
    """Defines a configuration class from which a DreamerV3 can be built.

    .. testcode::

        from ray.rllib.algorithms.dreamerv3 import DreamerV3Config
        config = (
            DreamerV3Config()
            .environment("CartPole-v1")
            .training(
                model_size="XS",
                training_ratio=1,
                # TODO
                model={
                    "batch_size_B": 1,
                    "batch_length_T": 1,
                    "horizon_H": 1,
                    "gamma": 0.997,
                    "model_size": "XS",
                },
            )
        )

        config = config.learners(num_learners=0)
        # Build a Algorithm object from the config and run 1 training iteration.
        algo = config.build()
        # algo.train()
        del algo
    """

    def __init__(self, algo_class=None):
        """Initializes a DreamerV3Config instance."""
        super().__init__(algo_class=algo_class or DreamerV3)

        # fmt: off
        # __sphinx_doc_begin__

        # DreamerV3 specific settings:
        self.model_size = "XS"
        self.training_ratio = 1024

        self.replay_buffer_config = {
            "type": "EpisodeReplayBuffer",
            "capacity": int(1e6),
        }
        self.world_model_lr = 1e-4
        self.actor_lr = 3e-5
        self.critic_lr = 3e-5
        self.batch_size_B = 16
        self.batch_length_T = 64
        self.horizon_H = 15
        self.gae_lambda = 0.95  # [1] eq. 7.
        self.entropy_scale = 3e-4  # [1] eq. 11.
        self.return_normalization_decay = 0.99  # [1] eq. 11 and 12.
        self.train_critic = True
        self.train_actor = True
        self.intrinsic_rewards_scale = 0.1
        self.world_model_grad_clip_by_global_norm = 1000.0
        self.critic_grad_clip_by_global_norm = 100.0
        self.actor_grad_clip_by_global_norm = 100.0
        self.symlog_obs = "auto"
        self.use_float16 = False
        self.use_curiosity = False

        # Reporting.
        # DreamerV3 is super sample efficient and only needs very few episodes
        # (normally) to learn. Leaving this at its default value would gravely
        # underestimate the learning performance over the course of an experiment.
        self.metrics_num_episodes_for_smoothing = 1
        self.report_individual_batch_item_stats = False
        self.report_dream_data = False
        self.report_images_and_videos = False

        # Override some of AlgorithmConfig's default values with DreamerV3-specific
        # values.
        self.lr = None
        self.gamma = 0.997  # [1] eq. 7.
        # Do not use! Set `batch_size_B` and `batch_length_T` instead.
        self.train_batch_size = None
        self.num_env_runners = 0
        self.rollout_fragment_length = 1
        # Dreamer only runs on the new API stack.
        self.enable_rl_module_and_learner = True
        self.enable_env_runner_and_connector_v2 = True
        # TODO (sven): DreamerV3 still uses its own EnvRunner class. This env-runner
        #  does not use connectors. We therefore should not attempt to merge/broadcast
        #  the connector states between EnvRunners (if >0). Note that this is only
        #  relevant if num_env_runners > 0, which is normally not the case when using
        #  this algo.
        self.use_worker_filter_stats = False
        # __sphinx_doc_end__
        # fmt: on

    @override(AlgorithmConfig)
    def build_env_to_module_connector(self, env, spaces, device):
        connector = super().build_env_to_module_connector(env, spaces, device)

        # Prepend the "is_first" connector such that the RSSM knows, when to insert
        # its (learned) internal state into the batch.
        # We have to do this before the `AddStatesFromEpisodesToBatch` piece
        # such that the column is properly batched/time-ranked.
        if self.add_default_connectors_to_learner_pipeline:
            connector.insert_before(
                AddStatesFromEpisodesToBatch,
                AddIsFirstsToBatch(),
            )
        return connector

    @property
    def batch_size_B_per_learner(self):
        """Returns the batch_size_B per Learner worker.

        Needed by some of the DreamerV3 loss math."""
        return self.batch_size_B // (self.num_learners or 1)

    @override(AlgorithmConfig)
    def training(
        self,
        *,
        model_size: Optional[str] = NotProvided,
        training_ratio: Optional[float] = NotProvided,
        batch_size_B: Optional[int] = NotProvided,
        batch_length_T: Optional[int] = NotProvided,
        horizon_H: Optional[int] = NotProvided,
        gae_lambda: Optional[float] = NotProvided,
        entropy_scale: Optional[float] = NotProvided,
        return_normalization_decay: Optional[float] = NotProvided,
        train_critic: Optional[bool] = NotProvided,
        train_actor: Optional[bool] = NotProvided,
        intrinsic_rewards_scale: Optional[float] = NotProvided,
        world_model_lr: Optional[LearningRateOrSchedule] = NotProvided,
        actor_lr: Optional[LearningRateOrSchedule] = NotProvided,
        critic_lr: Optional[LearningRateOrSchedule] = NotProvided,
        world_model_grad_clip_by_global_norm: Optional[float] = NotProvided,
        critic_grad_clip_by_global_norm: Optional[float] = NotProvided,
        actor_grad_clip_by_global_norm: Optional[float] = NotProvided,
        symlog_obs: Optional[Union[bool, str]] = NotProvided,
        use_float16: Optional[bool] = NotProvided,
        replay_buffer_config: Optional[dict] = NotProvided,
        use_curiosity: Optional[bool] = NotProvided,
        **kwargs,
    ) -> Self:
        """Sets the training related configuration.

        Args:
            model_size: The main switch for adjusting the overall model size. See [1]
                (table B) for more information on the effects of this setting on the
                model architecture.
                Supported values are "XS", "S", "M", "L", "XL" (as per the paper), as
                well as, "nano", "micro", "mini", and "XXS" (for RLlib's
                implementation). See ray.rllib.algorithms.dreamerv3.utils.
                __init__.py for the details on what exactly each size does to the layer
                sizes, number of layers, etc..
            training_ratio: The ratio of total steps trained (sum of the sizes of all
                batches ever sampled from the replay buffer) over the total env steps
                taken (in the actual environment, not the dreamed one). For example,
                if the training_ratio is 1024 and the batch size is 1024, we would take
                1 env step for every training update: 1024 / 1. If the training ratio
                is 512 and the batch size is 1024, we would take 2 env steps and then
                perform a single training update (on a 1024 batch): 1024 / 2.
            batch_size_B: The batch size (B) interpreted as number of rows (each of
                length `batch_length_T`) to sample from the replay buffer in each
                iteration.
            batch_length_T: The batch length (T) interpreted as the length of each row
                sampled from the replay buffer in each iteration. Note that
                `batch_size_B` rows will be sampled in each iteration. Rows normally
                contain consecutive data (consecutive timesteps from the same episode),
                but there might be episode boundaries in a row as well.
            horizon_H: The horizon (in timesteps) used to create dreamed data from the
                world model, which in turn is used to train/update both actor- and
                critic networks.
            gae_lambda: The lambda parameter used for computing the GAE-style
                value targets for the actor- and critic losses.
            entropy_scale: The factor with which to multiply the entropy loss term
                inside the actor loss.
            return_normalization_decay: The decay value to use when computing the
                running EMA values for return normalization (used in the actor loss).
            train_critic: Whether to train the critic network. If False, `train_actor`
                must also be False (cannot train actor w/o training the critic).
            train_actor: Whether to train the actor network. If True, `train_critic`
                must also be True (cannot train actor w/o training the critic).
            intrinsic_rewards_scale: The factor to multiply intrinsic rewards with
                before adding them to the extrinsic (environment) rewards.
            world_model_lr: The learning rate or schedule for the world model optimizer.
            actor_lr: The learning rate or schedule for the actor optimizer.
            critic_lr: The learning rate or schedule for the critic optimizer.
            world_model_grad_clip_by_global_norm: World model grad clipping value
                (by global norm).
            critic_grad_clip_by_global_norm: Critic grad clipping value
                (by global norm).
            actor_grad_clip_by_global_norm: Actor grad clipping value (by global norm).
            symlog_obs: Whether to symlog observations or not. If set to "auto"
                (default), will check for the environment's observation space and then
                only symlog if not an image space.
            use_float16: Whether to train with mixed float16 precision. In this mode,
                model parameters are stored as float32, but all computations are
                performed in float16 space (except for losses and distribution params
                and outputs).
            replay_buffer_config: Replay buffer config.
                Only serves in DreamerV3 to set the capacity of the replay buffer.
                Note though that in the paper ([1]) a size of 1M is used for all
                benchmarks and there doesn't seem to be a good reason to change this
                parameter.
                Examples:
                {
                "type": "EpisodeReplayBuffer",
                "capacity": 100000,
                }

        Returns:
            This updated AlgorithmConfig object.
        """
        # Not fully supported/tested yet.
        if use_curiosity is not NotProvided:
            raise ValueError(
                "`DreamerV3Config.curiosity` is not fully supported and tested yet! "
                "It thus remains disabled for now."
            )

        # Pass kwargs onto super's `training()` method.
        super().training(**kwargs)

        if model_size is not NotProvided:
            self.model_size = model_size
        if training_ratio is not NotProvided:
            self.training_ratio = training_ratio
        if batch_size_B is not NotProvided:
            self.batch_size_B = batch_size_B
        if batch_length_T is not NotProvided:
            self.batch_length_T = batch_length_T
        if horizon_H is not NotProvided:
            self.horizon_H = horizon_H
        if gae_lambda is not NotProvided:
            self.gae_lambda = gae_lambda
        if entropy_scale is not NotProvided:
            self.entropy_scale = entropy_scale
        if return_normalization_decay is not NotProvided:
            self.return_normalization_decay = return_normalization_decay
        if train_critic is not NotProvided:
            self.train_critic = train_critic
        if train_actor is not NotProvided:
            self.train_actor = train_actor
        if intrinsic_rewards_scale is not NotProvided:
            self.intrinsic_rewards_scale = intrinsic_rewards_scale
        if world_model_lr is not NotProvided:
            self.world_model_lr = world_model_lr
        if actor_lr is not NotProvided:
            self.actor_lr = actor_lr
        if critic_lr is not NotProvided:
            self.critic_lr = critic_lr
        if world_model_grad_clip_by_global_norm is not NotProvided:
            self.world_model_grad_clip_by_global_norm = (
                world_model_grad_clip_by_global_norm
            )
        if critic_grad_clip_by_global_norm is not NotProvided:
            self.critic_grad_clip_by_global_norm = critic_grad_clip_by_global_norm
        if actor_grad_clip_by_global_norm is not NotProvided:
            self.actor_grad_clip_by_global_norm = actor_grad_clip_by_global_norm
        if symlog_obs is not NotProvided:
            self.symlog_obs = symlog_obs
        if use_float16 is not NotProvided:
            self.use_float16 = use_float16
        if replay_buffer_config is not NotProvided:
            # Override entire `replay_buffer_config` if `type` key changes.
            # Update, if `type` key remains the same or is not specified.
            new_replay_buffer_config = deep_update(
                {"replay_buffer_config": self.replay_buffer_config},
                {"replay_buffer_config": replay_buffer_config},
                False,
                ["replay_buffer_config"],
                ["replay_buffer_config"],
            )
            self.replay_buffer_config = new_replay_buffer_config["replay_buffer_config"]

        return self

    @override(AlgorithmConfig)
    def reporting(
        self,
        *,
        report_individual_batch_item_stats: Optional[bool] = NotProvided,
        report_dream_data: Optional[bool] = NotProvided,
        report_images_and_videos: Optional[bool] = NotProvided,
        **kwargs,
    ):
        """Sets the reporting related configuration.

        Args:
            report_individual_batch_item_stats: Whether to include loss and other stats
                per individual timestep inside the training batch in the result dict
                returned by `training_step()`. If True, besides the `CRITIC_L_total`,
                the individual critic loss values per batch row and time axis step
                in the train batch (CRITIC_L_total_B_T) will also be part of the
                results.
            report_dream_data:  Whether to include the dreamed trajectory data in the
                result dict returned by `training_step()`. If True, however, will
                slice each reported item in the dream data down to the shape.
                (H, B, t=0, ...), where H is the horizon and B is the batch size. The
                original time axis will only be represented by the first timestep
                to not make this data too large to handle.
            report_images_and_videos: Whether to include any image/video data in the
                result dict returned by `training_step()`.
            **kwargs:

        Returns:
            This updated AlgorithmConfig object.
        """
        super().reporting(**kwargs)

        if report_individual_batch_item_stats is not NotProvided:
            self.report_individual_batch_item_stats = report_individual_batch_item_stats
        if report_dream_data is not NotProvided:
            self.report_dream_data = report_dream_data
        if report_images_and_videos is not NotProvided:
            self.report_images_and_videos = report_images_and_videos

        return self

    @override(AlgorithmConfig)
    def validate(self) -> None:
        # Call the super class' validation method first.
        super().validate()

        # Make sure, users are not using DreamerV3 yet for multi-agent:
        if self.is_multi_agent:
            self._value_error("DreamerV3 does NOT support multi-agent setups yet!")

        # Make sure, we are configure for the new API stack.
        if not self.enable_rl_module_and_learner:
            self._value_error(
                "DreamerV3 must be run with `config.api_stack("
                "enable_rl_module_and_learner=True)`!"
            )

        # If run on several Learners, the provided batch_size_B must be a multiple
        # of `num_learners`.
        if self.num_learners > 1 and (self.batch_size_B % self.num_learners != 0):
            self._value_error(
                f"Your `batch_size_B` ({self.batch_size_B}) must be a multiple of "
                f"`num_learners` ({self.num_learners}) in order for "
                "DreamerV3 to be able to split batches evenly across your Learner "
                "processes."
            )

        # Cannot train actor w/o critic.
        if self.train_actor and not self.train_critic:
            self._value_error(
                "Cannot train actor network (`train_actor=True`) w/o training critic! "
                "Make sure you either set `train_critic=True` or `train_actor=False`."
            )
        # Use DreamerV3 specific batch size settings.
        if self.train_batch_size is not None:
            self._value_error(
                "`train_batch_size` should NOT be set! Use `batch_size_B` and "
                "`batch_length_T` instead."
            )
        # Must be run with `EpisodeReplayBuffer` type.
        if self.replay_buffer_config.get("type") != "EpisodeReplayBuffer":
            self._value_error(
                "DreamerV3 must be run with the `EpisodeReplayBuffer` type! None "
                "other supported."
            )

    @override(AlgorithmConfig)
    def get_default_learner_class(self):
        if self.framework_str == "torch":
            from ray.rllib.algorithms.dreamerv3.torch.dreamerv3_torch_learner import (
                DreamerV3TorchLearner,
            )

            return DreamerV3TorchLearner
        else:
            raise ValueError(f"The framework {self.framework_str} is not supported.")

    @override(AlgorithmConfig)
    def get_default_rl_module_spec(self) -> RLModuleSpec:
        if self.framework_str == "torch":
            from ray.rllib.algorithms.dreamerv3.torch.dreamerv3_torch_rl_module import (
                DreamerV3TorchRLModule as module,
            )

        else:
            raise ValueError(f"The framework {self.framework_str} is not supported.")

        return RLModuleSpec(module_class=module, catalog_class=DreamerV3Catalog)

    @property
    @override(AlgorithmConfig)
    def _model_config_auto_includes(self) -> Dict[str, Any]:
        return super()._model_config_auto_includes | {
            "gamma": self.gamma,
            "horizon_H": self.horizon_H,
            "model_size": self.model_size,
            "symlog_obs": self.symlog_obs,
            "use_float16": self.use_float16,
            "batch_length_T": self.batch_length_T,
        }


class DreamerV3(Algorithm):
    """Implementation of the model-based DreamerV3 RL algorithm described in [1]."""

    # TODO (sven): Deprecate/do-over the Algorithm.compute_single_action() API.
    @override(Algorithm)
    def compute_single_action(self, *args, **kwargs):
        raise NotImplementedError(
            "DreamerV3 does not support the `compute_single_action()` API. Refer to the"
            " README here (https://github.com/ray-project/ray/tree/master/rllib/"
            "algorithms/dreamerv3) to find more information on how to run action "
            "inference with this algorithm."
        )

    @classmethod
    @override(Algorithm)
    def get_default_config(cls) -> DreamerV3Config:
        return DreamerV3Config()

    @override(Algorithm)
    def setup(self, config: AlgorithmConfig):
        super().setup(config)

        # Share RLModule between EnvRunner and single (local) Learner instance.
        # To avoid possibly expensive weight synching step.
        # if self.config.share_module_between_env_runner_and_learner:
        #    assert self.env_runner.module is None
        #    self.env_runner.module = self.learner_group._learner.module[
        #        DEFAULT_MODULE_ID
        #    ]

        # Create a replay buffer for storing actual env samples.
        self.replay_buffer = EpisodeReplayBuffer(
            capacity=self.config.replay_buffer_config["capacity"],
            batch_size_B=self.config.batch_size_B,
            batch_length_T=self.config.batch_length_T,
        )

    @override(Algorithm)
    def training_step(self) -> None:
        # 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.metrics.log_time((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.
                episodes, env_runner_results = synchronous_parallel_sample(
                    worker_set=self.env_runner_group,
                    max_agent_steps=(
                        self.config.rollout_fragment_length
                        * self.config.num_envs_per_env_runner
                    ),
                    sample_timeout_s=self.config.sample_timeout_s,
                    _uses_new_env_runners=True,
                    _return_metrics=True,
                )
                self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS)
                # 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=episodes)
                have_sampled = True

                # We took B x T env steps.
                env_steps_last_regular_sample = sum(len(eps) for eps in 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.metrics.peek(
                        (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME),
                        default=0,
                    )
                    == 0
                ):
                    _episodes, _env_runner_results = synchronous_parallel_sample(
                        worker_set=self.env_runner_group,
                        max_agent_steps=(
                            self.config.batch_size_B * self.config.batch_length_T
                            - env_steps_last_regular_sample
                        ),
                        sample_timeout_s=self.config.sample_timeout_s,
                        random_actions=True,
                        _uses_new_env_runners=True,
                        _return_metrics=True,
                    )
                    self.metrics.aggregate(_env_runner_results, key=ENV_RUNNER_RESULTS)
                    self.replay_buffer.add(episodes=_episodes)
                    total_sampled += sum(len(eps) for eps in _episodes)

        # Summarize environment interaction and buffer data.
        report_sampling_and_replay_buffer(
            metrics=self.metrics, replay_buffer=self.replay_buffer
        )
        # Get the replay buffer metrics.
        replay_buffer_results = self.local_replay_buffer.get_metrics()
        self.metrics.aggregate([replay_buffer_results], key=REPLAY_BUFFER_RESULTS)

        # Use self.spaces for the environment spaces of the env-runners
        single_observation_space, single_action_space = self.spaces[
            INPUT_ENV_SINGLE_SPACES
        ]

        # 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.metrics.log_time((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(single_action_space, gym.spaces.Discrete):
                    sample["actions_ints"] = sample[Columns.ACTIONS]
                    sample[Columns.ACTIONS] = one_hot(
                        sample["actions_ints"],
                        depth=single_action_space.n,
                    )

                # Perform the actual update via our learner group.
                learner_results = self.learner_group.update(
                    batch=SampleBatch(sample).as_multi_agent(),
                    # TODO(sven): Maybe we should do this broadcase of global timesteps
                    #  at the end, like for EnvRunner global env step counts. Maybe when
                    #  we request the state from the Learners, we can - at the same
                    #  time - send the current globally summed/reduced-timesteps.
                    timesteps={
                        NUM_ENV_STEPS_SAMPLED_LIFETIME: self.metrics.peek(
                            (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME),
                            default=0,
                        )
                    },
                )
                self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)

                sub_iter += 1
                self.metrics.log_value(
                    NUM_GRAD_UPDATES_LIFETIME, 1, reduce="lifetime_sum"
                )

        # Log videos showing how the decoder produces observation predictions
        # from the posterior states.
        # Only every n iterations and only for the first sampled batch row
        # (videos are `config.batch_length_T` frames long).
        report_predicted_vs_sampled_obs(
            # TODO (sven): DreamerV3 is single-agent only.
            metrics=self.metrics,
            sample=sample,
            batch_size_B=self.config.batch_size_B,
            batch_length_T=self.config.batch_length_T,
            symlog_obs=do_symlog_obs(
                single_observation_space,
                self.config.symlog_obs,
            ),
            do_report=(
                self.config.report_images_and_videos
                and self.training_iteration % 100 == 0
            ),
        )

        # Log videos showing some of the dreamed trajectories and compare them with the
        # actual trajectories from the train batch.
        # Only every n iterations and only for the first sampled batch row AND first ts.
        # (videos are `config.horizon_H` frames long originating from the observation
        # at B=0 and T=0 in the train batch).
        report_dreamed_eval_trajectory_vs_samples(
            metrics=self.metrics,
            sample=sample,
            burn_in_T=0,
            dreamed_T=self.config.horizon_H + 1,
            dreamer_model=self.env_runner.module.dreamer_model,
            symlog_obs=do_symlog_obs(
                single_observation_space,
                self.config.symlog_obs,
            ),
            do_report=(
                self.config.report_dream_data and self.training_iteration % 100 == 0
            ),
            framework=self.config.framework_str,
        )

        # Update weights - after learning on the LearnerGroup - on all EnvRunner
        # workers.
        with self.metrics.log_time((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.metrics.log_value(NUM_SYNCH_WORKER_WEIGHTS, 1, reduce="sum")
            self.env_runner_group.sync_weights(
                from_worker_or_learner_group=self.learner_group,
                inference_only=True,
            )

        # Add train results and the actual training ratio to stats. The latter should
        # be close to the configured `training_ratio`.
        self.metrics.log_value("actual_training_ratio", self.training_ratio, window=1)

    @property
    def training_ratio(self) -> float:
        """Returns the actual training ratio of this Algorithm (not the configured one).

        The training ratio is copmuted by dividing the total number of steps
        trained thus far (replayed from the buffer) over the total number of actual
        env steps taken thus far.
        """
        eps = 0.0001
        return self.metrics.peek(NUM_ENV_STEPS_TRAINED_LIFETIME, default=0) / (
            (
                self.metrics.peek(
                    (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME),
                    default=eps,
                )
                or eps
            )
        )

    # TODO (sven): Remove this once DreamerV3 is on the new SingleAgentEnvRunner.
    @PublicAPI
    def __setstate__(self, state) -> None:
        """Sts the algorithm to the provided state

        Args:
            state: The state dictionary to restore this `DreamerV3` instance to.
                `state` may have been returned by a call to an `Algorithm`'s
                `__getstate__()` method.
        """
        # Call the `Algorithm`'s `__setstate__()` method.
        super().__setstate__(state=state)

        # Assign the module to the local `EnvRunner` if sharing is enabled.
        # Note, in `Learner.restore_from_path()` the module is first deleted
        # and then a new one is built - therefore the worker has no
        # longer a copy of the learner.
        if self.config.share_module_between_env_runner_and_learner:
            assert id(self.env_runner.module) != id(
                self.learner_group._learner.module[DEFAULT_MODULE_ID]
            )
            self.env_runner.module = self.learner_group._learner.module[
                DEFAULT_MODULE_ID
            ]