mirrors/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
f1825c8ceba00fb3bf68d3c8d9de991c043b3a68
ray-project--ray/rllib/algorithms/dreamerv3/dreamerv3.py
T
wehub-resource-sync f1825c8ceb chore: import upstream snapshot with attribution
2026-07-13 13:17:40 +08:00

733 lines
33 KiB
Python
Raw Blame History

"""
[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
]
Reference in New Issue View Git Blame Copy Permalink