diff --git a/lib/python3.12/site-packages/ray/rllib/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1de2fc6e36a108c8d7cfb87b1dabaa50c153d5f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f7e0696a0d324fa2da15e1b2a1faa7fdab359407 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/__init__.py @@ -0,0 +1,38 @@ +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.algorithms.appo.appo import APPO, APPOConfig +from ray.rllib.algorithms.bc.bc import BC, BCConfig +from ray.rllib.algorithms.cql.cql import CQL, CQLConfig +from ray.rllib.algorithms.dqn.dqn import DQN, DQNConfig +from ray.rllib.algorithms.impala.impala import ( + IMPALA, + Impala, + IMPALAConfig, + ImpalaConfig, +) +from ray.rllib.algorithms.marwil.marwil import MARWIL, MARWILConfig +from ray.rllib.algorithms.ppo.ppo import PPO, PPOConfig +from ray.rllib.algorithms.sac.sac import SAC, SACConfig + +__all__ = [ + "Algorithm", + "AlgorithmConfig", + "APPO", + "APPOConfig", + "BC", + "BCConfig", + "CQL", + "CQLConfig", + "DQN", + "DQNConfig", + "IMPALA", + "IMPALAConfig", + "Impala", + "ImpalaConfig", + "MARWIL", + "MARWILConfig", + "PPO", + "PPOConfig", + "SAC", + "SACConfig", +] diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm.py b/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm.py new file mode 100644 index 0000000000000000000000000000000000000000..73cd7445d8adf6cce79e88e4ece4143d80498203 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm.py @@ -0,0 +1,4813 @@ +import concurrent +import copy +import functools +import importlib +import importlib.metadata +import json +import logging +import os +import pathlib +import re +import tempfile +import time +from collections import defaultdict +from datetime import datetime +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + DefaultDict, + Dict, + List, + Optional, + Set, + Tuple, + Type, + Union, +) + +import gymnasium as gym +import numpy as np +import pyarrow.fs +import tree # pip install dm_tree +from packaging import version + +import ray +import ray.cloudpickle as pickle +from ray._common.deprecation import ( + DEPRECATED_VALUE, + Deprecated, + deprecation_warning, +) +from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag +from ray.actor import ActorHandle +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.algorithms.registry import ALGORITHMS_CLASS_TO_NAME as ALL_ALGORITHMS +from ray.rllib.algorithms.utils import ( + AggregatorActor, + _get_env_runner_bundles, + _get_learner_bundles, + _get_main_process_bundle, + _get_offline_eval_runner_bundles, +) +from ray.rllib.callbacks.utils import make_callback +from ray.rllib.connectors.agent.obs_preproc import ObsPreprocessorConnector +from ray.rllib.connectors.connector_pipeline_v2 import ConnectorPipelineV2 +from ray.rllib.core import ( + COMPONENT_ENV_RUNNER, + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_EVAL_ENV_RUNNER, + COMPONENT_LEARNER, + COMPONENT_LEARNER_GROUP, + COMPONENT_METRICS_LOGGER, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + COMPONENT_RL_MODULE, + DEFAULT_MODULE_ID, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module import validate_module_id +from ray.rllib.core.rl_module.multi_rl_module import ( + MultiRLModule, + MultiRLModuleSpec, +) +from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec +from ray.rllib.env import INPUT_ENV_SPACES +from ray.rllib.env.env_context import EnvContext +from ray.rllib.env.env_runner import EnvRunner +from ray.rllib.env.env_runner_group import EnvRunnerGroup +from ray.rllib.env.utils import _gym_env_creator +from ray.rllib.evaluation.metrics import ( + collect_episodes, + summarize_episodes, +) +from ray.rllib.execution.rollout_ops import synchronous_parallel_sample +from ray.rllib.offline import get_dataset_and_shards +from ray.rllib.offline.estimators import ( + DirectMethod, + DoublyRobust, + ImportanceSampling, + OffPolicyEstimator, + WeightedImportanceSampling, +) +from ray.rllib.offline.offline_evaluator import OfflineEvaluator +from ray.rllib.policy.policy import Policy, PolicySpec +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch +from ray.rllib.utils import FilterManager, deep_update, force_list +from ray.rllib.utils.actor_manager import FaultTolerantActorManager +from ray.rllib.utils.annotations import ( + DeveloperAPI, + ExperimentalAPI, + OldAPIStack, + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, + PublicAPI, + override, +) +from ray.rllib.utils.checkpoints import ( + CHECKPOINT_VERSION, + CHECKPOINT_VERSION_LEARNER_AND_ENV_RUNNER, + Checkpointable, + get_checkpoint_info, + try_import_msgpack, +) +from ray.rllib.utils.debug import update_global_seed_if_necessary +from ray.rllib.utils.error import ERR_MSG_INVALID_ENV_DESCRIPTOR, EnvError +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.metrics import ( + AGGREGATOR_ACTOR_RESULTS, + ALL_MODULES, + DATASET_NUM_ITERS_EVALUATED, + ENV_RUNNER_RESULTS, + ENV_RUNNER_SAMPLING_TIMER, + EPISODE_LEN_MEAN, + EPISODE_RETURN_MEAN, + EVALUATION_ITERATION_TIMER, + EVALUATION_RESULTS, + FAULT_TOLERANCE_STATS, + LEARNER_RESULTS, + LEARNER_UPDATE_TIMER, + NUM_AGENT_STEPS_SAMPLED, + NUM_AGENT_STEPS_SAMPLED_LIFETIME, + NUM_AGENT_STEPS_SAMPLED_THIS_ITER, + NUM_AGENT_STEPS_TRAINED, + NUM_AGENT_STEPS_TRAINED_LIFETIME, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_FOR_EVALUATION_THIS_ITER, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_ENV_STEPS_SAMPLED_THIS_ITER, + NUM_ENV_STEPS_TRAINED, + NUM_ENV_STEPS_TRAINED_LIFETIME, + NUM_EPISODES, + NUM_EPISODES_LIFETIME, + NUM_TRAINING_STEP_CALLS_PER_ITERATION, + OFFLINE_EVAL_RUNNER_RESULTS, + OFFLINE_EVALUATION_ITERATION_TIMER, + RESTORE_ENV_RUNNERS_TIMER, + RESTORE_EVAL_ENV_RUNNERS_TIMER, + RESTORE_OFFLINE_EVAL_RUNNERS_TIMER, + STEPS_TRAINED_THIS_ITER_COUNTER, + SYNCH_ENV_CONNECTOR_STATES_TIMER, + SYNCH_EVAL_ENV_CONNECTOR_STATES_TIMER, + SYNCH_WORKER_WEIGHTS_TIMER, + TIMERS, + TRAINING_ITERATION_TIMER, + TRAINING_STEP_TIMER, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_INFO +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.metrics.ray_metrics import ( + DEFAULT_HISTOGRAM_BOUNDARIES_LONG_EVENTS, + DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + TimerAndPrometheusLogger, +) +from ray.rllib.utils.replay_buffers import MultiAgentReplayBuffer, ReplayBuffer +from ray.rllib.utils.runners.runner_group import RunnerGroup +from ray.rllib.utils.serialization import NOT_SERIALIZABLE, deserialize_type +from ray.rllib.utils.spaces import space_utils +from ray.rllib.utils.typing import ( + AgentConnectorDataType, + AgentID, + AgentToModuleMappingFn, + AlgorithmConfigDict, + EnvCreator, + EnvInfoDict, + EnvType, + EpisodeID, + ModuleID, + PartialAlgorithmConfigDict, + PolicyID, + PolicyState, + ResultDict, + SampleBatchType, + ShouldModuleBeUpdatedFn, + StateDict, + TensorStructType, + TensorType, +) +from ray.train.constants import DEFAULT_STORAGE_PATH +from ray.tune import Checkpoint +from ray.tune.execution.placement_groups import PlacementGroupFactory +from ray.tune.experiment.trial import ExportFormat +from ray.tune.logger import Logger, UnifiedLogger +from ray.tune.registry import ENV_CREATOR, _global_registry, get_trainable_cls +from ray.tune.resources import Resources +from ray.tune.result import TRAINING_ITERATION +from ray.tune.trainable import Trainable +from ray.util import log_once +from ray.util.metrics import Counter, Histogram +from ray.util.timer import _Timer + +if TYPE_CHECKING: + from ray.rllib.core.learner.learner_group import LearnerGroup + from ray.rllib.offline.offline_data import OfflineData + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +@PublicAPI +class Algorithm(Checkpointable, Trainable): + """An RLlib algorithm responsible for training one or more neural network models. + + You can write your own Algorithm classes by sub-classing from `Algorithm` + or any of its built-in subclasses. + Override the `training_step` method to implement your own algorithm logic. + Find the various built-in `training_step()` methods for different algorithms in + their respective [algo name].py files, for example: + `ray.rllib.algorithms.dqn.dqn.py` or `ray.rllib.algorithms.impala.impala.py`. + + The most important API methods an Algorithm exposes are `train()` for running a + single training iteration, `evaluate()` for running a single round of evaluation, + `save_to_path()` for creating a checkpoint, and `restore_from_path()` for loading a + state from an existing checkpoint. + """ + + #: The AlgorithmConfig instance of the Algorithm. + config: Optional[AlgorithmConfig] = None + #: The MetricsLogger instance of the Algorithm. RLlib uses this to log + #: metrics from within the `training_step()` method. Users can use it to log + #: metrics from within their custom Algorithm-based callbacks. + metrics: Optional[MetricsLogger] = None + #: The `EnvRunnerGroup` of the Algorithm. An `EnvRunnerGroup` is + #: composed of a single local `EnvRunner` (see: `self.env_runner`), serving as + #: the reference copy of the models to be trained and optionally one or more + #: remote `EnvRunners` used to generate training samples from the RL + #: environment, in parallel. EnvRunnerGroup is fault-tolerant and elastic. It + #: tracks health states for all the managed remote EnvRunner actors. As a + #: result, Algorithm should never access the underlying actor handles directly. + #: Instead, always access them via all the foreach APIs with assigned IDs of + #: the underlying EnvRunners. + env_runner_group: Optional[EnvRunnerGroup] = None + #: A special EnvRunnerGroup only used for evaluation, not to + #: collect training samples. + eval_env_runner_group: Optional[EnvRunnerGroup] = None + #: The `LearnerGroup` instance of the Algorithm, managing either + #: one local `Learner` or one or more remote `Learner` actors. Responsible for + #: updating the models from RL environment (episode) data. + learner_group: Optional["LearnerGroup"] = None + #: An optional OfflineData instance, used for offline RL. + offline_data: Optional["OfflineData"] = None + + # Whether to allow unknown top-level config keys. + _allow_unknown_configs = False + + # List of top-level keys with value=dict, for which new sub-keys are + # allowed to be added to the value dict. + _allow_unknown_subkeys = [ + "tf_session_args", + "local_tf_session_args", + "env_config", + "model", + "optimizer", + "custom_resources_per_env_runner", + "custom_resources_per_worker", + "evaluation_config", + "exploration_config", + "replay_buffer_config", + "extra_python_environs_for_worker", + "input_config", + "output_config", + ] + + # List of top level keys with value=dict, for which we always override the + # entire value (dict), iff the "type" key in that value dict changes. + _override_all_subkeys_if_type_changes = [ + "exploration_config", + "replay_buffer_config", + ] + + # List of keys that are always fully overridden if present in any dict or sub-dict + _override_all_key_list = ["off_policy_estimation_methods", "policies"] + + _progress_metrics = ( + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}", + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}", + f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}", + f"{NUM_ENV_STEPS_TRAINED_LIFETIME}", + f"{NUM_EPISODES_LIFETIME}", + f"{ENV_RUNNER_RESULTS}/{EPISODE_LEN_MEAN}", + ) + + # Backward compatibility with old checkpoint system (now through the + # `Checkpointable` API). + METADATA_FILE_NAME = "rllib_checkpoint.json" + STATE_FILE_NAME = "algorithm_state" + + @classmethod + @override(Checkpointable) + def from_checkpoint( + cls, + path: Union[str, Checkpoint], + filesystem: Optional["pyarrow.fs.FileSystem"] = None, + *, + # @OldAPIStack + policy_ids: Optional[Collection[PolicyID]] = None, + policy_mapping_fn: Optional[Callable[[AgentID, EpisodeID], PolicyID]] = None, + policies_to_train: Optional[ + Union[ + Collection[PolicyID], + Callable[[PolicyID, Optional[SampleBatchType]], bool], + ] + ] = None, + # deprecated args + checkpoint=DEPRECATED_VALUE, + **kwargs, + ) -> "Algorithm": + """Creates a new algorithm instance from a given checkpoint. + + Args: + path: The path (str) to the checkpoint directory to use or a Ray Train + Checkpoint instance to restore from. + filesystem: PyArrow FileSystem to use to access data at the `path`. If not + specified, this is inferred from the URI scheme of `path`. + policy_ids: Optional list of PolicyIDs to recover. This allows users to + restore an Algorithm with only a subset of the originally present + Policies. + policy_mapping_fn: An optional (updated) policy mapping function to use from + here on. + policies_to_train: An optional list of policy IDs to be trained or a + callable taking PolicyID and SampleBatchType and returning a bool + (trainable or not?). If None, will keep the existing setup in place. + Policies, whose IDs are not in the list (or for which the callable + returns False) will not be updated. + + Returns: + The instantiated Algorithm. + """ + if checkpoint != DEPRECATED_VALUE: + deprecation_warning( + old="Algorithm.from_checkpoint(checkpoint=...)", + new="Algorithm.from_checkpoint(path=...)", + error=True, + ) + checkpoint_info = get_checkpoint_info(path, filesystem) + + # New API stack -> Use Checkpointable's default implementation. + if checkpoint_info["checkpoint_version"] >= version.Version("2.0"): + # `path` is a Checkpoint instance: Translate to directory and continue. + if isinstance(path, Checkpoint): + path = path.to_directory() + return super().from_checkpoint(path, filesystem=filesystem, **kwargs) + + # Not possible for (v0.1) (algo class and config information missing + # or very hard to retrieve). + elif checkpoint_info["checkpoint_version"] == version.Version("0.1"): + raise ValueError( + "Cannot restore a v0 checkpoint using `Algorithm.from_checkpoint()`!" + "In this case, do the following:\n" + "1) Create a new Algorithm object using your original config.\n" + "2) Call the `restore()` method of this algo object passing it" + " your checkpoint dir or AIR Checkpoint object." + ) + elif checkpoint_info["checkpoint_version"] < version.Version("1.0"): + raise ValueError( + "`checkpoint_info['checkpoint_version']` in `Algorithm.from_checkpoint" + "()` must be 1.0 or later! You are using a checkpoint with " + f"version v{checkpoint_info['checkpoint_version']}." + ) + + # This is a msgpack checkpoint. + if checkpoint_info["format"] == "msgpack": + # User did not provide unserializable function with this call + # (`policy_mapping_fn`). Note that if `policies_to_train` is None, it + # defaults to training all policies (so it's ok to not provide this here). + if policy_mapping_fn is None: + # Only DEFAULT_POLICY_ID present in this algorithm, provide default + # implementations of these two functions. + if checkpoint_info["policy_ids"] == {DEFAULT_POLICY_ID}: + policy_mapping_fn = AlgorithmConfig.DEFAULT_POLICY_MAPPING_FN + # Provide meaningful error message. + else: + raise ValueError( + "You are trying to restore a multi-agent algorithm from a " + "`msgpack` formatted checkpoint, which do NOT store the " + "`policy_mapping_fn` or `policies_to_train` " + "functions! Make sure that when using the " + "`Algorithm.from_checkpoint()` utility, you also pass the " + "args: `policy_mapping_fn` and `policies_to_train` with your " + "call. You might leave `policies_to_train=None` in case " + "you would like to train all policies anyways." + ) + + state = Algorithm._checkpoint_info_to_algorithm_state( + checkpoint_info=checkpoint_info, + policy_ids=policy_ids, + policy_mapping_fn=policy_mapping_fn, + policies_to_train=policies_to_train, + ) + + return Algorithm.from_state(state) + + @PublicAPI + def __init__( + self, + config: Optional[AlgorithmConfig] = None, + env=None, # deprecated arg + logger_creator: Optional[Callable[[], Logger]] = None, + **kwargs, + ): + """Initializes an Algorithm instance. + + Args: + config: Algorithm-specific configuration object. + logger_creator: Callable that creates a ray.tune.Logger + object. If unspecified, a default logger is created. + **kwargs: Arguments passed to the Trainable base class. + """ + # Translate possible dict into an AlgorithmConfig object, as well as, + # resolving generic config objects into specific ones (e.g. passing + # an `AlgorithmConfig` super-class instance into a PPO constructor, + # which normally would expect a PPOConfig object). + if isinstance(config, dict): + default_config = self.get_default_config() + # `self.get_default_config()` also returned a dict -> + # Last resort: Create core AlgorithmConfig from merged dicts. + if isinstance(default_config, dict): + if "class" in config: + AlgorithmConfig.from_state(config) + else: + config = AlgorithmConfig.from_dict( + config_dict=self.merge_algorithm_configs( + default_config, config, True + ) + ) + + # Default config is an AlgorithmConfig -> update its properties + # from the given config dict. + else: + if isinstance(config, dict) and "class" in config: + config = default_config.from_state(config) + else: + config = default_config.update_from_dict(config) + else: + default_config = self.get_default_config() + # Given AlgorithmConfig is not of the same type as the default config: + # This could be the case e.g. if the user is building an algo from a + # generic AlgorithmConfig() object. + if not isinstance(config, type(default_config)): + config = default_config.update_from_dict(config.to_dict()) + else: + config = default_config.from_state(config.get_state()) + + # In case this algo is using a generic config (with no algo_class set), set it + # here. + if config.algo_class is None: + config.algo_class = type(self) + + if env is not None: + deprecation_warning( + old=f"algo = Algorithm(env='{env}', ...)", + new=f"algo = AlgorithmConfig().environment('{env}').build()", + error=False, + ) + config.environment(env) + + # Validate and freeze our AlgorithmConfig object (no more changes possible). + config.validate() + config.freeze() + + # Convert `env` provided in config into a concrete env creator callable, which + # takes an EnvContext (config dict) as arg and returning an RLlib supported Env + # type (e.g. a gym.Env). + self._env_id, self.env_creator = self._get_env_id_and_creator( + config.env, config + ) + env_descr = ( + self._env_id.__name__ if isinstance(self._env_id, type) else self._env_id + ) + + # Placeholder for a local replay buffer instance. + self.local_replay_buffer = None + + # Placeholder for our LearnerGroup responsible for updating the RLModule(s). + self.learner_group: Optional["LearnerGroup"] = None + + # The Algorithm's `MetricsLogger` object to collect stats from all its + # components (including timers, counters and other stats in its own + # `training_step()` and other methods) as well as custom callbacks. + self.metrics = MetricsLogger(root=True) + + # Create a default logger creator if no logger_creator is specified + if logger_creator is None: + # Default logdir prefix containing the agent's name and the + # env id. + timestr = datetime.today().strftime("%Y-%m-%d_%H-%M-%S") + env_descr_for_dir = re.sub("[/\\\\]", "-", str(env_descr)) + logdir_prefix = f"{type(self).__name__}_{env_descr_for_dir}_{timestr}" + + if not os.path.exists(DEFAULT_STORAGE_PATH): + # Possible race condition if dir is created several times on + # rollout workers + os.makedirs(DEFAULT_STORAGE_PATH, exist_ok=True) + logdir = tempfile.mkdtemp(prefix=logdir_prefix, dir=DEFAULT_STORAGE_PATH) + + # Allow users to more precisely configure the created logger + # via "logger_config.type". + if config.logger_config and "type" in config.logger_config: + + def default_logger_creator(config): + """Creates a custom logger with the default prefix.""" + cfg = config["logger_config"].copy() + cls = cfg.pop("type") + # Provide default for logdir, in case the user does + # not specify this in the "logger_config" dict. + logdir_ = cfg.pop("logdir", logdir) + return from_config(cls=cls, _args=[cfg], logdir=logdir_) + + # If no `type` given, use tune's UnifiedLogger as last resort. + else: + + def default_logger_creator(config): + """Creates a Unified logger with the default prefix.""" + return UnifiedLogger(config, logdir, loggers=None) + + logger_creator = default_logger_creator + + # Metrics-related properties. + self._timers = defaultdict(_Timer) + self._counters = defaultdict(int) + self._episode_history = [] + self._episodes_to_be_collected = [] + + # The fully qualified AlgorithmConfig used for evaluation + # (or None if evaluation not setup). + self.evaluation_config: Optional[AlgorithmConfig] = None + # Evaluation EnvRunnerGroup and metrics last returned by `self.evaluate()`. + self.eval_env_runner_group: Optional[EnvRunnerGroup] = None + + # Ray metrics - Algorithm + self._metrics_step_time: Optional[Histogram] = None + self._metrics_run_one_training_iteration_time: Optional[Histogram] = None + self._metrics_run_one_evaluation_time: Optional[Histogram] = None + self._metrics_compile_iteration_results_time: Optional[Histogram] = None + self._metrics_training_step_time: Optional[Histogram] = None + self._metrics_evaluate_time: Optional[Histogram] = None + self._metrics_evaluate_sync_env_runner_weights_time: Optional[Histogram] = None + self._metrics_evaluate_sync_connector_states_time: Optional[Histogram] = None + self._metrics_step_sync_env_runner_states_time: Optional[Histogram] = None + self._metrics_load_checkpoint_time: Optional[Histogram] = None + self._metrics_save_checkpoint_time: Optional[Histogram] = None + + # Ray metrics - Algorithm callbacks + self._metrics_callback_on_train_result_time: Optional[Histogram] = None + self._metrics_callback_on_evaluate_start_time: Optional[Histogram] = None + self._metrics_callback_on_evaluate_end_time: Optional[Histogram] = None + + # Ray metrics - IMPALA + self._metrics_impala_training_step_time: Optional[Histogram] = None + self._metrics_impala_training_step_aggregator_preprocessing_time: Optional[ + Histogram + ] = None + self._metrics_impala_training_step_learner_group_loop_time: Optional[ + Histogram + ] = None + self._metrics_impala_training_step_sync_env_runner_state_time: Optional[ + Histogram + ] = None + self._metrics_impala_sample_and_get_connector_states_time: Optional[ + Histogram + ] = None + self._metrics_impala_training_step_input_batches: Optional[Counter] = None + self._metrics_impala_training_step_zero_input_batches: Optional[Counter] = None + self._metrics_impala_training_step_env_steps_dropped: Optional[Counter] = None + + super().__init__( + config=config, + logger_creator=logger_creator, + **kwargs, + ) + + def _set_up_metrics(self): + self._metrics_step_time = Histogram( + name="rllib_algorithm_step_time", + description="Time spent in Algorithm.step()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_LONG_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_step_time.set_default_tags({"rllib": self.__class__.__name__}) + + self._metrics_run_one_training_iteration_time = Histogram( + name="rllib_algorithm_run_one_training_iteration_time", + description="Time spent in Algorithm._run_one_training_iteration()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_LONG_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_run_one_training_iteration_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_run_one_evaluation_time = Histogram( + name="rllib_algorithm_run_one_evaluation_time", + description="Time spent in Algorithm._run_one_evaluation()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_LONG_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_run_one_evaluation_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_compile_iteration_results_time = Histogram( + name="rllib_algorithm_compile_iteration_results_time", + description="Time spent in Algorithm._compile_iteration_results()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_compile_iteration_results_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_training_step_time = Histogram( + name="rllib_algorithm_training_step_time", + description="Time spent in Algorithm.training_step()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_LONG_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_training_step_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_evaluate_time = Histogram( + name="rllib_algorithm_evaluate_time", + description="Time spent in Algorithm.evaluate()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_LONG_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_evaluate_time.set_default_tags({"rllib": self.__class__.__name__}) + + self._metrics_evaluate_sync_env_runner_weights_time = Histogram( + name="rllib_algorithm_evaluate_sync_env_runner_weights_time", + description="Time spent on syncing weights to the eval EnvRunners in the Algorithm.evaluate()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_evaluate_sync_env_runner_weights_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_evaluate_sync_connector_states_time = Histogram( + name="rllib_algorithm_evaluate_sync_connector_states_time", + description="Time spent on syncing connector states in the Algorithm.evaluate()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_evaluate_sync_connector_states_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_step_sync_env_runner_states_time = Histogram( + name="rllib_algorithm_step_sync_env_runner_states_time", + description="Time spent in sync_env_runner_states code block of the Algorithm.step()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_step_sync_env_runner_states_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_load_checkpoint_time = Histogram( + name="rllib_algorithm_load_checkpoint_time", + description="Time spent in Algorithm.load_checkpoint()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_load_checkpoint_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_save_checkpoint_time = Histogram( + name="rllib_algorithm_save_checkpoint_time", + description="Time spent in Algorithm.save_checkpoint()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_save_checkpoint_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + # Ray metrics - Algorithm callbacks + self._metrics_callback_on_train_result_time = Histogram( + name="rllib_algorithm_callback_on_train_result_time", + description="Time spent in callback 'on_train_result()'", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_callback_on_train_result_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_callback_on_evaluate_start_time = Histogram( + name="rllib_algorithm_callback_on_evaluate_start_time", + description="Time spent in callback 'on_evaluate_start()'", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_callback_on_evaluate_start_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_callback_on_evaluate_end_time = Histogram( + name="rllib_algorithm_callback_on_evaluate_end_time", + description="Time spent in callback 'on_evaluate_end()'", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_callback_on_evaluate_end_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + @OverrideToImplementCustomLogic + @classmethod + def get_default_config(cls) -> AlgorithmConfig: + return AlgorithmConfig() + + @OverrideToImplementCustomLogic + def _remote_worker_ids_for_metrics(self) -> List[int]: + """Returns a list of remote worker IDs to fetch metrics from. + + Specific Algorithm implementations can override this method to + use a subset of the workers for metrics collection. + + Returns: + List of remote worker IDs to fetch metrics from. + """ + return self.env_runner_group.healthy_worker_ids() + + @OverrideToImplementCustomLogic_CallToSuperRecommended + @override(Trainable) + def setup(self, config: AlgorithmConfig) -> None: + # Setup our config: Merge the user-supplied config dict (which could + # be a partial config dict) with the class' default. + if not isinstance(config, AlgorithmConfig): + assert isinstance(config, PartialAlgorithmConfigDict) + config_obj = self.get_default_config() + if not isinstance(config_obj, AlgorithmConfig): + assert isinstance(config, PartialAlgorithmConfigDict) + config_obj = AlgorithmConfig().from_dict(config_obj) + config_obj.update_from_dict(config) + config_obj.env = self._env_id + self.config = config_obj + + # Set Algorithm's seed after we have - if necessary - enabled + # tf eager-execution. + update_global_seed_if_necessary(self.config.framework_str, self.config.seed) + + self._record_usage(self.config) + + # Create the callbacks object. + if self.config.enable_env_runner_and_connector_v2: + self.callbacks = [cls() for cls in force_list(self.config.callbacks_class)] + else: + self.callbacks = self.config.callbacks_class() + + if self.config.log_level in ["WARN", "ERROR"]: + logger.info( + f"Current log_level is {self.config.log_level}. For more information, " + "set 'log_level': 'INFO' / 'DEBUG' or use the -v and " + "-vv flags." + ) + if self.config.log_level: + logging.getLogger("ray.rllib").setLevel(self.config.log_level) + + # Create local replay buffer if necessary. + self.local_replay_buffer = self._create_local_replay_buffer_if_necessary( + self.config + ) + + # Create a dict, mapping ActorHandles to sets of open remote + # requests (object refs). This way, we keep track, of which actors + # inside this Algorithm (e.g. a remote EnvRunner) have + # already been sent how many (e.g. `sample()`) requests. + self.remote_requests_in_flight: DefaultDict[ + ActorHandle, Set[ray.ObjectRef] + ] = defaultdict(set) + + self.env_runner_group: Optional[EnvRunnerGroup] = None + # In case there is no local EnvRunner anymore, we need to handle connector + # pipelines directly here. + self.spaces: Optional[Dict] = None + self.env_to_module_connector: Optional[ConnectorPipelineV2] = None + self.module_to_env_connector: Optional[ConnectorPipelineV2] = None + + # Offline RL settings. + input_evaluation = self.config.get("input_evaluation") + if input_evaluation is not None and input_evaluation is not DEPRECATED_VALUE: + ope_dict = {str(ope): {"type": ope} for ope in input_evaluation} + deprecation_warning( + old="config.input_evaluation={}".format(input_evaluation), + new="config.evaluation(evaluation_config=config.overrides(" + f"off_policy_estimation_methods={ope_dict}" + "))", + error=True, + help="Running OPE during training is not recommended.", + ) + self.config.off_policy_estimation_methods = ope_dict + + # If an input path is available and we are on the new API stack generate + # an `OfflineData` instance. + if self.config.is_offline: + from ray.rllib.offline.offline_data import OfflineData + + # Use either user-provided `OfflineData` class or RLlib's default. + offline_data_class = self.config.offline_data_class or OfflineData + # Build the `OfflineData` class. + self.offline_data = offline_data_class(self.config) + # Otherwise set the attribute to `None`. + else: + self.offline_data = None + + if self.config.is_online or not self.config.enable_env_runner_and_connector_v2: + # Create a set of env runner actors via a EnvRunnerGroup. + self.env_runner_group = EnvRunnerGroup( + env_creator=self.env_creator, + validate_env=self.validate_env, + default_policy_class=self.get_default_policy_class(self.config), + config=self.config, + # New API stack: User decides whether to create local env runner. + # Old API stack: Always create local EnvRunner. + local_env_runner=( + True + if not self.config.enable_env_runner_and_connector_v2 + else self.config.create_local_env_runner + ), + logdir=self.logdir, + tune_trial_id=self.trial_id, + ) + + # Compile, validate, and freeze an evaluation config. + self.evaluation_config = self.config.get_evaluation_config_object() + self.evaluation_config.validate() + self.evaluation_config.freeze() + + # Evaluation EnvRunnerGroup setup. + # User would like to setup a separate evaluation worker set. + # Note: We skip EnvRunnerGroup creation if we need to do offline evaluation. + if self._should_create_evaluation_env_runners(self.evaluation_config): + _, env_creator = self._get_env_id_and_creator( + self.evaluation_config.env, self.evaluation_config + ) + + # Create a separate evaluation worker set for evaluation. + # If evaluation_num_env_runners=0, use the evaluation set's local + # worker for evaluation, otherwise, use its remote workers + # (parallelized evaluation). + self.eval_env_runner_group: EnvRunnerGroup = EnvRunnerGroup( + env_creator=env_creator, + validate_env=None, + default_policy_class=self.get_default_policy_class(self.config), + config=self.evaluation_config, + logdir=self.logdir, + tune_trial_id=self.trial_id, + # New API stack: User decides whether to create local env runner. + # Old API stack: Always create local EnvRunner. + local_env_runner=( + True + if not self.evaluation_config.enable_env_runner_and_connector_v2 + else self.evaluation_config.create_local_env_runner + ), + pg_offset=self.config.num_env_runners, + ) + + if self.env_runner_group: + self.spaces = self.env_runner_group.get_spaces() + elif self.eval_env_runner_group: + self.spaces = self.eval_env_runner_group.get_spaces() + + if self.env_runner is None and self.spaces is not None: + self.env_to_module_connector = self.config.build_env_to_module_connector( + spaces=self.spaces + ) + self.module_to_env_connector = self.config.build_module_to_env_connector( + spaces=self.spaces + ) + + self.evaluation_dataset = None + if ( + self.evaluation_config.off_policy_estimation_methods + and not self.evaluation_config.ope_split_batch_by_episode + ): + # the num worker is set to 0 to avoid creating shards. The dataset will not + # be repartioned to num_workers blocks. + logger.info("Creating evaluation dataset ...") + self.evaluation_dataset, _ = get_dataset_and_shards( + self.evaluation_config, num_workers=0 + ) + logger.info("Evaluation dataset created") + + self.reward_estimators: Dict[str, OffPolicyEstimator] = {} + ope_types = { + "is": ImportanceSampling, + "wis": WeightedImportanceSampling, + "dm": DirectMethod, + "dr": DoublyRobust, + } + for name, method_config in self.config.off_policy_estimation_methods.items(): + method_type = method_config.pop("type") + if method_type in ope_types: + deprecation_warning( + old=method_type, + new=str(ope_types[method_type]), + error=True, + ) + method_type = ope_types[method_type] + elif isinstance(method_type, str): + logger.log(0, "Trying to import from string: " + method_type) + mod, obj = method_type.rsplit(".", 1) + mod = importlib.import_module(mod) + method_type = getattr(mod, obj) + if isinstance(method_type, type) and issubclass( + method_type, OfflineEvaluator + ): + # TODO(kourosh) : Add an integration test for all these + # offline evaluators. + policy = self.get_policy() + if issubclass(method_type, OffPolicyEstimator): + method_config["gamma"] = self.config.gamma + self.reward_estimators[name] = method_type(policy, **method_config) + else: + raise ValueError( + f"Unknown off_policy_estimation type: {method_type}! Must be " + "either a class path or a sub-class of ray.rllib." + "offline.offline_evaluator::OfflineEvaluator" + ) + # TODO (Rohan138): Refactor this and remove deprecated methods + # Need to add back method_type in case Algorithm is restored from checkpoint + method_config["type"] = method_type + + if self.config.enable_rl_module_and_learner: + spaces = { + INPUT_ENV_SPACES: ( + self.config.observation_space, + self.config.action_space, + ) + } + if self.env_runner_group: + spaces.update(self.spaces) + elif self.eval_env_runner_group: + spaces.update(self.eval_env_runner_group.get_spaces()) + else: + # If the algorithm is online we use the spaces from as they are + # provided. + if self.config.is_online: + spaces.update( + { + DEFAULT_MODULE_ID: ( + self.config.observation_space, + self.config.action_space, + ), + } + ) + # Otherwise, when we are offline we need to check, if the learner connector + # is transforming the spaces. + elif self.config.is_offline: + # Build the learner connector with the input spaces from the environment. + learner_connector = self.config.build_learner_connector( + input_observation_space=spaces[INPUT_ENV_SPACES][0], + input_action_space=spaces[INPUT_ENV_SPACES][1], + ) + # Update the `spaces` dictionary by using the output spaces of the learner + # connector pipeline. + spaces.update( + { + DEFAULT_MODULE_ID: ( + learner_connector.observation_space, + learner_connector.action_space, + ), + } + ) + + module_spec: MultiRLModuleSpec = self.config.get_multi_rl_module_spec( + spaces=spaces, + inference_only=False, + ) + self.learner_group = self.config.build_learner_group( + rl_module_spec=module_spec + ) + + # Check if there are modules to load from the `module_spec`. + rl_module_ckpt_dirs = {} + multi_rl_module_ckpt_dir = module_spec.load_state_path + modules_to_load = module_spec.modules_to_load + for module_id, sub_module_spec in module_spec.rl_module_specs.items(): + if sub_module_spec.load_state_path: + rl_module_ckpt_dirs[module_id] = sub_module_spec.load_state_path + if multi_rl_module_ckpt_dir or rl_module_ckpt_dirs: + self.learner_group.load_module_state( + multi_rl_module_ckpt_dir=multi_rl_module_ckpt_dir, + modules_to_load=modules_to_load, + rl_module_ckpt_dirs=rl_module_ckpt_dirs, + ) + + # Sync the weights from the learner group to the EnvRunners. + rl_module_state = self.learner_group.get_state( + components=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE, + inference_only=True, + )[COMPONENT_LEARNER] + if self.env_runner_group: + self.env_runner_group.sync_env_runner_states( + config=self.config, + env_steps_sampled=self.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME), default=0 + ), + rl_module_state=rl_module_state, + env_to_module=self.env_to_module_connector, + module_to_env=self.module_to_env_connector, + ) + elif self.eval_env_runner_group: + self.eval_env_runner_group.sync_env_runner_states( + config=self.evaluation_config, + env_steps_sampled=self.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME), default=0 + ), + rl_module_state=rl_module_state, + env_to_module=self.env_to_module_connector, + module_to_env=self.module_to_env_connector, + ) + # TODO (simon): Update modules in DataWorkers. + + if self.offline_data: + # If the learners are remote we need to provide specific + # information and the learner's actor handles. + if self.learner_group.is_remote: + # If learners run on different nodes, locality hints help + # to use the nearest learner in the workers that do the + # data preprocessing. + learner_node_ids = self.learner_group.foreach_learner( + lambda _: ray.get_runtime_context().get_node_id() + ) + self.offline_data.locality_hints = [ + node_id.get() for node_id in learner_node_ids + ] + # Provide the actor handles for the learners for module + # updating during preprocessing. + self.offline_data.learner_handles = self.learner_group._workers + # Otherwise we can simply pass in the local learner. + else: + self.offline_data.learner_handles = [self.learner_group._learner] + # TODO (simon, sven): Replace these set-some-object's-attributes- + # directly? We should find some solution for this in the future, an API, + # or setting these in the OfflineData constructor? + # Provide the module_spec. Note, in the remote case this is needed + # because the learner module cannot be copied, but must be built. + self.offline_data.module_spec = module_spec + # Provide the `OfflineData` instance with space information. It might + # need it for reading recorded experiences. + self.offline_data.spaces = spaces + + if self._should_create_offline_evaluation_runners(self.evaluation_config): + from ray.rllib.offline.offline_evaluation_runner_group import ( + OfflineEvaluationRunnerGroup, + ) + + # If no inference-only `RLModule` should be used in offline evaluation, + # get the complete learner module. + if not self.evaluation_config.offline_eval_rl_module_inference_only: + rl_module_state = self.learner_group.get_state( + components=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE, + inference_only=False, + )[COMPONENT_LEARNER] + # Create the offline evaluation runner group. + self.offline_eval_runner_group: OfflineEvaluationRunnerGroup = OfflineEvaluationRunnerGroup( + config=self.evaluation_config, + # Do not create a local runner such that the dataset can be split. + local_runner=self.config.num_offline_eval_runners == 0, + # Provide the `RLModule`'s state for the `OfflinePreLearner`s. + module_state=rl_module_state[COMPONENT_RL_MODULE], + module_spec=module_spec, + # Note, even if no environment is run, the `MultiRLModule` needs + # spaces to construct the policy network. + spaces=spaces, + ) + + # Create an Aggregator actor set, if necessary. + self._aggregator_actor_manager = None + if self.config.enable_rl_module_and_learner and ( + self.config.num_aggregator_actors_per_learner > 0 + ): + rl_module_spec = self.config.get_multi_rl_module_spec( + spaces=self.spaces, + inference_only=False, + ) + agg_cls = ray.remote( + num_cpus=1, + max_restarts=-1, + )(AggregatorActor) + self._aggregator_actor_manager = FaultTolerantActorManager( + [ + agg_cls.remote(self.config, rl_module_spec) + for _ in range( + (self.config.num_learners or 1) + * self.config.num_aggregator_actors_per_learner + ) + ], + max_remote_requests_in_flight_per_actor=( + self.config.max_requests_in_flight_per_aggregator_actor + ), + ) + # Get the devices of each learner. + learner_locations = list( + enumerate( + self.learner_group.foreach_learner( + func=lambda _learner: (_learner.node, _learner.device), + ) + ) + ) + # Get the devices of each AggregatorActor. + aggregator_locations = list( + enumerate( + self._aggregator_actor_manager.foreach_actor( + func=lambda actor: (actor._node, actor._device) + ) + ) + ) + self._aggregator_actor_to_learner = {} + for agg_idx, aggregator_location in aggregator_locations: + aggregator_location = aggregator_location.get() + for learner_idx, learner_location in learner_locations: + # TODO (sven): Activate full comparison (including device) when Ray + # has figured out GPU pre-loading. + if learner_location.get()[0] == aggregator_location[0]: + # Round-robin, in case all Learners are on same device/node. + learner_locations = learner_locations[1:] + [ + learner_locations[0] + ] + self._aggregator_actor_to_learner[agg_idx] = learner_idx + break + if agg_idx not in self._aggregator_actor_to_learner: + raise RuntimeError( + "No Learner worker found that matches aggregation worker " + f"#{agg_idx}'s node ({aggregator_location[0]}) and device " + f"({aggregator_location[1]})! The Learner workers' locations " + f"are {learner_locations}." + ) + + # Make sure, each Learner index is mapped to from at least one + # AggregatorActor. + if not all( + learner_idx in self._aggregator_actor_to_learner.values() + for learner_idx in range(self.config.num_learners or 1) + ): + raise RuntimeError( + "Some Learner indices are not mapped to from any AggregatorActors! " + "Final AggregatorActor idx -> Learner idx mapping is: " + f"{self._aggregator_actor_to_learner}" + ) + + # Ray metrics + self._set_up_metrics() + + # Run `on_algorithm_init` callback after initialization is done. + make_callback( + "on_algorithm_init", + self.callbacks, + self.config.callbacks_on_algorithm_init, + kwargs=dict( + algorithm=self, + metrics_logger=self.metrics, + ), + ) + + @OverrideToImplementCustomLogic + @classmethod + def get_default_policy_class( + cls, + config: AlgorithmConfig, + ) -> Optional[Type[Policy]]: + """Returns a default Policy class to use, given a config. + + This class will be used by an Algorithm in case + the policy class is not provided by the user in any single- or + multi-agent PolicySpec. + + Note: This method is ignored when the RLModule API is enabled. + """ + return None + + @override(Trainable) + def step(self) -> ResultDict: + """Implements the main `Algorithm.train()` logic. + + Takes n attempts to perform a single training step. Thereby + catches RayErrors resulting from worker failures. After n attempts, + fails gracefully. + + Override this method in your Algorithm sub-classes if you would like to + handle worker failures yourself. + Otherwise, override only `training_step()` to implement the core + algorithm logic. + + Returns: + The results dict with stats/infos on sampling, training, + and - if required - evaluation. + """ + # Ray metrics + with TimerAndPrometheusLogger(self._metrics_step_time): + # Do we have to run `self.evaluate()` this iteration? + # `self.iteration` gets incremented after this function returns, + # meaning that e.g. the first time this function is called, + # self.iteration will be 0. + evaluate_this_iter = ( + self.config.evaluation_interval + and (self.iteration + 1) % self.config.evaluation_interval == 0 + ) + + evaluate_offline_this_iter = ( + self.config.offline_evaluation_interval + and (self.iteration + 1) % self.config.offline_evaluation_interval == 0 + ) + + # Results dict for training (and if appolicable: evaluation). + eval_results: ResultDict = {} + + # Parallel eval + training: Kick off evaluation-loop and parallel train() call. + if evaluate_this_iter and ( + self.config.evaluation_parallel_to_training + or self.config.offline_evaluation_parallel_to_training + ): + ( + train_results, + eval_results, + train_iter_ctx, + ) = self._run_one_training_iteration_and_evaluation_in_parallel() + + # - No evaluation necessary, just run the next training iteration. + # - We have to evaluate in this training iteration, but no parallelism -> + # evaluate after the training iteration is entirely done. + else: + if self.config.enable_env_runner_and_connector_v2: + train_results, train_iter_ctx = self._run_one_training_iteration() + else: + ( + train_results, + train_iter_ctx, + ) = self._run_one_training_iteration_old_api_stack() + + # Sequential: Train (already done above), then evaluate. + if evaluate_this_iter and not self.config.evaluation_parallel_to_training: + eval_results = self._run_one_evaluation(parallel_train_future=None) + + if evaluate_offline_this_iter: + offline_eval_results = self._run_one_offline_evaluation() + # If we already have online evaluation results merge the offline + # evaluation results. + if eval_results: + eval_results[EVALUATION_RESULTS].update( + offline_eval_results[EVALUATION_RESULTS] + ) + # Otherwise, just assign. + else: + eval_results = offline_eval_results + + # Sync EnvRunner workers. + # TODO (sven): For the new API stack, the common execution pattern for any algo + # should be: [sample + get_metrics + get_state] -> send all these in one remote + # call down to `training_step` (where episodes are sent as ray object + # references). Then distribute the episode refs to the learners, store metrics + # in special key in result dict and perform the connector merge/broadcast + # inside the `training_step` as well. See the new IMPALA for an example. + if self.config.enable_env_runner_and_connector_v2: + if ( + not self.config._dont_auto_sync_env_runner_states + and self.env_runner_group + ): + # Synchronize EnvToModule and ModuleToEnv connector states and broadcast + # new states back to all EnvRunners. + with self.metrics.log_time( + (TIMERS, SYNCH_ENV_CONNECTOR_STATES_TIMER) + ): + with TimerAndPrometheusLogger( + self._metrics_step_sync_env_runner_states_time + ): + self.env_runner_group.sync_env_runner_states( + config=self.config, + env_to_module=self.env_to_module_connector, + module_to_env=self.module_to_env_connector, + ) + # Compile final ResultDict from `train_results` and `eval_results`. Note + # that, as opposed to the old API stack, EnvRunner stats should already be + # in `train_results` and `eval_results`. + results = self._compile_iteration_results( + train_results=train_results, + eval_results=eval_results, + ) + else: + self._sync_filters_if_needed( + central_worker=self.env_runner_group.local_env_runner, + workers=self.env_runner_group, + config=self.config, + ) + # Get EnvRunner metrics and compile them into results. + episodes_this_iter = collect_episodes( + self.env_runner_group, + self._remote_worker_ids_for_metrics(), + timeout_seconds=self.config.metrics_episode_collection_timeout_s, + ) + results = self._compile_iteration_results_old_api_stack( + episodes_this_iter=episodes_this_iter, + step_ctx=train_iter_ctx, + iteration_results={**train_results, **eval_results}, + ) + + return results + + @PublicAPI + def evaluate_offline(self): + """Evaluates current policy offline under `evaluation_config` settings. + + Returns: + A ResultDict only containing the offline evaluation results from the current + iteration. + """ + + # First synchronize weights. + self.offline_eval_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=self.config.offline_eval_rl_module_inference_only, + ) + + # TODO (simon): Check, how we can sync without a local runner. Also, + # connectors are in the data pipeline not directly in the runner applied. + # if self.config.broadcast_offline_eval_runner_states: + # # TODO (simon): Create offline equivalent. + # with self.metrics.log_time(TIMERS, SYNCH_EVAL_ENV_CONNECTOR_STATES_TIMER): + # self.offline_eval_runner_group.sync_runner_states( + # from_runner= + # ) + make_callback( + "on_evaluate_offline_start", + callbacks_objects=self.callbacks, + callbacks_functions=self.config.callbacks_on_evaluate_offline_start, + kwargs=dict(algorithm=self, metrics_logger=self.metrics), + ) + + # Evaluate with fixed duration. + if self.offline_eval_runner_group.num_healthy_remote_runners > 0: + self._evaluate_offline_with_fixed_duration() + else: + self._evaluate_offline_on_local_runner() + # Reduce the evaluation results. + eval_results = self.metrics.peek( + (EVALUATION_RESULTS, OFFLINE_EVAL_RUNNER_RESULTS), default={} + ) + + # Trigger `on_evaluate_offline_end` callback. + make_callback( + "on_evaluate_offline_end", + callbacks_objects=self.callbacks, + callbacks_functions=self.config.callbacks_on_evaluate_offline_end, + kwargs=dict( + algorithm=self, + metrics_logger=self.metrics, + evaluation_metrics=eval_results, + ), + ) + + # Also return the results here for convenience. + return {OFFLINE_EVAL_RUNNER_RESULTS: eval_results} + + @PublicAPI + def evaluate( + self, + parallel_train_future: Optional[concurrent.futures.ThreadPoolExecutor] = None, + ) -> ResultDict: + """Evaluates current policy under `evaluation_config` settings. + + Args: + parallel_train_future: In case, we are training and avaluating in parallel, + this arg carries the currently running ThreadPoolExecutor object that + runs the training iteration. Use `parallel_train_future.done()` to + check, whether the parallel training job has completed and + `parallel_train_future.result()` to get its return values. + + Returns: + A ResultDict only containing the evaluation results from the current + iteration. + """ + with TimerAndPrometheusLogger(self._metrics_evaluate_time): + # Call the `_before_evaluate` hook. + self._before_evaluate() + + if self.evaluation_dataset is not None: + return self._run_offline_evaluation_old_api_stack() + + if self.config.enable_env_runner_and_connector_v2: + if ( + self.env_runner_group is not None + and self.env_runner_group.healthy_env_runner_ids() + ): + # TODO (sven): Replace this with a new ActorManager API: + # try_remote_request_till_success("get_state") -> tuple(int, + # remoteresult) + weights_src = self.env_runner_group._worker_manager._actors[ + self.env_runner_group.healthy_env_runner_ids()[0] + ] + else: + weights_src = self.learner_group + else: + weights_src = self.env_runner + + # Sync weights to the evaluation EnvRunners. + if self.eval_env_runner_group is not None: + with TimerAndPrometheusLogger( + self._metrics_evaluate_sync_env_runner_weights_time + ): + self.eval_env_runner_group.sync_weights( + from_worker_or_learner_group=weights_src, + inference_only=True, + ) + + # Merge (eval) EnvRunner states and broadcast the merged state back + # to the remote (eval) EnvRunner actors. + if self.config.enable_env_runner_and_connector_v2: + if self.evaluation_config.broadcast_env_runner_states: + with self.metrics.log_time( + (TIMERS, SYNCH_EVAL_ENV_CONNECTOR_STATES_TIMER) + ): + with TimerAndPrometheusLogger( + self._metrics_evaluate_sync_connector_states_time + ): + self.eval_env_runner_group.sync_env_runner_states( + config=self.evaluation_config, + from_worker=self.env_runner, + env_to_module=self.env_to_module_connector, + module_to_env=self.module_to_env_connector, + ) + else: + self._sync_filters_if_needed( + central_worker=self.env_runner_group.local_env_runner, + workers=self.eval_env_runner_group, + config=self.evaluation_config, + ) + # Sync weights to the local EnvRunner (if no eval EnvRunnerGroup). + elif self.config.enable_env_runner_and_connector_v2: + self.env_runner_group.sync_weights( + from_worker_or_learner_group=weights_src, + inference_only=True, + ) + + with TimerAndPrometheusLogger( + self._metrics_callback_on_evaluate_start_time + ): + make_callback( + "on_evaluate_start", + callbacks_objects=self.callbacks, + callbacks_functions=self.config.callbacks_on_evaluate_start, + kwargs=dict(algorithm=self, metrics_logger=self.metrics), + ) + + env_steps = agent_steps = 0 + batches = [] + + # We will use a user provided evaluation function. + if self.config.custom_evaluation_function: + if self.config.enable_env_runner_and_connector_v2: + ( + eval_results, + env_steps, + agent_steps, + ) = self._evaluate_with_custom_eval_function() + else: + eval_results = self.config.custom_evaluation_function() + # There is no eval EnvRunnerGroup -> Run on local EnvRunner. + elif self.eval_env_runner_group is None and self.env_runner: + ( + eval_results, + env_steps, + agent_steps, + batches, + ) = self._evaluate_on_local_env_runner(self.env_runner) + # There is only a local eval EnvRunner -> Run on that. + elif self.eval_env_runner_group.num_healthy_remote_workers() == 0: + ( + eval_results, + env_steps, + agent_steps, + batches, + ) = self._evaluate_on_local_env_runner(self.eval_env_runner) + # There are healthy remote evaluation workers -> Run on these. + elif self.eval_env_runner_group.num_healthy_remote_workers() > 0: + # Running in automatic duration mode (parallel with training step). + if self.config.evaluation_duration == "auto": + assert parallel_train_future is not None + ( + eval_results, + env_steps, + agent_steps, + batches, + ) = self._evaluate_with_auto_duration(parallel_train_future) + # Running with a fixed amount of data to sample. + else: + ( + eval_results, + env_steps, + agent_steps, + batches, + ) = self._evaluate_with_fixed_duration() + # Can't find a good way to run this evaluation -> Wait for next iteration. + else: + eval_results = {} + + if self.config.enable_env_runner_and_connector_v2: + eval_results = self.metrics.peek(key=EVALUATION_RESULTS, default={}) + if log_once("no_eval_results") and not eval_results: + logger.warning( + "No evaluation results found for this iteration. This can happen if the evaluation worker(s) is/are not healthy." + ) + else: + eval_results = {ENV_RUNNER_RESULTS: eval_results} + eval_results[NUM_AGENT_STEPS_SAMPLED_THIS_ITER] = agent_steps + eval_results[NUM_ENV_STEPS_SAMPLED_THIS_ITER] = env_steps + eval_results["timesteps_this_iter"] = env_steps + self._counters[ + NUM_ENV_STEPS_SAMPLED_FOR_EVALUATION_THIS_ITER + ] = env_steps + + # Compute off-policy estimates + if not self.config.custom_evaluation_function: + estimates = defaultdict(list) + # for each batch run the estimator's fwd pass + for name, estimator in self.reward_estimators.items(): + for batch in batches: + estimate_result = estimator.estimate( + batch, + split_batch_by_episode=self.config.ope_split_batch_by_episode, + ) + estimates[name].append(estimate_result) + + # collate estimates from all batches + if estimates: + eval_results["off_policy_estimator"] = {} + for name, estimate_list in estimates.items(): + avg_estimate = tree.map_structure( + lambda *x: np.mean(x, axis=0), *estimate_list + ) + eval_results["off_policy_estimator"][name] = avg_estimate + + # Trigger `on_evaluate_end` callback. + with TimerAndPrometheusLogger(self._metrics_callback_on_evaluate_end_time): + make_callback( + "on_evaluate_end", + callbacks_objects=self.callbacks, + callbacks_functions=self.config.callbacks_on_evaluate_end, + kwargs=dict( + algorithm=self, + metrics_logger=self.metrics, + evaluation_metrics=eval_results, + ), + ) + + # Also return the results here for convenience. + return eval_results + + def _evaluate_with_custom_eval_function(self) -> Tuple[ResultDict, int, int]: + logger.info( + f"Evaluating current state of {self} using the custom eval function " + f"{self.config.custom_evaluation_function}" + ) + if self.config.enable_env_runner_and_connector_v2: + ( + eval_results, + env_steps, + agent_steps, + ) = self.config.custom_evaluation_function(self, self.eval_env_runner_group) + if not env_steps or not agent_steps: + raise ValueError( + "Custom eval function must return " + "`Tuple[ResultDict, int, int]` with `int, int` being " + f"`env_steps` and `agent_steps`! Got {env_steps}, {agent_steps}." + ) + else: + eval_results = self.config.custom_evaluation_function() + if not eval_results or not isinstance(eval_results, dict): + raise ValueError( + "Custom eval function must return " + f"dict of metrics! Got {eval_results}." + ) + + return eval_results, env_steps, agent_steps + + def _evaluate_offline_on_local_runner(self): + # How many episodes/timesteps do we need to run? + unit = "batches" + duration = ( + self.config.offline_evaluation_duration + * self.config.dataset_num_iters_per_eval_runner + ) + + logger.info(f"Evaluating current state of {self} for {duration} {unit}.") + + results = self.offline_eval_runner_group.local_runner.run() + + self.metrics.aggregate( + [results], + key=(EVALUATION_RESULTS, OFFLINE_EVAL_RUNNER_RESULTS), + ) + + def _evaluate_on_local_env_runner(self, env_runner): + if hasattr(env_runner, "input_reader") and env_runner.input_reader is None: + raise ValueError( + "Can't evaluate on a local worker if this local worker does not have " + "an environment!\nTry one of the following:" + "\n1) Set `evaluation_interval` > 0 to force creating a separate " + "evaluation EnvRunnerGroup.\n2) Set `create_local_env_runner=True` to " + "force the local (non-eval) EnvRunner to have an environment to " + "evaluate on." + ) + elif self.config.evaluation_parallel_to_training: + raise ValueError( + "Cannot run on local evaluation worker parallel to training! Try " + "setting `evaluation_parallel_to_training=False`." + ) + + # How many episodes/timesteps do we need to run? + unit = self.config.evaluation_duration_unit + duration = self.config.evaluation_duration + eval_cfg = self.evaluation_config + + env_steps = agent_steps = 0 + + logger.info(f"Evaluating current state of {self} for {duration} {unit}.") + + all_batches = [] + if self.config.enable_env_runner_and_connector_v2: + episodes = env_runner.sample( + num_timesteps=duration if unit == "timesteps" else None, + num_episodes=duration if unit == "episodes" else None, + ) + agent_steps += sum(e.agent_steps() for e in episodes) + env_steps += sum(e.env_steps() for e in episodes) + elif unit == "episodes": + for _ in range(duration): + batch = env_runner.sample() + agent_steps += batch.agent_steps() + env_steps += batch.env_steps() + if self.reward_estimators: + all_batches.append(batch) + else: + batch = env_runner.sample() + agent_steps += batch.agent_steps() + env_steps += batch.env_steps() + if self.reward_estimators: + all_batches.append(batch) + + env_runner_results = env_runner.get_metrics() + + if not self.config.enable_env_runner_and_connector_v2: + env_runner_results = summarize_episodes( + env_runner_results, + env_runner_results, + keep_custom_metrics=eval_cfg.keep_per_episode_custom_metrics, + ) + else: + self.metrics.aggregate( + [env_runner_results], + key=(EVALUATION_RESULTS, ENV_RUNNER_RESULTS), + ) + env_runner_results = None + + return env_runner_results, env_steps, agent_steps, all_batches + + def _evaluate_with_auto_duration(self, parallel_train_future): + logger.info( + f"Evaluating current state of {self} for as long as the parallelly " + "running training step takes." + ) + + all_metrics = [] + all_batches = [] + + # How many episodes have we run (across all eval workers)? + num_healthy_workers = self.eval_env_runner_group.num_healthy_remote_workers() + # Do we have to force-reset the EnvRunners before the first round of `sample()` + # calls.? + force_reset = self.config.evaluation_force_reset_envs_before_iteration + + # Remote function used on healthy EnvRunners to sample, get metrics, and + # step counts. + def _env_runner_remote(worker, num, round, iter): + # Sample AND get_metrics, but only return metrics (and steps actually taken) + # to save time. + episodes = worker.sample( + num_timesteps=num, force_reset=force_reset and round == 0 + ) + metrics = worker.get_metrics() + env_steps = sum(e.env_steps() for e in episodes) + agent_steps = sum(e.agent_steps() for e in episodes) + return env_steps, agent_steps, metrics, iter + + env_steps = agent_steps = 0 + if self.config.enable_env_runner_and_connector_v2: + train_mean_time = self.metrics.peek( + (TIMERS, TRAINING_ITERATION_TIMER), default=0.0 + ) + else: + train_mean_time = self._timers[TRAINING_ITERATION_TIMER].mean + t0 = time.time() + algo_iteration = self.iteration + + _round = -1 + while ( + # In case all the remote evaluation workers die during a round of + # evaluation, we need to stop. + num_healthy_workers > 0 + # Run at least for one round AND at least for as long as the parallel + # training step takes. + and (_round == -1 or not parallel_train_future.done()) + ): + _round += 1 + # New API stack -> EnvRunners return Episodes. + if self.config.enable_env_runner_and_connector_v2: + # Compute rough number of timesteps it takes for a single EnvRunner + # to occupy the estimated (parallelly running) train step. + _num = min( + # Clamp number of steps to take between a max and a min. + self.config.evaluation_auto_duration_max_env_steps_per_sample, + max( + self.config.evaluation_auto_duration_min_env_steps_per_sample, + ( + # How much time do we have left? + (train_mean_time - (time.time() - t0)) + # Multiply by our own (eval) throughput to get the timesteps + # to do (per worker). + * self.metrics.peek( + ( + EVALUATION_RESULTS, + ENV_RUNNER_RESULTS, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + ), + throughput=True, + default=0.0, + ) + / num_healthy_workers + ), + ), + ) + + results = ( + self.eval_env_runner_group.foreach_env_runner_async_fetch_ready( + func=_env_runner_remote, + kwargs={"num": _num, "round": _round, "iter": algo_iteration}, + tag="_env_runner_remote", + ) + ) + + for env_s, ag_s, metrics, iter in results: + # Ignore eval results kicked off in an earlier iteration. + # (those results would be outdated and thus misleading). + if iter != self.iteration: + continue + env_steps += env_s + agent_steps += ag_s + all_metrics.append(metrics) + time.sleep(0.01) + + # Old API stack -> RolloutWorkers return batches. + else: + results = ( + self.eval_env_runner_group.foreach_env_runner_async_fetch_ready( + func=lambda w: (w.sample(), w.get_metrics(), algo_iteration), + tag="env_runner_sample_and_get_metrics", + ) + ) + + for batch, metrics, iter in results: + if iter != self.iteration: + continue + env_steps += batch.env_steps() + agent_steps += batch.agent_steps() + all_metrics.extend(metrics) + if self.reward_estimators: + # TODO: (kourosh) This approach will cause an OOM issue when + # the dataset gets huge (should be ok for now). + all_batches.append(batch) + + # Update correct number of healthy remote workers. + num_healthy_workers = ( + self.eval_env_runner_group.num_healthy_remote_workers() + ) + + if num_healthy_workers == 0: + logger.warning( + "Calling `sample()` on your remote evaluation worker(s) " + "resulted in all workers crashing! Make sure a) your environment is not" + " too unstable, b) you have enough evaluation workers " + "(`config.evaluation(evaluation_num_env_runners=...)`) to cover for " + "occasional losses, and c) you use the `config.fault_tolerance(" + "restart_failed_env_runners=True)` setting." + ) + + if not self.config.enable_env_runner_and_connector_v2: + env_runner_results = summarize_episodes( + all_metrics, + all_metrics, + keep_custom_metrics=( + self.evaluation_config.keep_per_episode_custom_metrics + ), + ) + num_episodes = env_runner_results[NUM_EPISODES] + else: + self.metrics.aggregate( + all_metrics, + key=(EVALUATION_RESULTS, ENV_RUNNER_RESULTS), + ) + num_episodes = self.metrics.peek( + (EVALUATION_RESULTS, ENV_RUNNER_RESULTS, NUM_EPISODES), + default=0, + ) + env_runner_results = None + + # Warn if results are empty, it could be that this is because the auto-time is + # not enough to run through one full episode. + if ( + self.config.evaluation_force_reset_envs_before_iteration + and num_episodes == 0 + ): + logger.warning( + "This evaluation iteration resulted in an empty set of episode summary " + "results! It's possible that the auto-duration time (roughly the mean " + "time it takes for the training step to finish) is not enough to finish" + " even a single episode. Your current mean training iteration time is " + f"{train_mean_time}sec. Try setting the min iteration time to a higher " + "value via the `config.reporting(min_time_s_per_iteration=...)` OR you " + "can also set `config.evaluation_force_reset_envs_before_iteration` to " + "False. However, keep in mind that then the evaluation results may " + "contain some episode stats generated with earlier weights versions." + ) + + return env_runner_results, env_steps, agent_steps, all_batches + + def _evaluate_offline_with_fixed_duration(self) -> None: + # How many batches do we need to run? + num_workers = self.config.num_offline_eval_runners + time_out = self.config.offline_evaluation_timeout_s + + def _offline_eval_runner_remote(runner, iter): + + metrics = runner.run() + + return metrics, iter + + all_metrics = [] + num_units_done = [] + + # How many episodes have we run (across all eval workers)? + num_units_done = 0 + num_healthy_workers = self.offline_eval_runner_group.num_healthy_remote_runners + + # TODO (simon): Note, agent steps might not be available, but only + # module steps. + + t_last_result = time.time() + _round = -1 + algo_iteration = self.iteration + + # In case all the remote evaluation workers die during a round of + # evaluation, we need to stop. + while num_healthy_workers > 0: + units_left_to_do = ( + self.config.offline_evaluation_duration * num_workers - num_units_done + ) + if units_left_to_do <= 0: + break + + _round += 1 + + self.offline_eval_runner_group.foreach_runner_async( + func=functools.partial( + _offline_eval_runner_remote, + iter=algo_iteration, + ), + ) + results = self.offline_eval_runner_group.fetch_ready_async_reqs( + return_obj_refs=False, timeout_seconds=0.01 + ) + # Make sure we properly time out if we have not received any results + # for more than `time_out` seconds. + time_now = time.time() + if not results and time_now - t_last_result > time_out: + break + elif results: + t_last_result = time_now + for wid, (met, iter) in results: + if iter != self.iteration: + continue + all_metrics.append(met) + # Note, the `dataset_num_iters_per_eval_runner` must be smaller than + # `offline_evaluation_duration` // `num_offline_eval_runners`. + num_units_done += ( + met[ALL_MODULES][DATASET_NUM_ITERS_EVALUATED].peek() + if DATASET_NUM_ITERS_EVALUATED in met[ALL_MODULES] + else 0 + ) + + # Update correct number of healthy remote workers. + num_healthy_workers = ( + self.offline_eval_runner_group.num_healthy_remote_runners + ) + + if num_healthy_workers == 0: + logger.warning( + "Calling `run()` on your remote offline evaluation runner(s) " + "resulted in all runners crashing! Make sure a) your dataset is not" + " corrupted, b) you have enough offline evaluation runners " + "(`config.evaluation(num_offline_eval_runners=...)`) to cover for " + "occasional losses, and c) you use the `config.fault_tolerance(" + "restart_failed_offline_eval_runners=True)` setting." + ) + + self.metrics.aggregate( + all_metrics, + key=(EVALUATION_RESULTS, OFFLINE_EVAL_RUNNER_RESULTS), + ) + + def _evaluate_with_fixed_duration(self): + # How many episodes/timesteps do we need to run? + unit = self.config.evaluation_duration_unit + eval_cfg = self.evaluation_config + num_workers = self.config.evaluation_num_env_runners + force_reset = self.config.evaluation_force_reset_envs_before_iteration + time_out = self.config.evaluation_sample_timeout_s + + # Remote function used on healthy EnvRunners to sample, get metrics, and + # step counts. + def _env_runner_remote(worker, num, round, iter, _force_reset): + # Sample AND get_metrics, but only return metrics (and steps actually taken) + # to save time. Also return the iteration to check, whether we should + # discard and outdated result (from a slow worker). + episodes = worker.sample( + num_timesteps=( + num[worker.worker_index] if unit == "timesteps" else None + ), + num_episodes=(num[worker.worker_index] if unit == "episodes" else None), + force_reset=_force_reset and round == 0, + ) + metrics = worker.get_metrics() + env_steps = sum(e.env_steps() for e in episodes) + agent_steps = sum(e.agent_steps() for e in episodes) + return env_steps, agent_steps, metrics, iter + + all_metrics = [] + all_batches = [] + + # How many episodes have we run (across all eval workers)? + num_units_done = 0 + num_healthy_workers = self.eval_env_runner_group.num_healthy_remote_workers() + + env_steps = agent_steps = 0 + + t_last_result = time.time() + _round = -1 + algo_iteration = self.iteration + + # In case all the remote evaluation workers die during a round of + # evaluation, we need to stop. + while num_healthy_workers > 0: + units_left_to_do = self.config.evaluation_duration - num_units_done + if units_left_to_do <= 0: + break + + _round += 1 + + # New API stack -> EnvRunners return Episodes. + if self.config.enable_env_runner_and_connector_v2: + _num = [None] + [ # [None]: skip idx=0 (local worker) + (units_left_to_do // num_healthy_workers) + + bool(i <= (units_left_to_do % num_healthy_workers)) + for i in range(1, num_workers + 1) + ] + + results = ( + self.eval_env_runner_group.foreach_env_runner_async_fetch_ready( + func=_env_runner_remote, + kwargs={ + "num": _num, + "round": _round, + "iter": algo_iteration, + "_force_reset": force_reset, + }, + tag="_env_runner_remote", + ) + ) + + # Make sure we properly time out if we have not received any results + # for more than `time_out` seconds. + time_now = time.time() + if not results and time_now - t_last_result > time_out: + break + elif results: + t_last_result = time_now + for env_s, ag_s, met, iter in results: + if iter != self.iteration: + continue + env_steps += env_s + agent_steps += ag_s + all_metrics.append(met) + num_units_done += ( + (met[NUM_EPISODES].peek() if NUM_EPISODES in met else 0) + if unit == "episodes" + else ( + env_s if self.config.count_steps_by == "env_steps" else ag_s + ) + ) + # Old API stack -> RolloutWorkers return batches. + else: + units_per_healthy_remote_worker = ( + 1 + if unit == "episodes" + else eval_cfg.rollout_fragment_length + * eval_cfg.num_envs_per_env_runner + ) + # Select proper number of evaluation workers for this round. + selected_eval_worker_ids = [ + worker_id + for i, worker_id in enumerate( + self.eval_env_runner_group.healthy_worker_ids() + ) + if i * units_per_healthy_remote_worker < units_left_to_do + ] + + results = ( + self.eval_env_runner_group.foreach_env_runner_async_fetch_ready( + func=lambda w: (w.sample(), w.get_metrics(), algo_iteration), + remote_worker_ids=selected_eval_worker_ids, + tag="env_runner_sample_and_get_metrics", + ) + ) + # Make sure we properly time out if we have not received any results + # for more than `time_out` seconds. + time_now = time.time() + if not results and time_now - t_last_result > time_out: + break + elif results: + t_last_result = time_now + for batch, metrics, iter in results: + if iter != self.iteration: + continue + env_steps += batch.env_steps() + agent_steps += batch.agent_steps() + all_metrics.extend(metrics) + if self.reward_estimators: + # TODO: (kourosh) This approach will cause an OOM issue when + # the dataset gets huge (should be ok for now). + all_batches.append(batch) + + # 1 episode per returned batch. + if unit == "episodes": + num_units_done += len(results) + # n timesteps per returned batch. + else: + num_units_done = ( + env_steps + if self.config.count_steps_by == "env_steps" + else agent_steps + ) + + # Update correct number of healthy remote workers. + num_healthy_workers = ( + self.eval_env_runner_group.num_healthy_remote_workers() + ) + + if num_healthy_workers == 0: + logger.warning( + "Calling `sample()` on your remote evaluation worker(s) " + "resulted in all workers crashing! Make sure a) your environment is not" + " too unstable, b) you have enough evaluation workers " + "(`config.evaluation(evaluation_num_env_runners=...)`) to cover for " + "occasional losses, and c) you use the `config.fault_tolerance(" + "restart_failed_env_runners=True)` setting." + ) + + if not self.config.enable_env_runner_and_connector_v2: + env_runner_results = summarize_episodes( + all_metrics, + all_metrics, + keep_custom_metrics=( + self.evaluation_config.keep_per_episode_custom_metrics + ), + ) + num_episodes = env_runner_results[NUM_EPISODES] + else: + self.metrics.aggregate( + all_metrics, + key=(EVALUATION_RESULTS, ENV_RUNNER_RESULTS), + ) + num_episodes = self.metrics.peek( + (EVALUATION_RESULTS, ENV_RUNNER_RESULTS, NUM_EPISODES), default=0 + ) + env_runner_results = None + + # Warn if results are empty, it could be that this is because the eval timesteps + # are not enough to run through one full episode. + if num_episodes == 0: + logger.warning( + "This evaluation iteration resulted in an empty set of episode summary " + "results! It's possible that your configured duration timesteps are not" + " enough to finish even a single episode. You have configured " + f"{self.config.evaluation_duration} " + f"{self.config.evaluation_duration_unit}. For 'timesteps', try " + "increasing this value via the `config.evaluation(evaluation_duration=" + "...)` OR change the unit to 'episodes' via `config.evaluation(" + "evaluation_duration_unit='episodes')` OR try increasing the timeout " + "threshold via `config.evaluation(evaluation_sample_timeout_s=...)` OR " + "you can also set `config.evaluation_force_reset_envs_before_iteration`" + " to False. However, keep in mind that in the latter case, the " + "evaluation results may contain some episode stats generated with " + "earlier weights versions." + ) + + return env_runner_results, env_steps, agent_steps, all_batches + + @OverrideToImplementCustomLogic + def restore_env_runners(self, env_runner_group: EnvRunnerGroup) -> List[int]: + """Try bringing back unhealthy EnvRunners and - if successful - sync with local. + + Algorithms that use custom EnvRunners may override this method to + disable the default, and create custom restoration logics. Note that "restoring" + does not include the actual restarting process, but merely what should happen + after such a restart of a (previously failed) worker. + + Args: + env_runner_group: The EnvRunnerGroup to restore. This may be the training or + the evaluation EnvRunnerGroup. + + Returns: + A list of EnvRunner indices that have been restored during the call of + this method. + """ + # This is really cheap, since probe_unhealthy_env_runners() is a no-op + # if there are no unhealthy workers. + restored = None + if self.config.is_online: + restored = env_runner_group.probe_unhealthy_env_runners() + + if not restored: + return [] + + # Count the restored workers. + self._counters["total_num_restored_workers"] += len(restored) + + from_env_runner = env_runner_group.local_env_runner or self.env_runner + + # Sync from local EnvRunner, if it exists. + if from_env_runner is not None: + # Get the state of the EnvRunner. + state = from_env_runner.get_state() + state_ref = ray.put(state) + + # Take out (old) connector states from local worker's state. + if not self.config.enable_env_runner_and_connector_v2: + for pol_states in state["policy_states"].values(): + pol_states.pop("connector_configs", None) + + elif self.config.is_multi_agent: + + multi_rl_module_spec = MultiRLModuleSpec.from_module( + from_env_runner.module + ) + + # Otherwise, sync from another EnvRunner that's still healthy. + else: + multi_rl_module_spec = ( + self.learner_group.foreach_learner( + lambda learner: MultiRLModuleSpec.from_module(learner.module) + ) + .result_or_errors[0] + .get() + ) + + # Sync the weights from the learner group to the EnvRunners. + state = self.learner_group.get_state( + components=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE, + inference_only=True, + )[COMPONENT_LEARNER] + state[ + COMPONENT_ENV_TO_MODULE_CONNECTOR + ] = self.env_to_module_connector.get_state() + state[ + COMPONENT_MODULE_TO_ENV_CONNECTOR + ] = self.module_to_env_connector.get_state() + state[NUM_ENV_STEPS_SAMPLED_LIFETIME] = self.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME), default=0 + ) + state_ref = ray.put(state) + + def _sync_env_runner(er): # noqa + # Remove modules (new API stack only), if necessary. + if ( + er.config.enable_env_runner_and_connector_v2 + and er.config.is_multi_agent + ): + for module_id, module in er.module._rl_modules.copy().items(): + if module_id not in multi_rl_module_spec.rl_module_specs: + er.module.remove_module(module_id, raise_err_if_not_found=True) + # Add modules, if necessary. + for mid, mod_spec in multi_rl_module_spec.rl_module_specs.items(): + if mid not in er.module: + er.module.add_module(mid, mod_spec.build(), override=False) + # Now that the MultiRLModule is fixed, update the state. + er.set_state(ray.get(state_ref)) + + # By default, entire local EnvRunner state is synced after restoration + # to bring the previously failed EnvRunner up to date. + env_runner_group.foreach_env_runner( + func=_sync_env_runner, + remote_worker_ids=restored, + # Don't update the local EnvRunner, b/c it's the one we are synching + # from. + local_env_runner=False, + timeout_seconds=self.config.env_runner_restore_timeout_s, + ) + + return restored + + @OverrideToImplementCustomLogic + def restore_offline_eval_runners(self, runner_group: RunnerGroup) -> List[int]: + if not runner_group or not runner_group.local_runner: + return [] + + restored = runner_group.probe_unhealthy_runners() + + if restored: + # Count the restored workers. + self._counters["total_num_restored_workers"] += len(restored) + + # Get the state of the correct (reference) worker. + from_runner = runner_group.healthy_runner_ids()[0] + state = runner_group.foreach_runner( + "get_state", + local_runner=False, + remote_worker_ids=from_runner, + )[0] + state_ref = ray.put(state) + + def _sync_runner(r): + r.set_state(ray.get(state_ref)) + + # By default, entire `Runner`` state is synced after restoration + # to bring the previously failed `Runner` up to date. + runner_group.foreach_runner( + func=_sync_runner, + remote_worker_ids=restored, + # Don't update the local `Runner`. + local_runner=False, + timeout_seconds=self.evaluation_config.offline_eval_runner_restore_timeout_s, + ) + # Restore the correct data iterator split stream. + # TODO (simon): Define a `restore` method in the `RunnerGroup` + # such that we do not have to check here for the group. + # Also get a different streaming split if a runner fails and is not + # recreated. + runner_group.foreach_runner( + func="set_dataset_iterator", + remote_worker_ids=restored, + local_runner=False, + timeout_seconds=self.evaluation_config.offline_eval_runner_restore_timeout_s, + kwargs={"iterator": runner_group._offline_data_iterators[restored]}, + ) + + return restored + + @OverrideToImplementCustomLogic + def training_step(self) -> None: + """Default single iteration logic of an algorithm. + + - Collect on-policy samples (SampleBatches) in parallel using the + Algorithm's EnvRunners (@ray.remote). + - Concatenate collected SampleBatches into one train batch. + - Note that we may have more than one policy in the multi-agent case: + Call the different policies' `learn_on_batch` (simple optimizer) OR + `load_batch_into_buffer` + `learn_on_loaded_batch` (multi-GPU + optimizer) methods to calculate loss and update the model(s). + - Return all collected metrics for the iteration. + + Returns: + For the new API stack, returns None. Results are compiled and extracted + automatically through a single `self.metrics.reduce()` call at the very end + of an iteration (which might contain more than one call to + `training_step()`). This way, we make sure that we account for all + results generated by each individual `training_step()` call. + For the old API stack, returns the results dict from executing the training + step. + """ + if not self.config.enable_env_runner_and_connector_v2: + raise NotImplementedError( + "The `Algorithm.training_step()` default implementation no longer " + "supports the old API stack! If you would like to continue " + "using these " + "old APIs with this default `training_step`, simply subclass " + "`Algorithm` and override its `training_step` method (copy/paste the " + "code and delete this error message)." + ) + + # Collect a list of Episodes from EnvRunners until we reach the train batch + # size. + with self.metrics.log_time((TIMERS, ENV_RUNNER_SAMPLING_TIMER)): + if self.config.count_steps_by == "agent_steps": + episodes, env_runner_results = synchronous_parallel_sample( + worker_set=self.env_runner_group, + max_agent_steps=self.config.total_train_batch_size, + sample_timeout_s=self.config.sample_timeout_s, + _uses_new_env_runners=True, + _return_metrics=True, + ) + else: + episodes, env_runner_results = synchronous_parallel_sample( + worker_set=self.env_runner_group, + max_env_steps=self.config.total_train_batch_size, + sample_timeout_s=self.config.sample_timeout_s, + _uses_new_env_runners=True, + _return_metrics=True, + ) + # Reduce EnvRunner metrics over the n EnvRunners. + self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS) + + with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)): + learner_results = self.learner_group.update( + episodes=episodes, + timesteps={ + NUM_ENV_STEPS_SAMPLED_LIFETIME: ( + self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME) + ), + }, + ) + self.metrics.log_dict(learner_results, key=LEARNER_RESULTS) + + # Update weights - after learning on the local worker - on all + # remote workers (only those RLModules that were actually trained). + with self.metrics.log_time((TIMERS, SYNCH_WORKER_WEIGHTS_TIMER)): + self.env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + policies=list(set(learner_results.keys()) - {ALL_MODULES}), + inference_only=True, + ) + + @PublicAPI + def get_module(self, module_id: ModuleID = DEFAULT_MODULE_ID) -> Optional[RLModule]: + """Returns the (single-agent) RLModule with `model_id` (None if ID not found). + + Args: + module_id: ID of the (single-agent) RLModule to return from the MARLModule + used by the local EnvRunner. + + Returns: + The RLModule found under the ModuleID key inside the local EnvRunner's + MultiRLModule. None if `module_id` doesn't exist. + """ + if self.env_runner is not None: + module = self.env_runner.module + else: + module = self.env_runner_group.foreach_env_runner( + lambda er: er.module, + remote_worker_ids=[1], + local_env_runner=False, + )[0] + + if isinstance(module, MultiRLModule): + return module.get(module_id) + else: + return module + + @PublicAPI + def add_module( + self, + module_id: ModuleID, + module_spec: RLModuleSpec, + *, + config_overrides: Optional[Dict] = None, + new_agent_to_module_mapping_fn: Optional[AgentToModuleMappingFn] = None, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + add_to_learners: bool = True, + add_to_env_runners: bool = True, + add_to_eval_env_runners: bool = True, + ) -> MultiRLModuleSpec: + """Adds a new (single-agent) RLModule to this Algorithm's MARLModule. + + Note that an Algorithm has up to 3 different components to which to add + the new module to: The LearnerGroup (with n Learners), the EnvRunnerGroup + (with m EnvRunners plus a local one) and - if applicable - the eval + EnvRunnerGroup (with o EnvRunners plus a local one). + + Args: + module_id: ID of the RLModule to add to the MARLModule. + IMPORTANT: Must not contain characters that + are also not allowed in Unix/Win filesystems, such as: `<>:"/|?*`, + or a dot, space or backslash at the end of the ID. + module_spec: The SingleAgentRLModuleSpec to use for constructing the new + RLModule. + config_overrides: The `AlgorithmConfig` overrides that should apply to + the new Module, if any. + new_agent_to_module_mapping_fn: An optional (updated) AgentID to ModuleID + mapping function to use from here on. Note that already ongoing + episodes will not change their mapping but will use the old mapping till + the end of the episode. + new_should_module_be_updated: An optional sequence of ModuleIDs or a + callable taking ModuleID and SampleBatchType and returning whether the + ModuleID should be updated (trained). + If None, will keep the existing setup in place. RLModules, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + add_to_learners: Whether to add the new RLModule to the LearnerGroup + (with its n Learners). + add_to_env_runners: Whether to add the new RLModule to the EnvRunnerGroup + (with its m EnvRunners plus the local one). + add_to_eval_env_runners: Whether to add the new RLModule to the eval + EnvRunnerGroup (with its o EnvRunners plus the local one). + + Returns: + The new MultiRLModuleSpec (after the RLModule has been added). + """ + validate_module_id(module_id, error=True) + + # The to-be-returned new MultiRLModuleSpec. + multi_rl_module_spec = None + + if not self.config.is_multi_agent: + raise RuntimeError( + "Can't add a new RLModule to a single-agent setup! Make sure that your " + "setup is already initially multi-agent by either defining >1 " + f"RLModules in your `rl_module_spec` or assigning a ModuleID other " + f"than {DEFAULT_MODULE_ID} to your (only) RLModule." + ) + + if not any([add_to_learners, add_to_env_runners, add_to_eval_env_runners]): + raise ValueError( + "At least one of `add_to_learners`, `add_to_env_runners`, or " + "`add_to_eval_env_runners` must be set to True!" + ) + + # Add to Learners and sync weights. + if add_to_learners: + multi_rl_module_spec = self.learner_group.add_module( + module_id=module_id, + module_spec=module_spec, + config_overrides=config_overrides, + new_should_module_be_updated=new_should_module_be_updated, + ) + + # Change our config (AlgorithmConfig) to contain the new Module. + # TODO (sven): This is a hack to manipulate the AlgorithmConfig directly, + # but we'll deprecate config.policies soon anyway. + self.config._is_frozen = False + self.config.policies[module_id] = PolicySpec() + if config_overrides is not None: + self.config.multi_agent( + algorithm_config_overrides_per_module={module_id: config_overrides} + ) + if new_agent_to_module_mapping_fn is not None: + self.config.multi_agent(policy_mapping_fn=new_agent_to_module_mapping_fn) + self.config.rl_module(rl_module_spec=multi_rl_module_spec) + if new_should_module_be_updated is not None: + self.config.multi_agent(policies_to_train=new_should_module_be_updated) + self.config.freeze() + + def _add(_env_runner, _module_spec=module_spec): + # Add the RLModule to the existing one on the EnvRunner. + _env_runner.module.add_module( + module_id=module_id, module=_module_spec.build() + ) + # Update the `agent_to_module_mapping_fn` on the EnvRunner. + if new_agent_to_module_mapping_fn is not None: + _env_runner.config.multi_agent( + policy_mapping_fn=new_agent_to_module_mapping_fn, + ) + # Update the `should_module_be_updated` on the EnvRunner. Note that + # even though this information is typically not needed by the EnvRunner, + # it's good practice to keep this setting updated everywhere either way. + if new_should_module_be_updated is not None: + _env_runner.config.multi_agent( + policies_to_train=new_should_module_be_updated, + ) + return MultiRLModuleSpec.from_module(_env_runner.module) + + # Add to (training) EnvRunners and sync weights. + if add_to_env_runners: + if multi_rl_module_spec is None: + multi_rl_module_spec = self.env_runner_group.foreach_env_runner(_add)[0] + else: + self.env_runner_group.foreach_env_runner(_add) + self.env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + # Add to eval EnvRunners and sync weights. + if add_to_eval_env_runners is True and self.eval_env_runner_group is not None: + if multi_rl_module_spec is None: + multi_rl_module_spec = self.eval_env_runner_group.foreach_env_runner( + _add + )[0] + else: + self.eval_env_runner_group.foreach_env_runner(_add) + self.eval_env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + + return multi_rl_module_spec + + @PublicAPI + def remove_module( + self, + module_id: ModuleID, + *, + new_agent_to_module_mapping_fn: Optional[AgentToModuleMappingFn] = None, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + remove_from_learners: bool = True, + remove_from_env_runners: bool = True, + remove_from_eval_env_runners: bool = True, + ) -> Optional[Policy]: + """Removes a new (single-agent) RLModule from this Algorithm's MARLModule. + + Args: + module_id: ID of the RLModule to remove from the MARLModule. + IMPORTANT: Must not contain characters that + are also not allowed in Unix/Win filesystems, such as: `<>:"/|?*`, + or a dot, space or backslash at the end of the ID. + new_agent_to_module_mapping_fn: An optional (updated) AgentID to ModuleID + mapping function to use from here on. Note that already ongoing + episodes will not change their mapping but will use the old mapping till + the end of the episode. + new_should_module_be_updated: An optional sequence of ModuleIDs or a + callable taking ModuleID and SampleBatchType and returning whether the + ModuleID should be updated (trained). + If None, will keep the existing setup in place. RLModules, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + remove_from_learners: Whether to remove the RLModule from the LearnerGroup + (with its n Learners). + remove_from_env_runners: Whether to remove the RLModule from the + EnvRunnerGroup (with its m EnvRunners plus the local one). + remove_from_eval_env_runners: Whether to remove the RLModule from the eval + EnvRunnerGroup (with its o EnvRunners plus the local one). + + Returns: + The new MultiRLModuleSpec (after the RLModule has been removed). + """ + # The to-be-returned new MultiRLModuleSpec. + multi_rl_module_spec = None + + # Remove RLModule from the LearnerGroup. + if remove_from_learners: + multi_rl_module_spec = self.learner_group.remove_module( + module_id=module_id, + new_should_module_be_updated=new_should_module_be_updated, + ) + + # Change our config (AlgorithmConfig) with the Module removed. + # TODO (sven): This is a hack to manipulate the AlgorithmConfig directly, + # but we'll deprecate config.policies soon anyway. + self.config._is_frozen = False + del self.config.policies[module_id] + self.config.algorithm_config_overrides_per_module.pop(module_id, None) + if new_agent_to_module_mapping_fn is not None: + self.config.multi_agent(policy_mapping_fn=new_agent_to_module_mapping_fn) + self.config.rl_module(rl_module_spec=multi_rl_module_spec) + if new_should_module_be_updated is not None: + self.config.multi_agent(policies_to_train=new_should_module_be_updated) + self.config.freeze() + + def _remove(_env_runner): + # Remove the RLModule from the existing one on the EnvRunner. + _env_runner.module.remove_module(module_id=module_id) + # Update the `agent_to_module_mapping_fn` on the EnvRunner. + if new_agent_to_module_mapping_fn is not None: + _env_runner.config.multi_agent( + policy_mapping_fn=new_agent_to_module_mapping_fn + ) + # Force reset all ongoing episodes on the EnvRunner to avoid having + # different ModuleIDs compute actions for the same AgentID in the same + # episode. + # TODO (sven): Create an API for this. + _env_runner._needs_initial_reset = True + + return MultiRLModuleSpec.from_module(_env_runner.module) + + # Remove from (training) EnvRunners and sync weights. + if remove_from_env_runners: + if multi_rl_module_spec is None: + multi_rl_module_spec = self.env_runner_group.foreach_env_runner( + _remove + )[0] + else: + self.env_runner_group.foreach_env_runner(_remove) + self.env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + + # Remove from (eval) EnvRunners and sync weights. + if ( + remove_from_eval_env_runners is True + and self.eval_env_runner_group is not None + ): + if multi_rl_module_spec is None: + multi_rl_module_spec = self.eval_env_runner_group.foreach_env_runner( + _remove + )[0] + else: + self.eval_env_runner_group.foreach_env_runner(_remove) + self.eval_env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + + return multi_rl_module_spec + + @OldAPIStack + def get_policy(self, policy_id: PolicyID = DEFAULT_POLICY_ID) -> Policy: + """Return policy for the specified id, or None. + + Args: + policy_id: ID of the policy to return. + """ + return self.env_runner.get_policy(policy_id) + + @PublicAPI + def get_weights(self, policies: Optional[List[PolicyID]] = None) -> dict: + """Return a dict mapping Module/Policy IDs to weights. + + Args: + policies: Optional list of policies to return weights for, + or None for all policies. + """ + # New API stack (get weights from LearnerGroup). + if self.learner_group is not None: + return self.learner_group.get_weights(module_ids=policies) + return self.env_runner.get_weights(policies) + + @PublicAPI + def set_weights(self, weights: Dict[PolicyID, dict]): + """Set RLModule/Policy weights by Module/Policy ID. + + Args: + weights: Dict mapping ModuleID/PolicyID to weights. + """ + # New API stack -> Use `set_state` API and specify the LearnerGroup state in the + # call, which will automatically take care of weight synching to all EnvRunners. + if self.learner_group is not None: + self.set_state( + { + COMPONENT_LEARNER_GROUP: { + COMPONENT_LEARNER: { + COMPONENT_RL_MODULE: weights, + }, + }, + }, + ) + self.env_runner_group.local_env_runner.set_weights(weights) + + @OldAPIStack + def add_policy( + self, + policy_id: PolicyID, + policy_cls: Optional[Type[Policy]] = None, + policy: Optional[Policy] = None, + *, + observation_space: Optional[gym.spaces.Space] = None, + action_space: Optional[gym.spaces.Space] = None, + config: Optional[Union[AlgorithmConfig, PartialAlgorithmConfigDict]] = None, + policy_state: Optional[PolicyState] = None, + policy_mapping_fn: Optional[Callable[[AgentID, EpisodeID], PolicyID]] = None, + policies_to_train: Optional[ + Union[ + Collection[PolicyID], + Callable[[PolicyID, Optional[SampleBatchType]], bool], + ] + ] = None, + add_to_env_runners: bool = True, + add_to_eval_env_runners: bool = True, + module_spec: Optional[RLModuleSpec] = None, + # Deprecated arg. + evaluation_workers=DEPRECATED_VALUE, + add_to_learners=DEPRECATED_VALUE, + ) -> Optional[Policy]: + """Adds a new policy to this Algorithm. + + Args: + policy_id: ID of the policy to add. + IMPORTANT: Must not contain characters that + are also not allowed in Unix/Win filesystems, such as: `<>:"/|?*`, + or a dot, space or backslash at the end of the ID. + policy_cls: The Policy class to use for constructing the new Policy. + Note: Only one of `policy_cls` or `policy` must be provided. + policy: The Policy instance to add to this algorithm. If not None, the + given Policy object will be directly inserted into the Algorithm's + local worker and clones of that Policy will be created on all remote + workers as well as all evaluation workers. + Note: Only one of `policy_cls` or `policy` must be provided. + observation_space: The observation space of the policy to add. + If None, try to infer this space from the environment. + action_space: The action space of the policy to add. + If None, try to infer this space from the environment. + config: The config object or overrides for the policy to add. + policy_state: Optional state dict to apply to the new + policy instance, right after its construction. + policy_mapping_fn: An optional (updated) policy mapping function + to use from here on. Note that already ongoing episodes will + not change their mapping but will use the old mapping till + the end of the episode. + policies_to_train: An optional list of policy IDs to be trained + or a callable taking PolicyID and SampleBatchType and + returning a bool (trainable or not?). + If None, will keep the existing setup in place. Policies, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + add_to_env_runners: Whether to add the new RLModule to the EnvRunnerGroup + (with its m EnvRunners plus the local one). + add_to_eval_env_runners: Whether to add the new RLModule to the eval + EnvRunnerGroup (with its o EnvRunners plus the local one). + module_spec: In the new RLModule API we need to pass in the module_spec for + the new module that is supposed to be added. Knowing the policy spec is + not sufficient. + + Returns: + The newly added policy (the copy that got added to the local + worker). If `workers` was provided, None is returned. + """ + if self.config.enable_env_runner_and_connector_v2: + raise ValueError( + "`Algorithm.add_policy()` is not supported on the new API stack w/ " + "EnvRunners! Use `Algorithm.add_module()` instead. Also see " + "`rllib/examples/self_play_league_based_with_open_spiel.py` for an " + "example." + ) + + if evaluation_workers != DEPRECATED_VALUE: + deprecation_warning( + old="Algorithm.add_policy(evaluation_workers=...)", + new="Algorithm.add_policy(add_to_eval_env_runners=...)", + error=True, + ) + if add_to_learners != DEPRECATED_VALUE: + deprecation_warning( + old="Algorithm.add_policy(add_to_learners=..)", + help="Hybrid API stack no longer supported by RLlib!", + error=True, + ) + + validate_module_id(policy_id, error=True) + + if add_to_env_runners is True: + self.env_runner_group.add_policy( + policy_id, + policy_cls, + policy, + observation_space=observation_space, + action_space=action_space, + config=config, + policy_state=policy_state, + policy_mapping_fn=policy_mapping_fn, + policies_to_train=policies_to_train, + module_spec=module_spec, + ) + + # Add to evaluation workers, if necessary. + if add_to_eval_env_runners is True and self.eval_env_runner_group is not None: + self.eval_env_runner_group.add_policy( + policy_id, + policy_cls, + policy, + observation_space=observation_space, + action_space=action_space, + config=config, + policy_state=policy_state, + policy_mapping_fn=policy_mapping_fn, + policies_to_train=policies_to_train, + module_spec=module_spec, + ) + + # Return newly added policy (from the local EnvRunner). + if add_to_env_runners: + return self.get_policy(policy_id) + elif add_to_eval_env_runners and self.eval_env_runner_group: + return self.eval_env_runner.policy_map[policy_id] + + @OldAPIStack + def remove_policy( + self, + policy_id: PolicyID = DEFAULT_POLICY_ID, + *, + policy_mapping_fn: Optional[Callable[[AgentID], PolicyID]] = None, + policies_to_train: Optional[ + Union[ + Collection[PolicyID], + Callable[[PolicyID, Optional[SampleBatchType]], bool], + ] + ] = None, + remove_from_env_runners: bool = True, + remove_from_eval_env_runners: bool = True, + # Deprecated args. + evaluation_workers=DEPRECATED_VALUE, + remove_from_learners=DEPRECATED_VALUE, + ) -> None: + """Removes a policy from this Algorithm. + + Args: + policy_id: ID of the policy to be removed. + policy_mapping_fn: An optional (updated) policy mapping function + to use from here on. Note that already ongoing episodes will + not change their mapping but will use the old mapping till + the end of the episode. + policies_to_train: An optional list of policy IDs to be trained + or a callable taking PolicyID and SampleBatchType and + returning a bool (trainable or not?). + If None, will keep the existing setup in place. Policies, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + remove_from_env_runners: Whether to remove the Policy from the + EnvRunnerGroup (with its m EnvRunners plus the local one). + remove_from_eval_env_runners: Whether to remove the RLModule from the eval + EnvRunnerGroup (with its o EnvRunners plus the local one). + """ + if evaluation_workers != DEPRECATED_VALUE: + deprecation_warning( + old="Algorithm.remove_policy(evaluation_workers=...)", + new="Algorithm.remove_policy(remove_from_eval_env_runners=...)", + error=False, + ) + remove_from_eval_env_runners = evaluation_workers + if remove_from_learners != DEPRECATED_VALUE: + deprecation_warning( + old="Algorithm.remove_policy(remove_from_learners=..)", + help="Hybrid API stack no longer supported by RLlib!", + error=True, + ) + + def fn(worker): + worker.remove_policy( + policy_id=policy_id, + policy_mapping_fn=policy_mapping_fn, + policies_to_train=policies_to_train, + ) + + # Update all EnvRunner workers. + if remove_from_env_runners: + self.env_runner_group.foreach_env_runner(fn, local_env_runner=True) + + # Update the evaluation worker set's workers, if required. + if remove_from_eval_env_runners and self.eval_env_runner_group is not None: + self.eval_env_runner_group.foreach_env_runner(fn, local_env_runner=True) + + @OldAPIStack + @staticmethod + def from_state(state: Dict) -> "Algorithm": + """Recovers an Algorithm from a state object. + + The `state` of an instantiated Algorithm can be retrieved by calling its + `get_state` method. It contains all information necessary + to create the Algorithm from scratch. No access to the original code (e.g. + configs, knowledge of the Algorithm's class, etc..) is needed. + + Args: + state: The state to recover a new Algorithm instance from. + + Returns: + A new Algorithm instance. + """ + algorithm_class: Type[Algorithm] = state.get("algorithm_class") + if algorithm_class is None: + raise ValueError( + "No `algorithm_class` key was found in given `state`! " + "Cannot create new Algorithm." + ) + # algo_class = get_trainable_cls(algo_class_name) + # Create the new algo. + config = state.get("config") + if not config: + raise ValueError("No `config` found in given Algorithm state!") + new_algo = algorithm_class(config=config) + # Set the new algo's state. + new_algo.__setstate__(state) + + # Return the new algo. + return new_algo + + @OldAPIStack + def export_policy_model( + self, + export_dir: str, + policy_id: PolicyID = DEFAULT_POLICY_ID, + onnx: Optional[int] = None, + ) -> None: + """Exports policy model with given policy_id to a local directory. + + Args: + export_dir: Writable local directory. + policy_id: Optional policy id to export. + onnx: If given, will export model in ONNX format. The + value of this parameter set the ONNX OpSet version to use. + If None, the output format will be DL framework specific. + """ + self.get_policy(policy_id).export_model(export_dir, onnx) + + @OldAPIStack + def export_policy_checkpoint( + self, + export_dir: str, + policy_id: PolicyID = DEFAULT_POLICY_ID, + ) -> None: + """Exports Policy checkpoint to a local directory and returns an AIR Checkpoint. + + Args: + export_dir: Writable local directory to store the AIR Checkpoint + information into. + policy_id: Optional policy ID to export. If not provided, will export + "default_policy". If `policy_id` does not exist in this Algorithm, + will raise a KeyError. + + Raises: + KeyError: if `policy_id` cannot be found in this Algorithm. + """ + policy = self.get_policy(policy_id) + if policy is None: + raise KeyError(f"Policy with ID {policy_id} not found in Algorithm!") + policy.export_checkpoint(export_dir) + + @override(Trainable) + def save_checkpoint(self, checkpoint_dir: str) -> None: + """Exports checkpoint to a local directory. + + The structure of an Algorithm checkpoint dir will be as follows:: + + policies/ + pol_1/ + policy_state.pkl + pol_2/ + policy_state.pkl + learner/ + learner_state.json + module_state/ + module_1/ + ... + optimizer_state/ + optimizers_module_1/ + ... + rllib_checkpoint.json + algorithm_state.pkl + + Note: `rllib_checkpoint.json` contains a "version" key (e.g. with value 0.1) + helping RLlib to remain backward compatible wrt. restoring from checkpoints from + Ray 2.0 onwards. + + Args: + checkpoint_dir: The directory where the checkpoint files will be stored. + """ + with TimerAndPrometheusLogger(self._metrics_save_checkpoint_time): + # New API stack: Delegate to the `Checkpointable` implementation of + # `save_to_path()` and return. + if self.config.enable_rl_module_and_learner: + self.save_to_path( + checkpoint_dir, + use_msgpack=self.config._use_msgpack_checkpoints, + ) + return + + checkpoint_dir = pathlib.Path(checkpoint_dir) + + state = self.__getstate__() + + # Extract policy states from worker state (Policies get their own + # checkpoint sub-dirs). + policy_states = {} + if "worker" in state and "policy_states" in state["worker"]: + policy_states = state["worker"].pop("policy_states", {}) + + # Add RLlib checkpoint version. + if self.config.enable_rl_module_and_learner: + state["checkpoint_version"] = CHECKPOINT_VERSION_LEARNER_AND_ENV_RUNNER + else: + state["checkpoint_version"] = CHECKPOINT_VERSION + + # Write state (w/o policies) to disk. + state_file = checkpoint_dir / "algorithm_state.pkl" + with open(state_file, "wb") as f: + pickle.dump(state, f) + + # Write rllib_checkpoint.json. + with open(checkpoint_dir / "rllib_checkpoint.json", "w") as f: + json.dump( + { + "type": "Algorithm", + "checkpoint_version": str(state["checkpoint_version"]), + "format": "cloudpickle", + "state_file": str(state_file), + "policy_ids": list(policy_states.keys()), + "ray_version": ray.__version__, + "ray_commit": ray.__commit__, + }, + f, + ) + + # Old API stack: Write individual policies to disk, each in their own + # sub-directory. + for pid, policy_state in policy_states.items(): + # From here on, disallow policyIDs that would not work as directory names. + validate_module_id(pid, error=True) + policy_dir = checkpoint_dir / "policies" / pid + os.makedirs(policy_dir, exist_ok=True) + policy = self.get_policy(pid) + policy.export_checkpoint(policy_dir, policy_state=policy_state) + + # If we are using the learner API (hybrid API stack) -> Save the learner group's + # state inside a "learner" subdir. Note that this is not in line with the + # new Checkpointable API, but makes this case backward compatible. + # The new Checkpointable API is only strictly applied anyways to the + # new API stack. + if self.config.enable_rl_module_and_learner: + learner_state_dir = os.path.join(checkpoint_dir, "learner") + self.learner_group.save_to_path(learner_state_dir) + + @override(Trainable) + def load_checkpoint(self, checkpoint_dir: str) -> None: + with TimerAndPrometheusLogger(self._metrics_load_checkpoint_time): + # New API stack: Delegate to the `Checkpointable` implementation of + # `restore_from_path()`. + if self.config.enable_rl_module_and_learner: + self.restore_from_path(checkpoint_dir) + else: + # Checkpoint is provided as a local directory. + # Restore from the checkpoint file or dir. + checkpoint_info = get_checkpoint_info(checkpoint_dir) + checkpoint_data = Algorithm._checkpoint_info_to_algorithm_state( + checkpoint_info + ) + self.__setstate__(checkpoint_data) + + # Call the `on_checkpoint_loaded` callback. + make_callback( + "on_checkpoint_loaded", + callbacks_objects=self.callbacks, + callbacks_functions=self.config.callbacks_on_checkpoint_loaded, + kwargs=dict(algorithm=self), + ) + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + if not self.config.enable_env_runner_and_connector_v2: + raise RuntimeError( + "Algorithm.get_state() not supported on the old API stack! " + "Use Algorithm.__getstate__() instead." + ) + + state = {} + + # Get (local) EnvRunner state (w/o RLModule). + if self.config.is_online: + if self.env_runner: + if self._check_component( + COMPONENT_ENV_RUNNER, components, not_components + ): + state[COMPONENT_ENV_RUNNER] = self.env_runner.get_state( + components=self._get_subcomponents( + COMPONENT_RL_MODULE, components + ), + not_components=force_list( + self._get_subcomponents(COMPONENT_RL_MODULE, not_components) + ) + # We don't want the RLModule state from the EnvRunners (it's + # `inference_only` anyway and already provided in full by the + # Learners). + + [COMPONENT_RL_MODULE], + **kwargs, + ) + else: + if self._check_component( + COMPONENT_ENV_TO_MODULE_CONNECTOR, components, not_components + ): + state[ + COMPONENT_ENV_TO_MODULE_CONNECTOR + ] = self.env_to_module_connector.get_state() + if self._check_component( + COMPONENT_MODULE_TO_ENV_CONNECTOR, components, not_components + ): + state[ + COMPONENT_MODULE_TO_ENV_CONNECTOR + ] = self.module_to_env_connector.get_state() + # Get (local) evaluation EnvRunner state (w/o RLModule). + if self.eval_env_runner and self._check_component( + COMPONENT_EVAL_ENV_RUNNER, components, not_components + ): + state[COMPONENT_EVAL_ENV_RUNNER] = self.eval_env_runner.get_state( + components=self._get_subcomponents(COMPONENT_RL_MODULE, components), + not_components=force_list( + self._get_subcomponents(COMPONENT_RL_MODULE, not_components) + ) + # We don't want the RLModule state from the EnvRunners (it's + # `inference_only` anyway and already provided in full by the Learners). + + [COMPONENT_RL_MODULE], + **kwargs, + ) + + # Get LearnerGroup state (w/ RLModule). + if self._check_component(COMPONENT_LEARNER_GROUP, components, not_components): + state[COMPONENT_LEARNER_GROUP] = self.learner_group.get_state( + components=self._get_subcomponents(COMPONENT_LEARNER_GROUP, components), + not_components=self._get_subcomponents( + COMPONENT_LEARNER_GROUP, not_components + ), + **kwargs, + ) + + # Get entire MetricsLogger state. + # TODO (sven): Make `MetricsLogger` a Checkpointable. + state[COMPONENT_METRICS_LOGGER] = self.metrics.get_state() + + # Save current `training_iteration`. + state[TRAINING_ITERATION] = self.training_iteration + + return state + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + # Set the (training) EnvRunners' states. + if COMPONENT_ENV_RUNNER in state: + if self.env_runner: + self.env_runner.set_state(state[COMPONENT_ENV_RUNNER]) + else: + self.env_to_module_connector.set_state( + state[COMPONENT_ENV_RUNNER][COMPONENT_ENV_TO_MODULE_CONNECTOR] + ) + self.module_to_env_connector.set_state( + state[COMPONENT_ENV_RUNNER][COMPONENT_MODULE_TO_ENV_CONNECTOR] + ) + self.env_runner_group.sync_env_runner_states( + config=self.config, + from_worker=self.env_runner, + env_to_module=self.env_to_module_connector, + module_to_env=self.module_to_env_connector, + ) + + # Set the (eval) EnvRunners' states. + if self.eval_env_runner_group and COMPONENT_EVAL_ENV_RUNNER in state: + if self.eval_env_runner: + self.eval_env_runner.set_state(state[COMPONENT_ENV_RUNNER]) + self.eval_env_runner_group.sync_env_runner_states( + config=self.evaluation_config, + from_worker=self.env_runner, + env_to_module=self.env_to_module_connector, + module_to_env=self.module_to_env_connector, + ) + + # Set the LearnerGroup's state. + if COMPONENT_LEARNER_GROUP in state: + self.learner_group.set_state(state[COMPONENT_LEARNER_GROUP]) + # Sync new weights to all EnvRunners. + self.env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + if self.eval_env_runner_group: + self.eval_env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + + # TODO (sven): Make `MetricsLogger` a Checkpointable. + if COMPONENT_METRICS_LOGGER in state: + self.metrics.set_state(state[COMPONENT_METRICS_LOGGER]) + + if TRAINING_ITERATION in state: + self._iteration = state[TRAINING_ITERATION] + + @override(Checkpointable) + def get_checkpointable_components(self) -> List[Tuple[str, "Checkpointable"]]: + components = [ + (COMPONENT_LEARNER_GROUP, self.learner_group), + ] + if self.config.is_online and self.env_runner: + components.append( + (COMPONENT_ENV_RUNNER, self.env_runner), + ) + elif self.config.is_online and not self.env_runner: + if self.env_to_module_connector: + components.append( + (COMPONENT_ENV_TO_MODULE_CONNECTOR, self.env_to_module_connector), + ) + if self.module_to_env_connector: + components.append( + (COMPONENT_MODULE_TO_ENV_CONNECTOR, self.module_to_env_connector), + ) + if self.eval_env_runner: + components.append( + ( + COMPONENT_EVAL_ENV_RUNNER, + self.eval_env_runner, + ) + ) + return components + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self) -> Tuple[Tuple, Dict[str, Any]]: + return ( + (self.config.get_state(),), # *args, + {}, # **kwargs + ) + + @override(Checkpointable) + def restore_from_path(self, path, *args, **kwargs): + # Override from parent method, b/c we might have to sync the EnvRunner weights + # after having restored/loaded the LearnerGroup state. + super().restore_from_path(path, *args, **kwargs) + + # Sync EnvRunners, if LearnerGroup's checkpoint can be found in path + # or user loaded a subcomponent within the LearnerGroup (for example a module). + path = pathlib.Path(path) + if (path / COMPONENT_LEARNER_GROUP).is_dir() or ( + "component" in kwargs and COMPONENT_LEARNER_GROUP in kwargs["component"] + ): + # Make also sure, all (training) EnvRunners get the just loaded weights, but + # only the inference-only ones. + self.env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + + # If we have remote `EnvRunner`s but no local `EnvRunner` we have to restore states + # from path. + if self.env_runner_group.num_remote_env_runners() > 0 and not self.env_runner: + if (path / COMPONENT_ENV_TO_MODULE_CONNECTOR).is_dir(): + self.env_to_module_connector.restore_from_path( + path / COMPONENT_ENV_TO_MODULE_CONNECTOR, *args, **kwargs + ) + + if (path / COMPONENT_MODULE_TO_ENV_CONNECTOR).is_dir(): + self.module_to_env_connector.restore_from_path( + path / COMPONENT_MODULE_TO_ENV_CONNECTOR, *args, **kwargs + ) + + self.env_runner_group.sync_env_runner_states( + config=self.config, + from_worker=None, + env_steps_sampled=self.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED) + ), + # connector_states=connector_states, + env_to_module=self.env_to_module_connector, + module_to_env=self.module_to_env_connector, + ) + # Otherwise get the connector states from the local `EnvRunner`. + elif self.env_runner_group.num_remote_env_runners() > 0 and self.env_runner: + self.env_runner_group.sync_env_runner_states( + config=self.config, + from_worker=self.env_runner, + ) + + @override(Trainable) + def log_result(self, result: ResultDict) -> None: + # Log after the callback is invoked, so that the user has a chance + # to mutate the result. + # TODO (sven): It might not make sense to pass in the MetricsLogger at this late + # point in time. In here, the result dict has already been "compiled" (reduced) + # by the MetricsLogger and there is probably no point in adding more Stats + # here. + with TimerAndPrometheusLogger(self._metrics_callback_on_train_result_time): + make_callback( + "on_train_result", + callbacks_objects=self.callbacks, + callbacks_functions=self.config.callbacks_on_train_result, + kwargs=dict( + algorithm=self, + metrics_logger=self.metrics, + result=result, + ), + ) + # Then log according to Trainable's logging logic. + Trainable.log_result(self, result) + + @override(Trainable) + def cleanup(self) -> None: + # Stop all Learners. + if hasattr(self, "learner_group") and self.learner_group is not None: + self.learner_group.shutdown() + + # Stop all aggregation actors. + if hasattr(self, "_aggregator_actor_manager") and ( + self._aggregator_actor_manager is not None + ): + self._aggregator_actor_manager.clear() + + # Stop all EnvRunners. + if hasattr(self, "env_runner_group") and self.env_runner_group is not None: + self.env_runner_group.stop() + if ( + hasattr(self, "eval_env_runner_group") + and self.eval_env_runner_group is not None + ): + self.eval_env_runner_group.stop() + + @OverrideToImplementCustomLogic + @classmethod + @override(Trainable) + def default_resource_request( + cls, config: Union[AlgorithmConfig, PartialAlgorithmConfigDict] + ) -> Union[Resources, PlacementGroupFactory]: + config = cls.get_default_config().update_from_dict(config) + config.validate() + config.freeze() + eval_config = config.get_evaluation_config_object() + eval_config.validate() + eval_config.freeze() + + if config.enable_rl_module_and_learner: + main_process = _get_main_process_bundle(config) + else: + main_process = { + "CPU": config.num_cpus_for_main_process, + "GPU": ( + 0 + if config._fake_gpus + else config.num_gpus + if not config.enable_rl_module_and_learner + else 0 + ), + } + + env_runner_bundles = _get_env_runner_bundles(config) + + if cls._should_create_evaluation_env_runners(eval_config): + eval_env_runner_bundles = _get_env_runner_bundles(eval_config) + else: + eval_env_runner_bundles = [] + + if cls._should_create_offline_evaluation_runners(eval_config): + offline_eval_runner_bundles = _get_offline_eval_runner_bundles(eval_config) + else: + offline_eval_runner_bundles = [] + + learner_bundles = [] + if config.enable_rl_module_and_learner: + learner_bundles = _get_learner_bundles(config) + + bundles = ( + [main_process] + + env_runner_bundles + + eval_env_runner_bundles + + offline_eval_runner_bundles + + learner_bundles + ) + + return PlacementGroupFactory( + bundles=bundles, + strategy=config.placement_strategy, + ) + + @DeveloperAPI + def _before_evaluate(self): + """Pre-evaluation callback.""" + pass + + @staticmethod + def _get_env_id_and_creator( + env_specifier: Union[str, EnvType, None], config: AlgorithmConfig + ) -> Tuple[Optional[str], EnvCreator]: + """Returns env_id and creator callable given original env id from config. + + Args: + env_specifier: An env class, an already tune registered env ID, a known + gym env name, or None (if no env is used). + config: The AlgorithmConfig object. + + Returns: + Tuple consisting of a) env ID string and b) env creator callable. + """ + # Environment is specified via a string. + if isinstance(env_specifier, str): + # An already registered env. + if _global_registry.contains(ENV_CREATOR, env_specifier): + return env_specifier, _global_registry.get(ENV_CREATOR, env_specifier) + + # A class path specifier. + elif "." in env_specifier: + + def env_creator_from_classpath(env_context): + try: + env_obj = from_config(env_specifier, env_context) + except ValueError: + raise EnvError( + ERR_MSG_INVALID_ENV_DESCRIPTOR.format(env_specifier) + ) + return env_obj + + return env_specifier, env_creator_from_classpath + # Try gym/PyBullet. + else: + return env_specifier, functools.partial( + _gym_env_creator, env_descriptor=env_specifier + ) + + elif isinstance(env_specifier, type): + env_id = env_specifier # .__name__ + + if config["remote_worker_envs"]: + # Check gym version (0.22 or higher?). + # If > 0.21, can't perform auto-wrapping of the given class as this + # would lead to a pickle error. + gym_version = importlib.metadata.version("gym") + if version.parse(gym_version) >= version.parse("0.22"): + raise ValueError( + "Cannot specify a gym.Env class via `config.env` while setting " + "`config.remote_worker_env=True` AND your gym version is >= " + "0.22! Try installing an older version of gym or set `config." + "remote_worker_env=False`." + ) + + @ray.remote(num_cpus=1) + class _wrapper(env_specifier): + # Add convenience `_get_spaces` and `_is_multi_agent` + # methods: + def _get_spaces(self): + return self.observation_space, self.action_space + + def _is_multi_agent(self): + from ray.rllib.env.multi_agent_env import MultiAgentEnv + + return isinstance(self, MultiAgentEnv) + + return env_id, lambda cfg: _wrapper.remote(cfg) + # gym.Env-subclass: Also go through our RLlib gym-creator. + elif issubclass(env_specifier, gym.Env): + return env_id, functools.partial( + _gym_env_creator, + env_descriptor=env_specifier, + ) + # All other env classes: Call c'tor directly. + else: + return env_id, lambda cfg: env_specifier(cfg) + + # No env -> Env creator always returns None. + elif env_specifier is None: + return None, lambda env_config: None + + else: + raise ValueError( + "{} is an invalid env specifier. ".format(env_specifier) + + "You can specify a custom env as either a class " + '(e.g., YourEnvCls) or a registered env id (e.g., "your_env").' + ) + + def _sync_filters_if_needed( + self, + *, + central_worker: EnvRunner, + workers: EnvRunnerGroup, + config: AlgorithmConfig, + ) -> None: + """Synchronizes the filter stats from `workers` to `central_worker`. + + .. and broadcasts the central_worker's filter stats back to all `workers` + (if configured). + + Args: + central_worker: The worker to sync/aggregate all `workers`' filter stats to + and from which to (possibly) broadcast the updated filter stats back to + `workers`. + workers: The EnvRunnerGroup, whose EnvRunners' filter stats should be used + for aggregation on `central_worker` and which (possibly) get updated + from `central_worker` after the sync. + config: The algorithm config instance. This is used to determine, whether + syncing from `workers` should happen at all and whether broadcasting + back to `workers` (after possible syncing) should happen. + """ + if central_worker and config.observation_filter != "NoFilter": + FilterManager.synchronize( + central_worker.filters, + workers, + update_remote=config.update_worker_filter_stats, + timeout_seconds=config.sync_filters_on_rollout_workers_timeout_s, + use_remote_data_for_update=config.use_worker_filter_stats, + ) + + @classmethod + @override(Trainable) + def resource_help(cls, config: Union[AlgorithmConfig, AlgorithmConfigDict]) -> str: + return ( + "\n\nYou can adjust the resource requests of RLlib Algorithms by calling " + "`AlgorithmConfig.env_runners(" + "num_env_runners=.., num_cpus_per_env_runner=.., " + "num_gpus_per_env_runner=.., ..)` and " + "`AgorithmConfig.learners(num_learners=.., num_gpus_per_learner=..)`. See " + "the `ray.rllib.algorithms.algorithm_config.AlgorithmConfig` classes " + "(each Algorithm has its own subclass of this class) for more info.\n\n" + f"The config of this Algorithm is: {config}" + ) + + @override(Trainable) + def get_auto_filled_metrics( + self, + now: Optional[datetime] = None, + time_this_iter: Optional[float] = None, + timestamp: Optional[int] = None, + debug_metrics_only: bool = False, + ) -> dict: + # Override this method to make sure, the `config` key of the returned results + # contains the proper Tune config dict (instead of an AlgorithmConfig object). + auto_filled = super().get_auto_filled_metrics( + now, time_this_iter, timestamp, debug_metrics_only + ) + if "config" not in auto_filled: + raise KeyError("`config` key not found in auto-filled results dict!") + + # If `config` key is no dict (but AlgorithmConfig object) -> + # make sure, it's a dict to not break Tune APIs. + if not isinstance(auto_filled["config"], dict): + assert isinstance(auto_filled["config"], AlgorithmConfig) + auto_filled["config"] = auto_filled["config"].to_dict() + return auto_filled + + @classmethod + def merge_algorithm_configs( + cls, + config1: AlgorithmConfigDict, + config2: PartialAlgorithmConfigDict, + _allow_unknown_configs: Optional[bool] = None, + ) -> AlgorithmConfigDict: + """Merges a complete Algorithm config dict with a partial override dict. + + Respects nested structures within the config dicts. The values in the + partial override dict take priority. + + Args: + config1: The complete Algorithm's dict to be merged (overridden) + with `config2`. + config2: The partial override config dict to merge on top of + `config1`. + _allow_unknown_configs: If True, keys in `config2` that don't exist + in `config1` are allowed and will be added to the final config. + + Returns: + The merged full algorithm config dict. + """ + config1 = copy.deepcopy(config1) + if "callbacks" in config2 and type(config2["callbacks"]) is dict: + deprecation_warning( + "callbacks dict interface", + "a class extending rllib.callbacks.callbacks.RLlibCallback; " + "see `rllib/examples/metrics/custom_metrics_and_callbacks.py` for an " + "example.", + error=True, + ) + + if _allow_unknown_configs is None: + _allow_unknown_configs = cls._allow_unknown_configs + return deep_update( + config1, + config2, + _allow_unknown_configs, + cls._allow_unknown_subkeys, + cls._override_all_subkeys_if_type_changes, + cls._override_all_key_list, + ) + + @staticmethod + @ExperimentalAPI + def validate_env(env: EnvType, env_context: EnvContext) -> None: + """Env validator function for this Algorithm class. + + Override this in child classes to define custom validation + behavior. + + Args: + env: The (sub-)environment to validate. This is normally a + single sub-environment (e.g. a gym.Env) within a vectorized + setup. + env_context: The EnvContext to configure the environment. + + Raises: + Exception: in case something is wrong with the given environment. + """ + pass + + def _run_one_training_iteration(self) -> Tuple[ResultDict, "TrainIterCtx"]: + """Runs one training iteration (`self.iteration` will be +1 after this). + + Calls `self.training_step()` repeatedly until the configured minimum time (sec), + minimum sample- or minimum training steps have been reached. + + Returns: + The ResultDict from the last call to `training_step()`. Note that even + though we only return the last ResultDict, the user still has full control + over the history and reduce behavior of individual metrics at the time these + metrics are logged with `self.metrics.log_...()`. + """ + with TimerAndPrometheusLogger(self._metrics_run_one_training_iteration_time): + with self.metrics.log_time((TIMERS, TRAINING_ITERATION_TIMER)): + # In case we are training (in a thread) parallel to evaluation, + # we may have to re-enable eager mode here (gets disabled in the + # thread). + if self.config.get("framework") == "tf2" and not tf.executing_eagerly(): + tf1.enable_eager_execution() + + has_run_once = False + # Create a step context ... + with TrainIterCtx(algo=self) as train_iter_ctx: + # .. so we can query it whether we should stop the iteration loop (e.g. + # when we have reached `min_time_s_per_iteration`). + while not train_iter_ctx.should_stop(has_run_once): + # Before training step, try to bring failed workers back. + with self.metrics.log_time((TIMERS, RESTORE_ENV_RUNNERS_TIMER)): + restored = self.restore_env_runners(self.env_runner_group) + # Fire the callback for re-created EnvRunners. + if restored: + self._make_on_env_runners_recreated_callbacks( + config=self.config, + env_runner_group=self.env_runner_group, + restored_env_runner_indices=restored, + ) + + # Try to train one step. + with self.metrics.log_time((TIMERS, TRAINING_STEP_TIMER)): + with TimerAndPrometheusLogger( + self._metrics_training_step_time + ): + training_step_return_value = self.training_step() + has_run_once = True + + # On the new API stack, results should NOT be returned anymore as + # a dict, but purely logged through the `MetricsLogger` API. This + # way, we make sure to never miss a single stats/counter/timer + # when calling `self.training_step()` more than once within the same + # iteration. + if training_step_return_value is not None: + raise ValueError( + "`Algorithm.training_step()` should NOT return a result " + "dict anymore on the new API stack! Instead, log all " + "results, timers, counters through the `self.metrics` " + "(MetricsLogger) instance of the Algorithm and return " + "None. The logged results are compiled automatically into " + "one single result dict per training iteration." + ) + + # TODO (sven): Resolve this metric through log_time's future + # ability to compute throughput. + self.metrics.log_value( + NUM_TRAINING_STEP_CALLS_PER_ITERATION, + 1, + reduce="sum", + clear_on_reduce=True, + ) + + if self.config.num_aggregator_actors_per_learner: + remote_aggregator_metrics = self._aggregator_actor_manager.foreach_actor_async_fetch_ready( + func=lambda actor: actor.get_metrics(), + tag="metrics", + timeout_seconds=0.0, + return_obj_refs=False, + # (Artur) TODO: In the future, we want to make aggregator actors fault tolerant and should make this configurable + ignore_ray_errors=False, + ) + + self.metrics.aggregate( + remote_aggregator_metrics, + key=AGGREGATOR_ACTOR_RESULTS, + ) + + # Only here (at the end of the iteration), compile the results into a single result dict. + # Calling compile here reduces the metrics into single values and adds throughputs to the results where applicable. + compiled_metrics = self.metrics.compile() + + return compiled_metrics, train_iter_ctx + + def _run_one_offline_evaluation(self): + """Runs offline evaluation step via `self.offline_evaluate()` and handling runner + failures. + + Returns: + The results dict from the offline evaluation call. + """ + # Restore crashed offline evaluation runners. + if self.offline_eval_runner_group is not None: + with self.metrics.log_time((TIMERS, RESTORE_OFFLINE_EVAL_RUNNERS_TIMER)): + restored = self.restore_offline_eval_runners( + self.offline_eval_runner_group + ) + if restored: + # Fire the callback for re-created workers. + make_callback( + "on_offline_eval_runners_recreated", + callbacks_objects=self.callbacks, + callbacks_functions=( + self.config.callbacks_on_offline_eval_runners_recreated + ), + kwargs=dict( + algorithm=self, + env_runner_group=self.offline_eval_runner_group, + env_runner_indices=restored, + ), + ) + + # Run one offline evaluation and time it. + with self.metrics.log_time((TIMERS, OFFLINE_EVALUATION_ITERATION_TIMER)): + eval_results = self.evaluate_offline() + + # After evaluation, do a round of health check on remote eval runners to see if + # any of the failed runners are back. + if self.offline_eval_runner_group is not None: + # Add number of healthy evaluation runners after this iteration. + eval_results[ + "num_healthy_offline_eval_runners" + ] = self.offline_eval_runner_group.num_healthy_remote_runners + eval_results[ + "offline_runners_actor_manager_num_outstanding_async_reqs" + ] = self.offline_eval_runner_group.num_in_flight_async_reqs + eval_results[ + "num_remote_offline_eval_runners_restarts" + ] = self.offline_eval_runner_group.num_remote_runner_restarts + + return {EVALUATION_RESULTS: eval_results} + + def _run_one_evaluation( + self, + parallel_train_future: Optional[concurrent.futures.ThreadPoolExecutor] = None, + ) -> ResultDict: + """Runs evaluation step via `self.evaluate()` and handling worker failures. + + Args: + parallel_train_future: In case, we are training and avaluating in parallel, + this arg carries the currently running ThreadPoolExecutor object that + runs the training iteration. Use `parallel_train_future.done()` to + check, whether the parallel training job has completed and + `parallel_train_future.result()` to get its return values. + + Returns: + The results dict from the evaluation call. + """ + with TimerAndPrometheusLogger(self._metrics_run_one_evaluation_time): + if self.eval_env_runner_group is not None: + if self.config.enable_env_runner_and_connector_v2: + with self.metrics.log_time( + (TIMERS, RESTORE_EVAL_ENV_RUNNERS_TIMER) + ): + restored = self.restore_env_runners(self.eval_env_runner_group) + else: + with self._timers["restore_eval_workers"]: + restored = self.restore_env_runners(self.eval_env_runner_group) + # Fire the callback for re-created EnvRunners. + if restored: + self._make_on_env_runners_recreated_callbacks( + config=self.evaluation_config, + env_runner_group=self.eval_env_runner_group, + restored_env_runner_indices=restored, + ) + + # Run `self.evaluate()` only once per training iteration. + if self.config.enable_env_runner_and_connector_v2: + with self.metrics.log_time((TIMERS, EVALUATION_ITERATION_TIMER)): + eval_results = self.evaluate( + parallel_train_future=parallel_train_future + ) + else: + with self._timers[EVALUATION_ITERATION_TIMER]: + eval_results = self.evaluate( + parallel_train_future=parallel_train_future + ) + self._timers[EVALUATION_ITERATION_TIMER].push_units_processed( + self._counters[NUM_ENV_STEPS_SAMPLED_FOR_EVALUATION_THIS_ITER] + ) + + # After evaluation, do a round of health check on remote eval workers to see if + # any of the failed workers are back. + if self.eval_env_runner_group is not None: + # Add number of healthy evaluation workers after this iteration. + eval_results[ + "num_healthy_workers" + ] = self.eval_env_runner_group.num_healthy_remote_workers() + eval_results[ + "actor_manager_num_outstanding_async_reqs" + ] = self.eval_env_runner_group.num_in_flight_async_reqs() + eval_results[ + "num_remote_worker_restarts" + ] = self.eval_env_runner_group.num_remote_worker_restarts() + + return {EVALUATION_RESULTS: eval_results} + + def _run_one_training_iteration_and_evaluation_in_parallel( + self, + ) -> Tuple[ResultDict, ResultDict, "TrainIterCtx"]: + """Runs one training iteration and one evaluation step in parallel. + + First starts the training iteration (via `self._run_one_training_iteration()`) + within a ThreadPoolExecutor, then runs the evaluation step in parallel. + In auto-duration mode (config.evaluation_duration=auto), makes sure the + evaluation step takes roughly the same time as the training iteration. + + Returns: + A tuple containing the training results, the evaluation results, and + the `TrainIterCtx` object returned by the training call. + """ + with concurrent.futures.ThreadPoolExecutor() as executor: + + if self.config.enable_env_runner_and_connector_v2: + parallel_train_future = executor.submit( + lambda: self._run_one_training_iteration() + ) + else: + parallel_train_future = executor.submit( + lambda: self._run_one_training_iteration_old_api_stack() + ) + + # Pass the train_future into `self._run_one_evaluation()` to allow it + # to run exactly as long as the training iteration takes in case + # evaluation_duration=auto. + evaluation_results = self._run_one_evaluation( + parallel_train_future=parallel_train_future + ) + # Collect the training results from the future. + train_results, train_iter_ctx = parallel_train_future.result() + + return train_results, evaluation_results, train_iter_ctx + + def _run_offline_evaluation_old_api_stack(self): + """Runs offline evaluation via `OfflineEvaluator.estimate_on_dataset()` API. + + This method will be used when `evaluation_dataset` is provided. + Note: This will only work if the policy is a single agent policy. + + Returns: + The results dict from the offline evaluation call. + """ + assert len(self.env_runner_group.local_env_runner.policy_map) == 1 + + parallelism = self.evaluation_config.evaluation_num_env_runners or 1 + offline_eval_results = {"off_policy_estimator": {}} + for evaluator_name, offline_evaluator in self.reward_estimators.items(): + offline_eval_results["off_policy_estimator"][ + evaluator_name + ] = offline_evaluator.estimate_on_dataset( + self.evaluation_dataset, + n_parallelism=parallelism, + ) + return offline_eval_results + + @classmethod + def _should_create_evaluation_env_runners(cls, eval_config: "AlgorithmConfig"): + """Determines whether we need to create evaluation workers. + + Returns False if we need to run offline evaluation + (with ope.estimate_on_dastaset API) or when local worker is to be used for + evaluation. Note: We only use estimate_on_dataset API with bandits for now. + That is when ope_split_batch_by_episode is False. + TODO: In future we will do the same for episodic RL OPE. + """ + run_offline_evaluation = ( + eval_config.off_policy_estimation_methods + and not eval_config.ope_split_batch_by_episode + ) + return not run_offline_evaluation and ( + eval_config.evaluation_num_env_runners > 0 + or eval_config.evaluation_interval + ) + + # TODO (simon, sven): Flexibilize the different env/offline components and move + # away from the currently hard-coded: (1) eval `EnvRunnerGroup`, (2) OfflineData + # and (3) `OfflineEvaluationRunnerGroup`. + @classmethod + def _should_create_offline_evaluation_runners(cls, eval_config: "AlgorithmConfig"): + """Determines whether we need to create offline evaluation workers.""" + + return ( + eval_config.offline_evaluation_interval is not None + or eval_config.num_offline_eval_runners > 0 + ) + + def _compile_iteration_results(self, *, train_results, eval_results): + with TimerAndPrometheusLogger(self._metrics_compile_iteration_results_time): + # Error if users still use `self._timers`. + if self._timers: + raise ValueError( + "`Algorithm._timers` is no longer supported on the new API stack! " + "Instead, use `Algorithm.metrics.log_time(" + "[some key (str) or nested key sequence (tuple)])`, e.g. inside your " + "custom `training_step()` method, do: " + "`with self.metrics.log_time(('timers', 'my_block_to_be_timed')): ...`" + ) + + # Return dict (shallow copy of `train_results`). + results: ResultDict = train_results.copy() + # Backward compatibility `NUM_ENV_STEPS_SAMPLED_LIFETIME` is now: + # `ENV_RUNNER_RESULTS/NUM_ENV_STEPS_SAMPLED_LIFETIME`. + results[NUM_ENV_STEPS_SAMPLED_LIFETIME] = results.get( + ENV_RUNNER_RESULTS, {} + ).get(NUM_ENV_STEPS_SAMPLED_LIFETIME, 0) + + # Evaluation results. + if eval_results: + assert ( + isinstance(eval_results, dict) + and len(eval_results) == 1 + and EVALUATION_RESULTS in eval_results + ) + results.update(eval_results) + + # EnvRunner actors fault tolerance stats. + if self.env_runner_group: + results[FAULT_TOLERANCE_STATS] = { + "num_healthy_workers": ( + self.env_runner_group.num_healthy_remote_workers() + ), + "num_remote_worker_restarts": ( + self.env_runner_group.num_remote_worker_restarts() + ), + } + results["env_runner_group"] = { + "actor_manager_num_outstanding_async_reqs": ( + self.env_runner_group.num_in_flight_async_reqs() + ), + } + + return results + + def _make_on_env_runners_recreated_callbacks( + self, + *, + config, + env_runner_group, + restored_env_runner_indices, + ): + make_callback( + "on_env_runners_recreated", + callbacks_objects=self.callbacks, + callbacks_functions=(config.callbacks_on_env_runners_recreated), + kwargs=dict( + algorithm=self, + env_runner_group=env_runner_group, + env_runner_indices=restored_env_runner_indices, + is_evaluation=config.in_evaluation, + ), + ) + # TODO (sven): Deprecate this call. + make_callback( + "on_workers_recreated", + callbacks_objects=self.callbacks, + kwargs=dict( + algorithm=self, + worker_set=env_runner_group, + worker_ids=restored_env_runner_indices, + is_evaluation=config.in_evaluation, + ), + ) + + def __repr__(self): + if self.config.enable_rl_module_and_learner: + return ( + f"{type(self).__name__}(" + f"env={self.config.env}; env-runners={self.config.num_env_runners}; " + f"learners={self.config.num_learners}; " + f"multi-agent={self.config.is_multi_agent}" + f")" + ) + else: + return type(self).__name__ + + @property + def env_runner(self): + """The local EnvRunner instance within the algo's EnvRunnerGroup.""" + if self.env_runner_group: + return self.env_runner_group.local_env_runner + return None + + @property + def eval_env_runner(self): + """The local EnvRunner instance within the algo's evaluation EnvRunnerGroup.""" + if self.eval_env_runner_group: + return self.eval_env_runner_group.local_env_runner + return None + + def _record_usage(self, config): + """Record the framework and algorithm used. + + Args: + config: Algorithm config dict. + """ + record_extra_usage_tag(TagKey.RLLIB_FRAMEWORK, config["framework"]) + record_extra_usage_tag(TagKey.RLLIB_NUM_WORKERS, str(config["num_env_runners"])) + alg = self.__class__.__name__ + # We do not want to collect user defined algorithm names. + if alg not in ALL_ALGORITHMS: + alg = "USER_DEFINED" + record_extra_usage_tag(TagKey.RLLIB_ALGORITHM, alg) + + @OldAPIStack + def _export_model( + self, export_formats: List[str], export_dir: str + ) -> Dict[str, str]: + ExportFormat.validate(export_formats) + exported = {} + if ExportFormat.CHECKPOINT in export_formats: + path = os.path.join(export_dir, ExportFormat.CHECKPOINT) + self.export_policy_checkpoint(path) + exported[ExportFormat.CHECKPOINT] = path + if ExportFormat.MODEL in export_formats: + path = os.path.join(export_dir, ExportFormat.MODEL) + self.export_policy_model(path) + exported[ExportFormat.MODEL] = path + if ExportFormat.ONNX in export_formats: + path = os.path.join(export_dir, ExportFormat.ONNX) + self.export_policy_model(path, onnx=int(os.getenv("ONNX_OPSET", "11"))) + exported[ExportFormat.ONNX] = path + return exported + + @OldAPIStack + def __getstate__(self) -> Dict: + """Returns current state of Algorithm, sufficient to restore it from scratch. + + Returns: + The current state dict of this Algorithm, which can be used to sufficiently + restore the algorithm from scratch without any other information. + """ + if self.config.enable_env_runner_and_connector_v2: + raise RuntimeError( + "Algorithm.__getstate__() not supported anymore on the new API stack! " + "Use Algorithm.get_state() instead." + ) + + # Add config to state so complete Algorithm can be reproduced w/o it. + state = { + "algorithm_class": type(self), + "config": self.config.get_state(), + } + + if hasattr(self, "env_runner_group"): + state["worker"] = self.env_runner_group.local_env_runner.get_state() + + # Also store eval `policy_mapping_fn` (in case it's different from main + # one). Note, the new `EnvRunner API` has no policy mapping function. + if ( + hasattr(self, "eval_env_runner_group") + and self.eval_env_runner_group is not None + ): + state["eval_policy_mapping_fn"] = self.eval_env_runner.policy_mapping_fn + + # Save counters. + state["counters"] = self._counters + + # TODO: Experimental functionality: Store contents of replay buffer + # to checkpoint, only if user has configured this. + if self.local_replay_buffer is not None and self.config.get( + "store_buffer_in_checkpoints" + ): + state["local_replay_buffer"] = self.local_replay_buffer.get_state() + + # Save current `training_iteration`. + state[TRAINING_ITERATION] = self.training_iteration + + return state + + @OldAPIStack + def __setstate__(self, state) -> None: + """Sets the algorithm to the provided state. + + Args: + state: The state dict to restore this Algorithm instance to. `state` may + have been returned by a call to an Algorithm's `__getstate__()` method. + """ + if self.config.enable_env_runner_and_connector_v2: + raise RuntimeError( + "Algorithm.__setstate__() not supported anymore on the new API stack! " + "Use Algorithm.set_state() instead." + ) + + # Old API stack: The local worker stores its state (together with all the + # Module information) in state['worker']. + if hasattr(self, "env_runner_group") and "worker" in state and state["worker"]: + self.env_runner.set_state(state["worker"]) + remote_state_ref = ray.put(state["worker"]) + self.env_runner_group.foreach_env_runner( + lambda w: w.set_state(ray.get(remote_state_ref)), + local_env_runner=False, + ) + if self.eval_env_runner_group: + # Avoid `state` being pickled into the remote function below. + _eval_policy_mapping_fn = state.get("eval_policy_mapping_fn") + + def _setup_eval_worker(w): + w.set_state(ray.get(remote_state_ref)) + # Override `policy_mapping_fn` as it might be different for eval + # workers. + w.set_policy_mapping_fn(_eval_policy_mapping_fn) + + # If evaluation workers are used, also restore the policies + # there in case they are used for evaluation purpose. + self.eval_env_runner_group.foreach_env_runner(_setup_eval_worker) + + # Restore replay buffer data. + if self.local_replay_buffer is not None: + # TODO: Experimental functionality: Restore contents of replay + # buffer from checkpoint, only if user has configured this. + if self.config.store_buffer_in_checkpoints: + if "local_replay_buffer" in state: + self.local_replay_buffer.set_state(state["local_replay_buffer"]) + else: + logger.warning( + "`store_buffer_in_checkpoints` is True, but no replay " + "data found in state!" + ) + elif "local_replay_buffer" in state and log_once( + "no_store_buffer_in_checkpoints_but_data_found" + ): + logger.warning( + "`store_buffer_in_checkpoints` is False, but some replay " + "data found in state!" + ) + + if "counters" in state: + self._counters = state["counters"] + + if TRAINING_ITERATION in state: + self._iteration = state[TRAINING_ITERATION] + + @OldAPIStack + @staticmethod + def _checkpoint_info_to_algorithm_state( + checkpoint_info: dict, + *, + policy_ids: Optional[Collection[PolicyID]] = None, + policy_mapping_fn: Optional[Callable[[AgentID, EpisodeID], PolicyID]] = None, + policies_to_train: Optional[ + Union[ + Collection[PolicyID], + Callable[[PolicyID, Optional[SampleBatchType]], bool], + ] + ] = None, + ) -> Dict: + """Converts a checkpoint info or object to a proper Algorithm state dict. + + The returned state dict can be used inside self.__setstate__(). + + Args: + checkpoint_info: A checkpoint info dict as returned by + `ray.rllib.utils.checkpoints.get_checkpoint_info( + [checkpoint dir or AIR Checkpoint])`. + policy_ids: Optional list/set of PolicyIDs. If not None, only those policies + listed here will be included in the returned state. Note that + state items such as filters, the `is_policy_to_train` function, as + well as the multi-agent `policy_ids` dict will be adjusted as well, + based on this arg. + policy_mapping_fn: An optional (updated) policy mapping function + to include in the returned state. + policies_to_train: An optional list of policy IDs to be trained + or a callable taking PolicyID and SampleBatchType and + returning a bool (trainable or not?) to include in the returned state. + + Returns: + The state dict usable within the `self.__setstate__()` method. + """ + if checkpoint_info["type"] != "Algorithm": + raise ValueError( + "`checkpoint` arg passed to " + "`Algorithm._checkpoint_info_to_algorithm_state()` must be an " + f"Algorithm checkpoint (but is {checkpoint_info['type']})!" + ) + + msgpack = None + if checkpoint_info.get("format") == "msgpack": + msgpack = try_import_msgpack(error=True) + + with open(checkpoint_info["state_file"], "rb") as f: + if msgpack is not None: + data = f.read() + state = msgpack.unpackb(data, raw=False) + else: + state = pickle.load(f) + + # Old API stack: Policies are in separate sub-dirs. + if ( + checkpoint_info["checkpoint_version"] > version.Version("0.1") + and state.get("worker") is not None + and state.get("worker") + ): + worker_state = state["worker"] + + # Retrieve the set of all required policy IDs. + policy_ids = set( + policy_ids if policy_ids is not None else worker_state["policy_ids"] + ) + + # Remove those policies entirely from filters that are not in + # `policy_ids`. + worker_state["filters"] = { + pid: filter + for pid, filter in worker_state["filters"].items() + if pid in policy_ids + } + + # Get Algorithm class. + if isinstance(state["algorithm_class"], str): + # Try deserializing from a full classpath. + # Or as a last resort: Tune registered algorithm name. + state["algorithm_class"] = deserialize_type( + state["algorithm_class"] + ) or get_trainable_cls(state["algorithm_class"]) + # Compile actual config object. + default_config = state["algorithm_class"].get_default_config() + if isinstance(default_config, AlgorithmConfig): + new_config = default_config.update_from_dict(state["config"]) + else: + new_config = Algorithm.merge_algorithm_configs( + default_config, state["config"] + ) + + # Remove policies from multiagent dict that are not in `policy_ids`. + new_policies = new_config.policies + if isinstance(new_policies, (set, list, tuple)): + new_policies = {pid for pid in new_policies if pid in policy_ids} + else: + new_policies = { + pid: spec for pid, spec in new_policies.items() if pid in policy_ids + } + new_config.multi_agent( + policies=new_policies, + policies_to_train=policies_to_train, + **( + {"policy_mapping_fn": policy_mapping_fn} + if policy_mapping_fn is not None + else {} + ), + ) + state["config"] = new_config + + # Prepare local `worker` state to add policies' states into it, + # read from separate policy checkpoint files. + worker_state["policy_states"] = {} + for pid in policy_ids: + policy_state_file = os.path.join( + checkpoint_info["checkpoint_dir"], + "policies", + pid, + "policy_state." + + ("msgpck" if checkpoint_info["format"] == "msgpack" else "pkl"), + ) + if not os.path.isfile(policy_state_file): + raise ValueError( + "Given checkpoint does not seem to be valid! No policy " + f"state file found for PID={pid}. " + f"The file not found is: {policy_state_file}." + ) + + with open(policy_state_file, "rb") as f: + if msgpack is not None: + worker_state["policy_states"][pid] = msgpack.load(f) + else: + worker_state["policy_states"][pid] = pickle.load(f) + + # These two functions are never serialized in a msgpack checkpoint (which + # does not store code, unlike a cloudpickle checkpoint). Hence the user has + # to provide them with the `Algorithm.from_checkpoint()` call. + if policy_mapping_fn is not None: + worker_state["policy_mapping_fn"] = policy_mapping_fn + if ( + policies_to_train is not None + # `policies_to_train` might be left None in case all policies should be + # trained. + or worker_state["is_policy_to_train"] == NOT_SERIALIZABLE + ): + worker_state["is_policy_to_train"] = policies_to_train + + if state["config"].enable_rl_module_and_learner: + state["learner_state_dir"] = os.path.join( + checkpoint_info["checkpoint_dir"], "learner" + ) + + return state + + @OldAPIStack + def _create_local_replay_buffer_if_necessary( + self, config: PartialAlgorithmConfigDict + ) -> Optional[MultiAgentReplayBuffer]: + """Create a MultiAgentReplayBuffer instance if necessary. + + Args: + config: Algorithm-specific configuration data. + + Returns: + MultiAgentReplayBuffer instance based on algorithm config. + None, if local replay buffer is not needed. + """ + if not config.get("replay_buffer_config") or config["replay_buffer_config"].get( + "no_local_replay_buffer" + ): + return + + # Add parameters, if necessary. + if "EpisodeReplayBuffer" in config["replay_buffer_config"]["type"]: + # TODO (simon): Subclassing needs a proper class and therefore + # we need at this moment the string checking. Because we add + # this keyword argument the old stack ReplayBuffer constructors + # will exit with an error b/c tje keyword argument is unknown to them. + config["replay_buffer_config"][ + "metrics_num_episodes_for_smoothing" + ] = self.config.metrics_num_episodes_for_smoothing + + return from_config(ReplayBuffer, config["replay_buffer_config"]) + + @OldAPIStack + def _run_one_training_iteration_old_api_stack(self): + with self._timers[TRAINING_ITERATION_TIMER]: + if self.config.get("framework") == "tf2" and not tf.executing_eagerly(): + tf1.enable_eager_execution() + + results = {} + training_step_results = None + with TrainIterCtx(algo=self) as train_iter_ctx: + while not train_iter_ctx.should_stop(training_step_results): + with self._timers["restore_workers"]: + restored = self.restore_env_runners(self.env_runner_group) + # Fire the callback for re-created EnvRunners. + if restored: + self._make_on_env_runners_recreated_callbacks( + config=self.config, + env_runner_group=self.env_runner_group, + restored_env_runner_indices=restored, + ) + + with self._timers[TRAINING_STEP_TIMER]: + training_step_results = self.training_step() + + if training_step_results: + results = training_step_results + + return results, train_iter_ctx + + @OldAPIStack + def _compile_iteration_results_old_api_stack( + self, *, episodes_this_iter, step_ctx, iteration_results + ): + # Results to be returned. + results: ResultDict = {} + + # Evaluation results. + if "evaluation" in iteration_results: + eval_results = iteration_results.pop("evaluation") + iteration_results.pop(EVALUATION_RESULTS, None) + results["evaluation"] = results[EVALUATION_RESULTS] = eval_results + + # Custom metrics and episode media. + results["custom_metrics"] = iteration_results.pop("custom_metrics", {}) + results["episode_media"] = iteration_results.pop("episode_media", {}) + + # Learner info. + results["info"] = {LEARNER_INFO: iteration_results} + + # Calculate how many (if any) of older, historical episodes we have to add to + # `episodes_this_iter` in order to reach the required smoothing window. + episodes_for_metrics = episodes_this_iter[:] + missing = self.config.metrics_num_episodes_for_smoothing - len( + episodes_this_iter + ) + # We have to add some older episodes to reach the smoothing window size. + if missing > 0: + episodes_for_metrics = self._episode_history[-missing:] + episodes_this_iter + assert ( + len(episodes_for_metrics) + <= self.config.metrics_num_episodes_for_smoothing + ) + # Note that when there are more than `metrics_num_episodes_for_smoothing` + # episodes in `episodes_for_metrics`, leave them as-is. In this case, we'll + # compute the stats over that larger number. + + # Add new episodes to our history and make sure it doesn't grow larger than + # needed. + self._episode_history.extend(episodes_this_iter) + self._episode_history = self._episode_history[ + -self.config.metrics_num_episodes_for_smoothing : + ] + results[ENV_RUNNER_RESULTS] = summarize_episodes( + episodes_for_metrics, + episodes_this_iter, + self.config.keep_per_episode_custom_metrics, + ) + + results[ + "num_healthy_workers" + ] = self.env_runner_group.num_healthy_remote_workers() + results[ + "actor_manager_num_outstanding_async_reqs" + ] = self.env_runner_group.num_in_flight_async_reqs() + results[ + "num_remote_worker_restarts" + ] = self.env_runner_group.num_remote_worker_restarts() + + # Train-steps- and env/agent-steps this iteration. + for c in [ + NUM_AGENT_STEPS_SAMPLED, + NUM_AGENT_STEPS_TRAINED, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_TRAINED, + ]: + results[c] = self._counters[c] + time_taken_sec = step_ctx.get_time_taken_sec() + if self.config.count_steps_by == "agent_steps": + results[NUM_AGENT_STEPS_SAMPLED + "_this_iter"] = step_ctx.sampled + results[NUM_AGENT_STEPS_TRAINED + "_this_iter"] = step_ctx.trained + results[NUM_AGENT_STEPS_SAMPLED + "_throughput_per_sec"] = ( + step_ctx.sampled / time_taken_sec + ) + results[NUM_AGENT_STEPS_TRAINED + "_throughput_per_sec"] = ( + step_ctx.trained / time_taken_sec + ) + # TODO: For CQL and other algos, count by trained steps. + results["timesteps_total"] = self._counters[NUM_AGENT_STEPS_SAMPLED] + else: + results[NUM_ENV_STEPS_SAMPLED + "_this_iter"] = step_ctx.sampled + results[NUM_ENV_STEPS_TRAINED + "_this_iter"] = step_ctx.trained + results[NUM_ENV_STEPS_SAMPLED + "_throughput_per_sec"] = ( + step_ctx.sampled / time_taken_sec + ) + results[NUM_ENV_STEPS_TRAINED + "_throughput_per_sec"] = ( + step_ctx.trained / time_taken_sec + ) + # TODO: For CQL and other algos, count by trained steps. + results["timesteps_total"] = self._counters[NUM_ENV_STEPS_SAMPLED] + + # Forward compatibility with new API stack. + results[NUM_ENV_STEPS_SAMPLED_LIFETIME] = results["timesteps_total"] + results[NUM_AGENT_STEPS_SAMPLED_LIFETIME] = self._counters[ + NUM_AGENT_STEPS_SAMPLED + ] + + # TODO: Backward compatibility. + results[STEPS_TRAINED_THIS_ITER_COUNTER] = step_ctx.trained + results["agent_timesteps_total"] = self._counters[NUM_AGENT_STEPS_SAMPLED] + + # Process timer results. + timers = {} + for k, timer in self._timers.items(): + timers["{}_time_ms".format(k)] = round(timer.mean * 1000, 3) + if timer.has_units_processed(): + timers["{}_throughput".format(k)] = round(timer.mean_throughput, 3) + results["timers"] = timers + + # Process counter results. + counters = {} + for k, counter in self._counters.items(): + counters[k] = counter + results["counters"] = counters + # TODO: Backward compatibility. + results["info"].update(counters) + + return results + + @OldAPIStack + @Deprecated( + help="`Algorithm.compute_single_action` should no longer be used. Get the " + "RLModule instance through `Algorithm.get_module([module ID])`, then compute " + "actions through `RLModule.forward_inference({'obs': [obs batch]})`.", + error=False, + ) + def compute_single_action( + self, + observation: Optional[TensorStructType] = None, + state: Optional[List[TensorStructType]] = None, + *, + prev_action: Optional[TensorStructType] = None, + prev_reward: Optional[float] = None, + info: Optional[EnvInfoDict] = None, + input_dict: Optional[SampleBatch] = None, + policy_id: PolicyID = DEFAULT_POLICY_ID, + full_fetch: bool = False, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + episode=None, + unsquash_action: Optional[bool] = None, + clip_action: Optional[bool] = None, + ) -> Union[ + TensorStructType, + Tuple[TensorStructType, List[TensorType], Dict[str, TensorType]], + ]: + if unsquash_action is None: + unsquash_action = self.config.normalize_actions + elif clip_action is None: + clip_action = self.config.clip_actions + + err_msg = ( + "Provide either `input_dict` OR [`observation`, ...] as " + "args to `Algorithm.compute_single_action()`!" + ) + if input_dict is not None: + assert ( + observation is None + and prev_action is None + and prev_reward is None + and state is None + ), err_msg + observation = input_dict[Columns.OBS] + else: + assert observation is not None, err_msg + + policy = self.get_policy(policy_id) + if policy is None: + raise KeyError( + f"PolicyID '{policy_id}' not found in PolicyMap of the " + f"Algorithm's local worker!" + ) + pp = policy.agent_connectors[ObsPreprocessorConnector] + + if not isinstance(observation, (np.ndarray, dict, tuple)): + try: + observation = np.asarray(observation) + except Exception: + raise ValueError( + f"Observation type {type(observation)} cannot be converted to " + f"np.ndarray." + ) + if pp: + assert len(pp) == 1, "Only one preprocessor should be in the pipeline" + pp = pp[0] + + if not pp.is_identity(): + pp.in_eval() + if observation is not None: + _input_dict = {Columns.OBS: observation} + elif input_dict is not None: + _input_dict = {Columns.OBS: input_dict[Columns.OBS]} + else: + raise ValueError( + "Either observation or input_dict must be provided." + ) + + acd = AgentConnectorDataType("0", "0", _input_dict) + pp.reset(env_id="0") + ac_o = pp([acd])[0] + observation = ac_o.data[Columns.OBS] + + if input_dict is not None: + input_dict[Columns.OBS] = observation + action, state, extra = policy.compute_single_action( + input_dict=input_dict, + explore=explore, + timestep=timestep, + episode=episode, + ) + else: + action, state, extra = policy.compute_single_action( + obs=observation, + state=state, + prev_action=prev_action, + prev_reward=prev_reward, + info=info, + explore=explore, + timestep=timestep, + episode=episode, + ) + + if unsquash_action: + action = space_utils.unsquash_action(action, policy.action_space_struct) + elif clip_action: + action = space_utils.clip_action(action, policy.action_space_struct) + + if state or full_fetch: + return action, state, extra + else: + return action + + @OldAPIStack + @Deprecated( + help="`Algorithm.compute_actions` should no longer be used. Get the RLModule " + "instance through `Algorithm.get_module([module ID])`, then compute actions " + "through `RLModule.forward_inference({'obs': [obs batch]})`.", + error=False, + ) + def compute_actions( + self, + observations: TensorStructType, + state: Optional[List[TensorStructType]] = None, + *, + prev_action: Optional[TensorStructType] = None, + prev_reward: Optional[TensorStructType] = None, + info: Optional[EnvInfoDict] = None, + policy_id: PolicyID = DEFAULT_POLICY_ID, + full_fetch: bool = False, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + episodes=None, + unsquash_actions: Optional[bool] = None, + clip_actions: Optional[bool] = None, + ): + if unsquash_actions is None: + unsquash_actions = self.config.normalize_actions + elif clip_actions is None: + clip_actions = self.config.clip_actions + + state_defined = state is not None + policy = self.get_policy(policy_id) + filtered_obs, filtered_state = [], [] + for agent_id, ob in observations.items(): + worker = self.env_runner_group.local_env_runner + if worker.preprocessors.get(policy_id) is not None: + preprocessed = worker.preprocessors[policy_id].transform(ob) + else: + preprocessed = ob + filtered = worker.filters[policy_id](preprocessed, update=False) + filtered_obs.append(filtered) + if state is None: + continue + elif agent_id in state: + filtered_state.append(state[agent_id]) + else: + filtered_state.append(policy.get_initial_state()) + + obs_batch = np.stack(filtered_obs) + if state is None: + state = [] + else: + state = list(zip(*filtered_state)) + state = [np.stack(s) for s in state] + + input_dict = {Columns.OBS: obs_batch} + + if prev_action is not None: + input_dict[SampleBatch.PREV_ACTIONS] = prev_action + if prev_reward is not None: + input_dict[SampleBatch.PREV_REWARDS] = prev_reward + if info: + input_dict[Columns.INFOS] = info + for i, s in enumerate(state): + input_dict[f"state_in_{i}"] = s + + actions, states, infos = policy.compute_actions_from_input_dict( + input_dict=input_dict, + explore=explore, + timestep=timestep, + episodes=episodes, + ) + + single_actions = space_utils.unbatch(actions) + actions = {} + for key, a in zip(observations, single_actions): + if unsquash_actions: + a = space_utils.unsquash_action(a, policy.action_space_struct) + elif clip_actions: + a = space_utils.clip_action(a, policy.action_space_struct) + actions[key] = a + + unbatched_states = {} + for idx, agent_id in enumerate(observations): + unbatched_states[agent_id] = [s[idx] for s in states] + + if state_defined or full_fetch: + return actions, unbatched_states, infos + else: + return actions + + @Deprecated(new="Algorithm.restore_env_runners", error=True) + def restore_workers(self, *args, **kwargs): + pass + + @Deprecated( + new="Algorithm.env_runner_group", + error=True, + ) + @property + def workers(self): + return self.env_runner_group + + @Deprecated( + new="Algorithm.eval_env_runner_group", + error=True, + ) + @property + def evaluation_workers(self): + return self.eval_env_runner_group + + +class TrainIterCtx: + def __init__(self, algo: Algorithm): + self.algo = algo + self.time_start = None + self.time_stop = None + + def __enter__(self): + # Before first call to `step()`, `results` is expected to be None -> + # Start with self.failures=-1 -> set to 0 before the very first call + # to `self.step()`. + self.failures = -1 + + self.time_start = time.time() + self.sampled = 0 + self.trained = 0 + if self.algo.config.enable_env_runner_and_connector_v2: + self.init_env_steps_sampled = self.algo.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME), default=0 + ) + self.init_env_steps_trained = self.algo.metrics.peek( + (LEARNER_RESULTS, ALL_MODULES, NUM_ENV_STEPS_TRAINED_LIFETIME), + default=0, + ) + self.init_agent_steps_sampled = sum( + self.algo.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_AGENT_STEPS_SAMPLED_LIFETIME), default={} + ).values() + ) + self.init_agent_steps_trained = sum( + self.algo.metrics.peek( + (LEARNER_RESULTS, NUM_AGENT_STEPS_TRAINED_LIFETIME), default={} + ).values() + ) + else: + self.init_env_steps_sampled = self.algo._counters[NUM_ENV_STEPS_SAMPLED] + self.init_env_steps_trained = self.algo._counters[NUM_ENV_STEPS_TRAINED] + self.init_agent_steps_sampled = self.algo._counters[NUM_AGENT_STEPS_SAMPLED] + self.init_agent_steps_trained = self.algo._counters[NUM_AGENT_STEPS_TRAINED] + self.failure_tolerance = ( + self.algo.config.num_consecutive_env_runner_failures_tolerance + ) + return self + + def __exit__(self, *args): + self.time_stop = time.time() + + def get_time_taken_sec(self) -> float: + """Returns the time we spent in the context in seconds.""" + return self.time_stop - self.time_start + + def should_stop(self, results): + # Before first call to `step()`. + if results in [None, False]: + # Fail after n retries. + self.failures += 1 + if self.failures > self.failure_tolerance: + raise RuntimeError( + "More than `num_consecutive_env_runner_failures_tolerance=" + f"{self.failure_tolerance}` consecutive worker failures! " + "Exiting." + ) + # Continue to very first `step()` call or retry `step()` after + # a (tolerable) failure. + return False + + # Stopping criteria. + if self.algo.config.enable_env_runner_and_connector_v2: + if self.algo.config.count_steps_by == "agent_steps": + self.sampled = ( + sum( + self.algo.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_AGENT_STEPS_SAMPLED_LIFETIME), + default={}, + ).values() + ) + - self.init_agent_steps_sampled + ) + self.trained = ( + sum( + self.algo.metrics.peek( + (LEARNER_RESULTS, NUM_AGENT_STEPS_TRAINED_LIFETIME), + default={}, + ).values() + ) + - self.init_agent_steps_trained + ) + else: + self.sampled = ( + self.algo.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME), default=0 + ) + - self.init_env_steps_sampled + ) + self.trained = ( + self.algo.metrics.peek( + (LEARNER_RESULTS, ALL_MODULES, NUM_ENV_STEPS_TRAINED_LIFETIME), + default=0, + ) + - self.init_env_steps_trained + ) + else: + if self.algo.config.count_steps_by == "agent_steps": + self.sampled = ( + self.algo._counters[NUM_AGENT_STEPS_SAMPLED] + - self.init_agent_steps_sampled + ) + self.trained = ( + self.algo._counters[NUM_AGENT_STEPS_TRAINED] + - self.init_agent_steps_trained + ) + else: + self.sampled = ( + self.algo._counters[NUM_ENV_STEPS_SAMPLED] + - self.init_env_steps_sampled + ) + self.trained = ( + self.algo._counters[NUM_ENV_STEPS_TRAINED] + - self.init_env_steps_trained + ) + + min_t = self.algo.config.min_time_s_per_iteration + min_sample_ts = self.algo.config.min_sample_timesteps_per_iteration + min_train_ts = self.algo.config.min_train_timesteps_per_iteration + + # Repeat if not enough time has passed or if not enough + # env|train timesteps have been processed (or these min + # values are not provided by the user). + if ( + (not min_t or time.time() - self.time_start >= min_t) + and (not min_sample_ts or self.sampled >= min_sample_ts) + and (not min_train_ts or self.trained >= min_train_ts) + ): + return True + else: + return False diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm_config.py b/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ecf3c16b6404d6e239d9bda58b89f8ad754f29d6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm_config.py @@ -0,0 +1,6325 @@ +import copy +import dataclasses +import logging +import math +import sys +from enum import Enum +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + Dict, + List, + Optional, + Tuple, + Type, + Union, +) + +import gymnasium as gym +import tree +from packaging import version +from typing_extensions import Self + +import ray +from ray._common.deprecation import ( + DEPRECATED_VALUE, + Deprecated, + deprecation_warning, +) +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.differentiable_learner_config import ( + DifferentiableLearnerConfig, +) +from ray.rllib.core.rl_module import validate_module_id +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env import INPUT_ENV_SINGLE_SPACES, INPUT_ENV_SPACES +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.wrappers.atari_wrappers import is_atari +from ray.rllib.evaluation.collectors.sample_collector import SampleCollector +from ray.rllib.evaluation.collectors.simple_list_collector import SimpleListCollector +from ray.rllib.models import MODEL_DEFAULTS +from ray.rllib.offline.input_reader import InputReader +from ray.rllib.offline.io_context import IOContext +from ray.rllib.policy.policy import Policy, PolicySpec +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID +from ray.rllib.utils import deep_update, force_list, merge_dicts +from ray.rllib.utils.annotations import ( + OldAPIStack, + OverrideToImplementCustomLogic_CallToSuperRecommended, +) +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.from_config import NotProvided, from_config +from ray.rllib.utils.schedules.scheduler import Scheduler +from ray.rllib.utils.serialization import ( + NOT_SERIALIZABLE, + deserialize_type, + serialize_type, +) +from ray.rllib.utils.test_utils import check +from ray.rllib.utils.torch_utils import TORCH_COMPILE_REQUIRED_VERSION +from ray.rllib.utils.typing import ( + AgentID, + AlgorithmConfigDict, + EnvConfigDict, + EnvType, + LearningRateOrSchedule, + ModuleID, + MultiAgentPolicyConfigDict, + PartialAlgorithmConfigDict, + PolicyID, + RLModuleSpecType, + SampleBatchType, +) +from ray.tune.logger import Logger +from ray.tune.registry import get_trainable_cls +from ray.tune.result import TRIAL_INFO +from ray.tune.tune import _Config +from ray.util import log_once +from ray.util.placement_group import PlacementGroup + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm import Algorithm + from ray.rllib.core.learner import Learner + from ray.rllib.core.learner.differentiable_learner import DifferentiableLearner + from ray.rllib.core.learner.learner_group import LearnerGroup + from ray.rllib.core.learner.torch.torch_meta_learner import TorchMetaLearner + from ray.rllib.core.rl_module.rl_module import RLModule + from ray.rllib.utils.typing import EpisodeType + +logger = logging.getLogger(__name__) + + +def _check_rl_module_spec(module_spec: RLModuleSpecType) -> None: + if not isinstance(module_spec, (RLModuleSpec, MultiRLModuleSpec)): + raise ValueError( + "rl_module_spec must be an instance of " + "RLModuleSpec or MultiRLModuleSpec." + f"Got {type(module_spec)} instead." + ) + + +class AlgorithmConfig(_Config): + """A RLlib AlgorithmConfig builds an RLlib Algorithm from a given configuration. + + .. testcode:: + + from ray.rllib.algorithms.ppo import PPOConfig + from ray.rllib.algorithms.callbacks import MemoryTrackingCallbacks + # Construct a generic config object, specifying values within different + # sub-categories, e.g. "training". + config = ( + PPOConfig() + .training(gamma=0.9, lr=0.01) + .environment(env="CartPole-v1") + .env_runners(num_env_runners=0) + .callbacks(MemoryTrackingCallbacks) + ) + # A config object can be used to construct the respective Algorithm. + rllib_algo = config.build() + + .. testcode:: + + from ray.rllib.algorithms.ppo import PPOConfig + from ray import tune + # In combination with a tune.grid_search: + config = PPOConfig() + config.training(lr=tune.grid_search([0.01, 0.001])) + # Use `to_dict()` method to get the legacy plain python config dict + # for usage with `tune.Tuner().fit()`. + tune.Tuner("PPO", param_space=config.to_dict()) + """ + + @staticmethod + def DEFAULT_AGENT_TO_MODULE_MAPPING_FN(agent_id, episode): + # The default agent ID to module ID mapping function to use in the multi-agent + # case if None is provided. + # Map any agent ID to "default_policy". + return DEFAULT_MODULE_ID + + # @OldAPIStack + # TODO (sven): Deprecate in new API stack. + @staticmethod + def DEFAULT_POLICY_MAPPING_FN(aid, episode, worker, **kwargs): + # The default policy mapping function to use if None provided. + # Map any agent ID to "default_policy". + return DEFAULT_POLICY_ID + + @classmethod + def from_dict(cls, config_dict: dict) -> Self: + """Creates an AlgorithmConfig from a legacy python config dict. + + .. testcode:: + + from ray.rllib.algorithms.ppo.ppo import PPOConfig + # pass a RLlib config dict + ppo_config = PPOConfig.from_dict({}) + ppo = ppo_config.build(env="Pendulum-v1") + + Args: + config_dict: The legacy formatted python config dict for some algorithm. + + Returns: + A new AlgorithmConfig object that matches the given python config dict. + """ + # Create a default config object of this class. + config_obj = cls() + # Remove `_is_frozen` flag from config dict in case the AlgorithmConfig that + # the dict was derived from was already frozen (we don't want to copy the + # frozenness). + config_dict.pop("_is_frozen", None) + config_obj.update_from_dict(config_dict) + return config_obj + + @classmethod + def overrides(cls, **kwargs): + """Generates and validates a set of config key/value pairs (passed via kwargs). + + Validation whether given config keys are valid is done immediately upon + construction (by comparing against the properties of a default AlgorithmConfig + object of this class). + Allows combination with a full AlgorithmConfig object to yield a new + AlgorithmConfig object. + + Used anywhere, we would like to enable the user to only define a few config + settings that would change with respect to some main config, e.g. in multi-agent + setups and evaluation configs. + + .. testcode:: + + from ray.rllib.algorithms.ppo import PPOConfig + from ray.rllib.policy.policy import PolicySpec + config = ( + PPOConfig() + .multi_agent( + policies={ + "pol0": PolicySpec(config=PPOConfig.overrides(lambda_=0.95)) + }, + ) + ) + + + .. testcode:: + + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.algorithms.ppo import PPOConfig + config = ( + PPOConfig() + .evaluation( + evaluation_num_env_runners=1, + evaluation_interval=1, + evaluation_config=AlgorithmConfig.overrides(explore=False), + ) + ) + + Returns: + A dict mapping valid config property-names to values. + + Raises: + KeyError: In case a non-existing property name (kwargs key) is being + passed in. Valid property names are taken from a default + AlgorithmConfig object of `cls`. + """ + default_config = cls() + config_overrides = {} + for key, value in kwargs.items(): + if not hasattr(default_config, key): + raise KeyError( + f"Invalid property name {key} for config class {cls.__name__}!" + ) + # Allow things like "lambda" as well. + key = cls._translate_special_keys(key, warn_deprecated=True) + config_overrides[key] = value + + return config_overrides + + def __init__(self, algo_class: Optional[type] = None): + """Initializes an AlgorithmConfig instance. + + Args: + algo_class: An optional Algorithm class that this config class belongs to. + Used (if provided) to build a respective Algorithm instance from this + config. + """ + # Define all settings and their default values. + + # Define the default RLlib Algorithm class that this AlgorithmConfig is applied + # to. + self.algo_class = algo_class + + # `self.python_environment()` + self.extra_python_environs_for_driver = {} + self.extra_python_environs_for_worker = {} + + # `self.resources()` + self.placement_strategy = "PACK" + self.num_gpus = 0 # @OldAPIStack + self._fake_gpus = False # @OldAPIStack + self.num_cpus_for_main_process = 1 + + # `self.framework()` + self.framework_str = "torch" + self.eager_tracing = True + self.eager_max_retraces = 20 + self.tf_session_args = { + # note: overridden by `local_tf_session_args` + "intra_op_parallelism_threads": 2, + "inter_op_parallelism_threads": 2, + "gpu_options": { + "allow_growth": True, + }, + "log_device_placement": False, + "device_count": {"CPU": 1}, + # Required by multi-GPU (num_gpus > 1). + "allow_soft_placement": True, + } + self.local_tf_session_args = { + # Allow a higher level of parallelism by default, but not unlimited + # since that can cause crashes with many concurrent drivers. + "intra_op_parallelism_threads": 8, + "inter_op_parallelism_threads": 8, + } + # Torch compile settings + self.torch_compile_learner = False + self.torch_compile_learner_what_to_compile = ( + TorchCompileWhatToCompile.FORWARD_TRAIN + ) + # AOT Eager is a dummy backend and doesn't result in speedups. + self.torch_compile_learner_dynamo_backend = ( + "aot_eager" if sys.platform == "darwin" else "inductor" + ) + self.torch_compile_learner_dynamo_mode = None + self.torch_compile_worker = False + # AOT Eager is a dummy backend and doesn't result in speedups. + self.torch_compile_worker_dynamo_backend = ( + "aot_eager" if sys.platform == "darwin" else "onnxrt" + ) + self.torch_compile_worker_dynamo_mode = None + # Default kwargs for `torch.nn.parallel.DistributedDataParallel`. + self.torch_ddp_kwargs = {} + # Default setting for skipping `nan` gradient updates. + self.torch_skip_nan_gradients = False + + # `self.environment()` + self.env = None + self.env_config = {} + self.observation_space = None + self.action_space = None + self.clip_rewards = None + self.normalize_actions = True + self.clip_actions = False + self._is_atari = None + self.disable_env_checking = False + # Deprecated settings: + self.render_env = False + self.action_mask_key = "action_mask" + + # `self.env_runners()` + self.env_runner_cls = None + self.num_env_runners = 0 + self.create_local_env_runner = True + self.num_envs_per_env_runner = 1 + # TODO (sven): Once new ormsgpack system in place, replace the string + # with proper `gym.envs.registration.VectorizeMode.SYNC`. + self.gym_env_vectorize_mode = "sync" + self.num_cpus_per_env_runner = 1 + self.num_gpus_per_env_runner = 0 + self.custom_resources_per_env_runner = {} + self.validate_env_runners_after_construction = True + self.episodes_to_numpy = True + self.max_requests_in_flight_per_env_runner = 1 + self.sample_timeout_s = 60.0 + self.create_env_on_local_worker = False + self._env_to_module_connector = None + self.add_default_connectors_to_env_to_module_pipeline = True + self._module_to_env_connector = None + self.add_default_connectors_to_module_to_env_pipeline = True + self.merge_env_runner_states = "training_only" + self.broadcast_env_runner_states = True + self.episode_lookback_horizon = 1 + # TODO (sven): Rename into `sample_timesteps` (or `sample_duration` + # and `sample_duration_unit` (replacing batch_mode), like we do it + # in the evaluation config). + self.rollout_fragment_length = 200 + # TODO (sven): Rename into `sample_mode`. + self.batch_mode = "truncate_episodes" + self.compress_observations = False + # @OldAPIStack + self.remote_worker_envs = False + self.remote_env_batch_wait_ms = 0 + self.enable_tf1_exec_eagerly = False + self.sample_collector = SimpleListCollector + self.preprocessor_pref = "deepmind" + self.observation_filter = "NoFilter" + self.update_worker_filter_stats = True + self.use_worker_filter_stats = True + self.sampler_perf_stats_ema_coef = None + self._is_online = True + + # `self.learners()` + self.num_learners = 0 + self.num_gpus_per_learner = 0 + self.num_cpus_per_learner = "auto" + self.num_aggregator_actors_per_learner = 0 + self.max_requests_in_flight_per_aggregator_actor = 3 + self.local_gpu_idx = 0 + # TODO (sven): This probably works even without any restriction + # (allowing for any arbitrary number of requests in-flight). Test with + # 3 first, then with unlimited, and if both show the same behavior on + # an async algo, remove this restriction entirely. + self.max_requests_in_flight_per_learner = 3 + + # `self.training()` + self.gamma = 0.99 + self.lr = 0.001 + self.grad_clip = None + self.grad_clip_by = "global_norm" + # Simple logic for now: If None, use `train_batch_size`. + self._train_batch_size_per_learner = None + self.train_batch_size = 32 # @OldAPIStack + + # These setting have been adopted from the original PPO batch settings: + # num_sgd_iter, minibatch_size, and shuffle_sequences. + self.num_epochs = 1 + self.minibatch_size = None + self.shuffle_batch_per_epoch = False + + # TODO (sven): Unsolved problem with RLModules sometimes requiring settings from + # the main AlgorithmConfig. We should not require the user to provide those + # settings in both, the AlgorithmConfig (as property) AND the model config + # dict. We should generally move to a world, in which there exists an + # AlgorithmConfig that a) has-a user provided model config object and b) + # is given a chance to compile a final model config (dict or object) that is + # then passed into the RLModule/Catalog. This design would then match our + # "compilation" pattern, where we compile automatically those settings that + # should NOT be touched by the user. + # In case, an Algorithm already uses the above described pattern (and has + # `self.model` as a @property, ignore AttributeError (for trying to set this + # property). + try: + self.model = copy.deepcopy(MODEL_DEFAULTS) + except AttributeError: + pass + + self._learner_connector = None + self.add_default_connectors_to_learner_pipeline = True + self.learner_config_dict = {} + self.optimizer = {} # @OldAPIStack + self._learner_class = None + + # `self.callbacks()` + # TODO (sven): Set this default to None, once the old API stack has been + # deprecated. + self.callbacks_class = RLlibCallback + self.callbacks_on_algorithm_init = None + self.callbacks_on_env_runners_recreated = None + self.callbacks_on_offline_eval_runners_recreated = None + self.callbacks_on_checkpoint_loaded = None + self.callbacks_on_environment_created = None + self.callbacks_on_episode_created = None + self.callbacks_on_episode_start = None + self.callbacks_on_episode_step = None + self.callbacks_on_episode_end = None + self.callbacks_on_evaluate_start = None + self.callbacks_on_evaluate_end = None + self.callbacks_on_evaluate_offline_start = None + self.callbacks_on_evaluate_offline_end = None + self.callbacks_on_sample_end = None + self.callbacks_on_train_result = None + + # `self.explore()` + self.explore = True + # This is not compatible with RLModules, which have a method + # `forward_exploration` to specify custom exploration behavior. + if not hasattr(self, "exploration_config"): + # Helper to keep track of the original exploration config when dis-/enabling + # rl modules. + self._prior_exploration_config = None + self.exploration_config = {} + + # `self.api_stack()` + self.enable_rl_module_and_learner = True + self.enable_env_runner_and_connector_v2 = True + self.api_stack( + enable_rl_module_and_learner=True, + enable_env_runner_and_connector_v2=True, + ) + + # `self.multi_agent()` + # TODO (sven): Prepare multi-agent setup for logging each agent's and each + # RLModule's steps taken thus far (and passing this information into the + # EnvRunner metrics and the RLModule's forward pass). Thereby, deprecate the + # `count_steps_by` config setting AND - at the same time - allow users to + # specify the batch size unit instead (agent- vs env steps). + self.count_steps_by = "env_steps" + # self.agent_to_module_mapping_fn = self.DEFAULT_AGENT_TO_MODULE_MAPPING_FN + # Soon to be Deprecated. + self.policies = {DEFAULT_POLICY_ID: PolicySpec()} + self.policy_map_capacity = 100 + self.policy_mapping_fn = self.DEFAULT_POLICY_MAPPING_FN + self.policies_to_train = None + self.policy_states_are_swappable = False + self.observation_fn = None + + # `self.offline_data()` + self.input_ = "sampler" + self.offline_data_class = None + self.offline_data_class = None + self.input_read_method = "read_parquet" + self.input_read_method_kwargs = {} + self.input_read_schema = {} + self.input_read_episodes = False + self.input_read_sample_batches = False + self.input_read_batch_size = None + self.input_filesystem = None + self.input_filesystem_kwargs = {} + self.input_compress_columns = [Columns.OBS, Columns.NEXT_OBS] + self.input_spaces_jsonable = True + self.materialize_data = False + self.materialize_mapped_data = True + self.map_batches_kwargs = {} + self.iter_batches_kwargs = {} + # Use always the final observation until the user explicitly ask + # to ignore it. + self.ignore_final_observation = False + self.prelearner_class = None + self.prelearner_buffer_class = None + self.prelearner_buffer_kwargs = {} + self.prelearner_module_synch_period = 10 + self.dataset_num_iters_per_learner = None + self.input_config = {} + self.actions_in_input_normalized = False + self.postprocess_inputs = False + self.shuffle_buffer_size = 0 + self.output = None + self.output_config = {} + self.output_compress_columns = [Columns.OBS, Columns.NEXT_OBS] + self.output_max_file_size = 64 * 1024 * 1024 + self.output_max_rows_per_file = None + self.output_write_remaining_data = False + self.output_write_method = "write_parquet" + self.output_write_method_kwargs = {} + self.output_filesystem = None + self.output_filesystem_kwargs = {} + self.output_write_episodes = True + self.offline_sampling = False + + # `self.evaluation()` + self.evaluation_interval = None + self.evaluation_duration = 10 + self.evaluation_duration_unit = "episodes" + self.evaluation_sample_timeout_s = 120.0 + self.evaluation_auto_duration_min_env_steps_per_sample = 100 + self.evaluation_auto_duration_max_env_steps_per_sample = 2000 + self.evaluation_parallel_to_training = False + self.evaluation_force_reset_envs_before_iteration = True + self.evaluation_config = None + self.off_policy_estimation_methods = {} + self.ope_split_batch_by_episode = True + self.evaluation_num_env_runners = 0 + self.custom_evaluation_function = None + # TODO: Set this flag still in the config or - much better - in the + # RolloutWorker as a property. + self.in_evaluation = False + # TODO (sven): Deprecate this setting (it's not user-accessible right now any + # way). Replace by logic within `training_step` to merge and broadcast the + # EnvRunner (connector) states. + self.sync_filters_on_rollout_workers_timeout_s = 10.0 + # Offline evaluation. + self.offline_evaluation_interval = None + self.num_offline_eval_runners = 0 + self.offline_evaluation_type: str = None + self.offline_eval_runner_class = None + # TODO (simon): Only `_offline_evaluate_with_fixed_duration` works. Also, + # decide, if we use `offline_evaluation_duration` or + # `dataset_num_iters_per_offline_eval_runner`. Should the user decide here? + # The latter will be much faster, but runs per runner call all evaluation. + self.offline_loss_for_module_fn = None + self.offline_evaluation_duration = 1 + self.offline_evaluation_parallel_to_training = False + self.offline_evaluation_timeout_s = 120.0 + self.num_cpus_per_offline_eval_runner = 1 + self.num_gpus_per_offline_eval_runner = 0 + self.custom_resources_per_offline_eval_runner = {} + self.restart_failed_offline_eval_runners = True + self.ignore_offline_eval_runner_failures = False + self.max_num_offline_eval_runner_restarts = 1000 + self.offline_eval_runner_restore_timeout_s = 1800.0 + self.max_requests_in_flight_per_offline_eval_runner = 1 + self.validate_offline_eval_runners_after_construction = True + self.offline_eval_runner_health_probe_timeout_s = 30.0 + self.offline_eval_rl_module_inference_only = False + self.broadcast_offline_eval_runner_states = False + self.offline_eval_batch_size_per_runner = 256 + self.dataset_num_iters_per_eval_runner = 1 + + # `self.reporting()` + self.keep_per_episode_custom_metrics = False + self.metrics_episode_collection_timeout_s = 60.0 + self.metrics_num_episodes_for_smoothing = 100 + self.min_time_s_per_iteration = None + self.min_train_timesteps_per_iteration = 0 + self.min_sample_timesteps_per_iteration = 0 + self.log_gradients = False + + # `self.checkpointing()` + self.export_native_model_files = False + self.checkpoint_trainable_policies_only = False + + # `self.debugging()` + self.logger_creator = None + self.logger_config = None + self.log_level = "WARN" + self.log_sys_usage = True + self.fake_sampler = False + self.seed = None + + # `self.fault_tolerance()` + self.restart_failed_env_runners = True + self.ignore_env_runner_failures = False + # By default, restart failed worker a thousand times. + # This should be enough to handle normal transient failures. + # This also prevents infinite number of restarts in case the worker or env has + # a bug. + self.max_num_env_runner_restarts = 1000 + # Small delay between worker restarts. In case EnvRunners or eval EnvRunners + # have remote dependencies, this delay can be adjusted to make sure we don't + # flood them with re-connection requests, and allow them enough time to recover. + # This delay also gives Ray time to stream back error logging and exceptions. + self.delay_between_env_runner_restarts_s = 60.0 + self.restart_failed_sub_environments = False + self.num_consecutive_env_runner_failures_tolerance = 100 + self.env_runner_health_probe_timeout_s = 30.0 + self.env_runner_restore_timeout_s = 1800.0 + + # `self.rl_module()` + self._model_config = {} + self._rl_module_spec = None + # Module ID specific config overrides. + self.algorithm_config_overrides_per_module = {} + # Cached, actual AlgorithmConfig objects derived from + # `self.algorithm_config_overrides_per_module`. + self._per_module_overrides: Dict[ModuleID, "AlgorithmConfig"] = {} + + # `self.experimental()` + self._validate_config = True + self._use_msgpack_checkpoints = False + self._torch_grad_scaler_class = None + self._torch_lr_scheduler_classes = None + self._tf_policy_handles_more_than_one_loss = False + self._disable_preprocessor_api = False + self._disable_action_flattening = False + self._disable_initialize_loss_from_dummy_batch = False + self._dont_auto_sync_env_runner_states = False + + # Has this config object been frozen (cannot alter its attributes anymore). + self._is_frozen = False + + # TODO: Remove, once all deprecation_warning calls upon using these keys + # have been removed. + # === Deprecated keys === + self.env_task_fn = DEPRECATED_VALUE + self.enable_connectors = DEPRECATED_VALUE + self.simple_optimizer = DEPRECATED_VALUE + self.monitor = DEPRECATED_VALUE + self.evaluation_num_episodes = DEPRECATED_VALUE + self.metrics_smoothing_episodes = DEPRECATED_VALUE + self.timesteps_per_iteration = DEPRECATED_VALUE + self.min_iter_time_s = DEPRECATED_VALUE + self.collect_metrics_timeout = DEPRECATED_VALUE + self.min_time_s_per_reporting = DEPRECATED_VALUE + self.min_train_timesteps_per_reporting = DEPRECATED_VALUE + self.min_sample_timesteps_per_reporting = DEPRECATED_VALUE + self.input_evaluation = DEPRECATED_VALUE + self.policy_map_cache = DEPRECATED_VALUE + self.worker_cls = DEPRECATED_VALUE + self.synchronize_filters = DEPRECATED_VALUE + self.enable_async_evaluation = DEPRECATED_VALUE + self.custom_async_evaluation_function = DEPRECATED_VALUE + self._enable_rl_module_api = DEPRECATED_VALUE + self.auto_wrap_old_gym_envs = DEPRECATED_VALUE + self.always_attach_evaluation_results = DEPRECATED_VALUE + + # The following values have moved because of the new ReplayBuffer API + self.buffer_size = DEPRECATED_VALUE + self.prioritized_replay = DEPRECATED_VALUE + self.learning_starts = DEPRECATED_VALUE + self.replay_batch_size = DEPRECATED_VALUE + # -1 = DEPRECATED_VALUE is a valid value for replay_sequence_length + self.replay_sequence_length = None + self.replay_mode = DEPRECATED_VALUE + self.prioritized_replay_alpha = DEPRECATED_VALUE + self.prioritized_replay_beta = DEPRECATED_VALUE + self.prioritized_replay_eps = DEPRECATED_VALUE + self.min_time_s_per_reporting = DEPRECATED_VALUE + self.min_train_timesteps_per_reporting = DEPRECATED_VALUE + self.min_sample_timesteps_per_reporting = DEPRECATED_VALUE + self._disable_execution_plan_api = DEPRECATED_VALUE + + def to_dict(self) -> AlgorithmConfigDict: + """Converts all settings into a legacy config dict for backward compatibility. + + Returns: + A complete AlgorithmConfigDict, usable in backward-compatible Tune/RLlib + use cases. + """ + config = copy.deepcopy(vars(self)) + config.pop("algo_class") + config.pop("_is_frozen") + + # Worst naming convention ever: NEVER EVER use reserved key-words... + if "lambda_" in config: + assert hasattr(self, "lambda_") + config["lambda"] = self.lambda_ + config.pop("lambda_") + if "input_" in config: + assert hasattr(self, "input_") + config["input"] = self.input_ + config.pop("input_") + + # Convert `policies` (PolicySpecs?) into dict. + # Convert policies dict such that each policy ID maps to a old-style. + # 4-tuple: class, obs-, and action space, config. + if "policies" in config and isinstance(config["policies"], dict): + policies_dict = {} + for policy_id, policy_spec in config.pop("policies").items(): + if isinstance(policy_spec, PolicySpec): + policies_dict[policy_id] = policy_spec.get_state() + else: + policies_dict[policy_id] = policy_spec + config["policies"] = policies_dict + + # Switch out deprecated vs new config keys. + config["callbacks"] = config.pop("callbacks_class", None) + config["create_env_on_driver"] = config.pop("create_env_on_local_worker", 1) + config["custom_eval_function"] = config.pop("custom_evaluation_function", None) + config["framework"] = config.pop("framework_str", None) + + # Simplify: Remove all deprecated keys that have as value `DEPRECATED_VALUE`. + # These would be useless in the returned dict anyways. + for dep_k in [ + "monitor", + "evaluation_num_episodes", + "metrics_smoothing_episodes", + "timesteps_per_iteration", + "min_iter_time_s", + "collect_metrics_timeout", + "buffer_size", + "prioritized_replay", + "learning_starts", + "replay_batch_size", + "replay_mode", + "prioritized_replay_alpha", + "prioritized_replay_beta", + "prioritized_replay_eps", + "min_time_s_per_reporting", + "min_train_timesteps_per_reporting", + "min_sample_timesteps_per_reporting", + "input_evaluation", + "_enable_new_api_stack", + ]: + if config.get(dep_k) == DEPRECATED_VALUE: + config.pop(dep_k, None) + + return config + + def update_from_dict( + self, + config_dict: PartialAlgorithmConfigDict, + ) -> Self: + """Modifies this AlgorithmConfig via the provided python config dict. + + Warns if `config_dict` contains deprecated keys. + Silently sets even properties of `self` that do NOT exist. This way, this method + may be used to configure custom Policies which do not have their own specific + AlgorithmConfig classes, e.g. + `ray.rllib.examples.policy.random_policy::RandomPolicy`. + + Args: + config_dict: The old-style python config dict (PartialAlgorithmConfigDict) + to use for overriding some properties defined in there. + + Returns: + This updated AlgorithmConfig object. + """ + eval_call = {} + + # We deal with this special key before all others because it may influence + # stuff like "exploration_config". + # Namely, we want to re-instantiate the exploration config this config had + # inside `self.experimental()` before potentially overwriting it in the + # following. + enable_new_api_stack = config_dict.get( + "enable_rl_module_and_learner", + config_dict.get("enable_env_runner_and_connector_v2"), + ) + if enable_new_api_stack is not None: + self.api_stack( + enable_rl_module_and_learner=enable_new_api_stack, + enable_env_runner_and_connector_v2=enable_new_api_stack, + ) + + # Modify our properties one by one. + for key, value in config_dict.items(): + key = self._translate_special_keys(key, warn_deprecated=False) + + # Ray Tune saves additional data under this magic keyword. + # This should not get treated as AlgorithmConfig field. + if key == TRIAL_INFO: + continue + + if key in ["_enable_new_api_stack"]: + # We've dealt with this above. + continue + # Set our multi-agent settings. + elif key == "multiagent": + kwargs = { + k: value[k] + for k in [ + "policies", + "policy_map_capacity", + "policy_mapping_fn", + "policies_to_train", + "policy_states_are_swappable", + "observation_fn", + "count_steps_by", + ] + if k in value + } + self.multi_agent(**kwargs) + # Some keys specify config sub-dicts and therefore should go through the + # correct methods to properly `.update()` those from given config dict + # (to not lose any sub-keys). + elif key == "callbacks_class" and value != NOT_SERIALIZABLE: + # For backward compatibility reasons, only resolve possible + # classpath if value is a str type. + if isinstance(value, str): + value = deserialize_type(value, error=True) + self.callbacks(callbacks_class=value) + elif key == "env_config": + self.environment(env_config=value) + elif key.startswith("evaluation_"): + eval_call[key] = value + elif key == "exploration_config": + if enable_new_api_stack: + self.exploration_config = value + continue + if isinstance(value, dict) and "type" in value: + value["type"] = deserialize_type(value["type"]) + self.env_runners(exploration_config=value) + elif key == "model": + # Resolve possible classpath. + if isinstance(value, dict) and value.get("custom_model"): + value["custom_model"] = deserialize_type(value["custom_model"]) + self.training(**{key: value}) + elif key == "optimizer": + self.training(**{key: value}) + elif key == "replay_buffer_config": + if isinstance(value, dict) and "type" in value: + value["type"] = deserialize_type(value["type"]) + self.training(**{key: value}) + elif key == "sample_collector": + # Resolve possible classpath. + value = deserialize_type(value) + self.env_runners(sample_collector=value) + # Set the property named `key` to `value`. + else: + setattr(self, key, value) + + self.evaluation(**eval_call) + + return self + + def get_state(self) -> Dict[str, Any]: + """Returns a dict state that can be pickled. + + Returns: + A dictionary containing all attributes of the instance. + """ + + state = self.__dict__.copy() + state["class"] = type(self) + state.pop("algo_class") + state.pop("_is_frozen") + state = {k: v for k, v in state.items() if v != DEPRECATED_VALUE} + + # Convert `policies` (PolicySpecs?) into dict. + # Convert policies dict such that each policy ID maps to a old-style. + # 4-tuple: class, obs-, and action space, config. + # TODO (simon, sven): Remove when deprecating old stack. + if "policies" in state and isinstance(state["policies"], dict): + policies_dict = {} + for policy_id, policy_spec in state.pop("policies").items(): + if isinstance(policy_spec, PolicySpec): + policies_dict[policy_id] = policy_spec.get_state() + else: + policies_dict[policy_id] = policy_spec + state["policies"] = policies_dict + + # state = self._serialize_dict(state) + + return state + + @classmethod + def from_state(cls, state: Dict[str, Any]) -> Union[Self, Any]: + """Returns an instance constructed from the state. + + Args: + state: A dictionary containing the state of an `AlgorithmConfig`. + See `AlgorithmConfig.get_state` for creating a state. + The constructed class will be of ``state["class"]``. + + Returns: + An `AlgorithmConfig` instance with attributes from the `state`. + """ + + # As ctor could be any other class add Any to the return type to indicate this. + ctor = state["class"] + config = ctor() + + config.__dict__.update(state) + + return config + + # TODO(sven): We might want to have a `deserialize` method as well. Right now, + # simply using the from_dict() API works in this same (deserializing) manner, + # whether the dict used is actually code-free (already serialized) or not + # (i.e. a classic RLlib config dict with e.g. "callbacks" key still pointing to + # a class). + def serialize(self) -> Dict[str, Any]: + """Returns a mapping from str to JSON'able values representing this config. + + The resulting values don't have any code in them. + Classes (such as `callbacks_class`) are converted to their full + classpath, e.g. `ray.rllib.callbacks.callbacks.RLlibCallback`. + Actual code such as lambda functions ware written as their source + code (str) plus any closure information for properly restoring the + code inside the AlgorithmConfig object made from the returned dict data. + Dataclass objects get converted to dicts. + + Returns: + A dict mapping from str to JSON'able values. + """ + config = self.to_dict() + return self._serialize_dict(config) + + def copy(self, copy_frozen: Optional[bool] = None) -> Self: + """Creates a deep copy of this config and (un)freezes if necessary. + + Args: + copy_frozen: Whether the created deep copy is frozen or not. If None, + keep the same frozen status that `self` currently has. + + Returns: + A deep copy of `self` that is (un)frozen. + """ + cp = copy.deepcopy(self) + if copy_frozen is True: + cp.freeze() + elif copy_frozen is False: + cp._is_frozen = False + if isinstance(cp.evaluation_config, AlgorithmConfig): + cp.evaluation_config._is_frozen = False + return cp + + def freeze(self) -> None: + """Freezes this config object, such that no attributes can be set anymore. + + Algorithms should use this method to make sure that their config objects + remain read-only after this. + """ + if self._is_frozen: + return + self._is_frozen = True + + # Also freeze underlying eval config, if applicable. + if isinstance(self.evaluation_config, AlgorithmConfig): + self.evaluation_config.freeze() + + # TODO: Flip out all set/dict/list values into frozen versions + # of themselves? This way, users won't even be able to alter those values + # directly anymore. + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def validate(self) -> None: + """Validates all values in this config.""" + + # Validation is blocked. + if not self._validate_config: + return + + self._validate_env_runner_settings() + self._validate_callbacks_settings() + self._validate_framework_settings() + self._validate_resources_settings() + self._validate_multi_agent_settings() + self._validate_input_settings() + self._validate_evaluation_settings() + self._validate_offline_settings() + self._validate_new_api_stack_settings() + self._validate_to_be_deprecated_settings() + + def build_algo( + self, + env: Optional[Union[str, EnvType]] = None, + logger_creator: Optional[Callable[[], Logger]] = None, + use_copy: bool = True, + ) -> "Algorithm": + """Builds an Algorithm from this AlgorithmConfig (or a copy thereof). + + Args: + env: Name of the environment to use (e.g. a gym-registered str), + a full class path (e.g. + "ray.rllib.examples.envs.classes.random_env.RandomEnv"), or an Env + class directly. Note that this arg can also be specified via + the "env" key in `config`. + logger_creator: Callable that creates a ray.tune.Logger + object. If unspecified, a default logger is created. + use_copy: Whether to deepcopy `self` and pass the copy to the Algorithm + (instead of `self`) as config. This is useful in case you would like to + recycle the same AlgorithmConfig over and over, e.g. in a test case, in + which we loop over different DL-frameworks. + + Returns: + A ray.rllib.algorithms.algorithm.Algorithm object. + """ + if env is not None: + self.env = env + if self.evaluation_config is not None: + self.evaluation_config["env"] = env + if logger_creator is not None: + self.logger_creator = logger_creator + + algo_class = self.algo_class + if isinstance(self.algo_class, str): + algo_class = get_trainable_cls(self.algo_class) + + return algo_class( + config=self if not use_copy else copy.deepcopy(self), + logger_creator=self.logger_creator, + ) + + def build_env_to_module_connector( + self, + env=None, + spaces=None, + device=None, + ) -> ConnectorV2: + from ray.rllib.connectors.env_to_module import ( + AddObservationsFromEpisodesToBatch, + AddStatesFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AgentToModuleMapping, + BatchIndividualItems, + EnvToModulePipeline, + NumpyToTensor, + ) + + custom_connectors = [] + # Create an env-to-module connector pipeline (including RLlib's default + # env->module connector piece) and return it. + if self._env_to_module_connector is not None: + try: + val_ = self._env_to_module_connector(env, spaces, device) + # Try deprecated signature, if necessary. + except TypeError as e: + if "positional argument" in e.args[0]: + if log_once("env-to-module-wrong-signature"): + logger.error( + "Your `config.env_to_module_connector` function seems to " + "have a wrong or outdated signature! It should be: " + "`def myfunc(env, spaces, device): ...`, where any of " + "these arguments are optional and may be None.\n" + "`env` is the (vectorized) gym env.\n" + "`spaces` is a dict of structure `{'__env__': ([" + "vectorized env obs. space, vectorized env act. space])," + "'__env_single__': ([env obs. space, env act. space])}`.\n" + "`device` is a (torch) device.\n" + ) + val_ = self._env_to_module_connector(env) + else: + raise e + + # ConnectorV2 (piece or pipeline). + if isinstance(val_, ConnectorV2): + custom_connectors = [val_] + # Sequence of individual ConnectorV2 pieces. + elif isinstance(val_, (list, tuple)): + custom_connectors = list(val_) + # Unsupported return value. + else: + raise ValueError( + "`AlgorithmConfig.env_runners(env_to_module_connector=..)` must " + "return a ConnectorV2 object or a list thereof to be added to a " + f"connector pipeline! Your function returned {val_}." + ) + + if env is not None: + obs_space = getattr(env, "single_observation_space", env.observation_space) + elif spaces is not None and INPUT_ENV_SINGLE_SPACES in spaces: + obs_space = spaces[INPUT_ENV_SINGLE_SPACES][0] + else: + obs_space = self.observation_space + if obs_space is None and self.is_multi_agent: + obs_space = gym.spaces.Dict( + { + aid: env.envs[0].unwrapped.get_observation_space(aid) + for aid in env.envs[0].unwrapped.possible_agents + } + ) + if env is not None: + act_space = getattr(env, "single_action_space", env.action_space) + elif spaces is not None and INPUT_ENV_SINGLE_SPACES in spaces: + act_space = spaces[INPUT_ENV_SINGLE_SPACES][1] + else: + act_space = self.action_space + if act_space is None and self.is_multi_agent: + act_space = gym.spaces.Dict( + { + aid: env.envs[0].unwrapped.get_action_space(aid) + for aid in env.envs[0].unwrapped.possible_agents + } + ) + pipeline = EnvToModulePipeline( + input_observation_space=obs_space, + input_action_space=act_space, + connectors=custom_connectors, + ) + + if self.add_default_connectors_to_env_to_module_pipeline: + # Append OBS handling. + pipeline.append(AddObservationsFromEpisodesToBatch()) + # Append time-rank handler. + pipeline.append(AddTimeDimToBatchAndZeroPad()) + # Append STATE_IN/STATE_OUT handler. + pipeline.append(AddStatesFromEpisodesToBatch()) + # If multi-agent -> Map from AgentID-based data to ModuleID based data. + if self.is_multi_agent: + pipeline.append( + AgentToModuleMapping( + rl_module_specs=( + self.rl_module_spec.rl_module_specs + if isinstance(self.rl_module_spec, MultiRLModuleSpec) + else set(self.policies) + ), + agent_to_module_mapping_fn=self.policy_mapping_fn, + ) + ) + # Batch all data. + pipeline.append(BatchIndividualItems(multi_agent=self.is_multi_agent)) + # Convert to Tensors. + pipeline.append(NumpyToTensor(device=device)) + + return pipeline + + def build_module_to_env_connector(self, env=None, spaces=None) -> ConnectorV2: + from ray.rllib.connectors.module_to_env import ( + GetActions, + ListifyDataForVectorEnv, + ModuleToAgentUnmapping, + ModuleToEnvPipeline, + NormalizeAndClipActions, + RemoveSingleTsTimeRankFromBatch, + TensorToNumpy, + UnBatchToIndividualItems, + ) + + custom_connectors = [] + # Create a module-to-env connector pipeline (including RLlib's default + # module->env connector piece) and return it. + if self._module_to_env_connector is not None: + try: + val_ = self._module_to_env_connector(env, spaces) + # Try deprecated signature, if necessary. + except TypeError as e: + if "positional argument" in e.args[0]: + if log_once("module-to-env-wrong-signature"): + logger.error( + "Your `config.module_to_env_connector` function seems to " + "have a wrong or outdated signature! It should be: " + "`def myfunc(env, spaces): ...`, where any of " + "these arguments are optional and may be None.\n" + "`env` is the (vectorized) gym env.\n" + "`spaces` is a dict of structure `{'__env__': ([" + "vectorized env obs. space, vectorized env act. space])," + "'__env_single__': ([env obs. space, env act. space])}`.\n" + ) + val_ = self._module_to_env_connector(env) + + # ConnectorV2 (piece or pipeline). + if isinstance(val_, ConnectorV2): + custom_connectors = [val_] + # Sequence of individual ConnectorV2 pieces. + elif isinstance(val_, (list, tuple)): + custom_connectors = list(val_) + # Unsupported return value. + else: + raise ValueError( + "`AlgorithmConfig.env_runners(module_to_env_connector=..)` must " + "return a ConnectorV2 object or a list thereof to be added to a " + f"connector pipeline! Your function returned {val_}." + ) + + if env is not None: + obs_space = getattr(env, "single_observation_space", env.observation_space) + elif spaces is not None and INPUT_ENV_SINGLE_SPACES in spaces: + obs_space = spaces[INPUT_ENV_SINGLE_SPACES][0] + else: + obs_space = self.observation_space + if obs_space is None and self.is_multi_agent: + obs_space = gym.spaces.Dict( + { + aid: env.envs[0].unwrapped.get_observation_space(aid) + for aid in env.envs[0].unwrapped.possible_agents + } + ) + if env is not None: + act_space = getattr(env, "single_action_space", env.action_space) + elif spaces is not None and INPUT_ENV_SINGLE_SPACES in spaces: + act_space = spaces[INPUT_ENV_SINGLE_SPACES][1] + else: + act_space = self.action_space + if act_space is None and self.is_multi_agent: + act_space = gym.spaces.Dict( + { + aid: env.envs[0].unwrapped.get_action_space(aid) + for aid in env.envs[0].unwrapped.possible_agents + } + ) + pipeline = ModuleToEnvPipeline( + input_observation_space=obs_space, + input_action_space=act_space, + connectors=custom_connectors, + ) + + if self.add_default_connectors_to_module_to_env_pipeline: + # Prepend: Anything that has to do with plain data processing (not + # particularly with the actions). + + # Remove extra time-rank, if applicable. + pipeline.prepend(RemoveSingleTsTimeRankFromBatch()) + + # If multi-agent -> Map from ModuleID-based data to AgentID based data. + if self.is_multi_agent: + pipeline.prepend(ModuleToAgentUnmapping()) + + # Unbatch all data. + pipeline.prepend(UnBatchToIndividualItems()) + + # Convert to numpy. + pipeline.prepend(TensorToNumpy()) + + # Sample actions from ACTION_DIST_INPUTS (if ACTIONS not present). + pipeline.prepend(GetActions()) + + # Append: Anything that has to do with action sampling. + # Unsquash/clip actions based on config and action space. + pipeline.append( + NormalizeAndClipActions( + normalize_actions=self.normalize_actions, + clip_actions=self.clip_actions, + ) + ) + # Listify data from ConnectorV2-data format to normal lists that we can + # index into by env vector index. These lists contain individual items + # for single-agent and multi-agent dicts for multi-agent. + pipeline.append(ListifyDataForVectorEnv()) + + return pipeline + + def build_learner_connector( + self, + input_observation_space, + input_action_space, + device=None, + ) -> ConnectorV2: + from ray.rllib.connectors.learner import ( + AddColumnsFromEpisodesToTrainBatch, + AddObservationsFromEpisodesToBatch, + AddStatesFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AgentToModuleMapping, + BatchIndividualItems, + LearnerConnectorPipeline, + NumpyToTensor, + ) + + custom_connectors = [] + # Create a learner connector pipeline (including RLlib's default + # learner connector piece) and return it. + if self._learner_connector is not None: + val_ = self._learner_connector( + input_observation_space, + input_action_space, + # device, # TODO (sven): Also pass device into custom builder. + ) + + # ConnectorV2 (piece or pipeline). + if isinstance(val_, ConnectorV2): + custom_connectors = [val_] + # Sequence of individual ConnectorV2 pieces. + elif isinstance(val_, (list, tuple)): + custom_connectors = list(val_) + # Unsupported return value. + else: + raise ValueError( + "`AlgorithmConfig.learners(learner_connector=..)` must return " + "a ConnectorV2 object or a list thereof to be added to a connector " + f"pipeline! Your function returned {val_}." + ) + + pipeline = LearnerConnectorPipeline( + connectors=custom_connectors, + input_observation_space=input_observation_space, + input_action_space=input_action_space, + ) + if self.add_default_connectors_to_learner_pipeline: + # Append OBS handling. + pipeline.append( + AddObservationsFromEpisodesToBatch(as_learner_connector=True) + ) + # Append all other columns handling. + pipeline.append(AddColumnsFromEpisodesToTrainBatch()) + # Append time-rank handler. + pipeline.append(AddTimeDimToBatchAndZeroPad(as_learner_connector=True)) + # Append STATE_IN/STATE_OUT handler. + pipeline.append(AddStatesFromEpisodesToBatch(as_learner_connector=True)) + # If multi-agent -> Map from AgentID-based data to ModuleID based data. + if self.is_multi_agent: + pipeline.append( + AgentToModuleMapping( + rl_module_specs=( + self.rl_module_spec.rl_module_specs + if isinstance(self.rl_module_spec, MultiRLModuleSpec) + else set(self.policies) + ), + agent_to_module_mapping_fn=self.policy_mapping_fn, + ) + ) + # Batch all data. + pipeline.append(BatchIndividualItems(multi_agent=self.is_multi_agent)) + # Convert to Tensors. + pipeline.append(NumpyToTensor(as_learner_connector=True, device=device)) + return pipeline + + def build_learner_group( + self, + *, + env: Optional[EnvType] = None, + spaces: Optional[Dict[ModuleID, Tuple[gym.Space, gym.Space]]] = None, + rl_module_spec: Optional[RLModuleSpecType] = None, + placement_group: Optional["PlacementGroup"] = None, + ) -> "LearnerGroup": + """Builds and returns a new LearnerGroup object based on settings in `self`. + + Args: + env: An optional EnvType object (e.g. a gym.Env) useful for extracting space + information for the to-be-constructed RLModule inside the LearnerGroup's + Learner workers. Note that if RLlib cannot infer any space information + either from this `env` arg, from the optional `spaces` arg or from + `self`, the LearnerGroup cannot be created. + spaces: An optional dict mapping ModuleIDs to + (observation-space, action-space)-tuples for the to-be-constructed + RLModule inside the LearnerGroup's Learner workers. Note that if RLlib + cannot infer any space information either from this `spces` arg, + from the optional `env` arg or from `self`, the LearnerGroup cannot + be created. + rl_module_spec: An optional (single-agent or multi-agent) RLModuleSpec to + use for the constructed LearnerGroup. If None, RLlib tries to infer + the RLModuleSpec using the other information given and stored in this + `AlgorithmConfig` object. + + Returns: + The newly created `LearnerGroup` object. + """ + from ray.rllib.core.learner.learner_group import LearnerGroup + + # If `spaces` or `env` provided -> Create a MultiRLModuleSpec first to be + # passed into the LearnerGroup constructor. + if rl_module_spec is None: + rl_module_spec = self.get_multi_rl_module_spec(env=env, spaces=spaces) + + # Construct the actual LearnerGroup. + learner_group = LearnerGroup( + config=self.copy(), + module_spec=rl_module_spec, + placement_group=placement_group, + ) + + return learner_group + + def build_learner( + self, + *, + env: Optional[EnvType] = None, + spaces: Optional[Dict[PolicyID, Tuple[gym.Space, gym.Space]]] = None, + ) -> "Learner": + """Builds and returns a new Learner object based on settings in `self`. + + This Learner object already has its `build()` method called, meaning + its RLModule is already constructed. + + Args: + env: An optional EnvType object (e.g. a gym.Env) useful for extracting space + information for the to-be-constructed RLModule inside the Learner. + Note that if RLlib cannot infer any space information + either from this `env` arg, from the optional `spaces` arg or from + `self`, the Learner cannot be created. + spaces: An optional dict mapping ModuleIDs to + (observation-space, action-space)-tuples for the to-be-constructed + RLModule inside the Learner. Note that if RLlib cannot infer any + space information either from this `spces` arg, from the optional + `env` arg or from `self`, the Learner cannot be created. + + Returns: + The newly created (and already built) Learner object. + """ + # If `spaces` or `env` provided -> Create a MultiRLModuleSpec first to be + # passed into the LearnerGroup constructor. + rl_module_spec = None + if env is not None or spaces is not None: + rl_module_spec = self.get_multi_rl_module_spec(env=env, spaces=spaces) + # Construct the actual Learner object. + learner = self.learner_class(config=self, module_spec=rl_module_spec) + # `build()` the Learner (internal structures such as RLModule, etc..). + learner.build() + + return learner + + def get_config_for_module(self, module_id: ModuleID) -> Self: + """Returns an AlgorithmConfig object, specific to the given module ID. + + In a multi-agent setup, individual modules might override one or more + AlgorithmConfig properties (e.g. `train_batch_size`, `lr`) using the + `overrides()` method. + + In order to retrieve a full AlgorithmConfig instance (with all these overrides + already translated and built-in), users can call this method with the respective + module ID. + + Args: + module_id: The module ID for which to get the final AlgorithmConfig object. + + Returns: + A new AlgorithmConfig object for the specific module ID. + """ + # ModuleID NOT found in cached ModuleID, but in overrides dict. + # Create new algo config object and cache it. + if ( + module_id not in self._per_module_overrides + and module_id in self.algorithm_config_overrides_per_module + ): + self._per_module_overrides[module_id] = self.copy().update_from_dict( + self.algorithm_config_overrides_per_module[module_id] + ) + + # Return the module specific algo config object. + if module_id in self._per_module_overrides: + return self._per_module_overrides[module_id] + # No overrides for ModuleID -> return self. + else: + return self + + def python_environment( + self, + *, + extra_python_environs_for_driver: Optional[dict] = NotProvided, + extra_python_environs_for_worker: Optional[dict] = NotProvided, + ) -> Self: + """Sets the config's python environment settings. + + Args: + extra_python_environs_for_driver: Any extra python env vars to set in the + algorithm's process, e.g., {"OMP_NUM_THREADS": "16"}. + extra_python_environs_for_worker: The extra python environments need to set + for worker processes. + + Returns: + This updated AlgorithmConfig object. + """ + if extra_python_environs_for_driver is not NotProvided: + self.extra_python_environs_for_driver = extra_python_environs_for_driver + if extra_python_environs_for_worker is not NotProvided: + self.extra_python_environs_for_worker = extra_python_environs_for_worker + return self + + def resources( + self, + *, + num_cpus_for_main_process: Optional[int] = NotProvided, + num_gpus: Optional[Union[float, int]] = NotProvided, # @OldAPIStack + _fake_gpus: Optional[bool] = NotProvided, # @OldAPIStack + placement_strategy: Optional[str] = NotProvided, + # Deprecated args. + num_cpus_per_worker=DEPRECATED_VALUE, # moved to `env_runners` + num_gpus_per_worker=DEPRECATED_VALUE, # moved to `env_runners` + custom_resources_per_worker=DEPRECATED_VALUE, # moved to `env_runners` + num_learner_workers=DEPRECATED_VALUE, # moved to `learners` + num_cpus_per_learner_worker=DEPRECATED_VALUE, # moved to `learners` + num_gpus_per_learner_worker=DEPRECATED_VALUE, # moved to `learners` + local_gpu_idx=DEPRECATED_VALUE, # moved to `learners` + num_cpus_for_local_worker=DEPRECATED_VALUE, + ) -> Self: + """Specifies resources allocated for an Algorithm and its ray actors/workers. + + Args: + num_cpus_for_main_process: Number of CPUs to allocate for the main algorithm + process that runs `Algorithm.training_step()`. + Note: This is only relevant when running RLlib through Tune. Otherwise, + `Algorithm.training_step()` runs in the main program (driver). + num_gpus: Number of GPUs to allocate to the algorithm process. + Note that not all algorithms can take advantage of GPUs. + Support for multi-GPU is currently only available for + tf-[PPO/IMPALA/DQN/PG]. This can be fractional (e.g., 0.3 GPUs). + _fake_gpus: Set to True for debugging (multi-)?GPU funcitonality on a + CPU machine. GPU towers are simulated by graphs located on + CPUs in this case. Use `num_gpus` to test for different numbers of + fake GPUs. + placement_strategy: The strategy for the placement group factory returned by + `Algorithm.default_resource_request()`. A PlacementGroup defines, which + devices (resources) should always be co-located on the same node. + For example, an Algorithm with 2 EnvRunners and 1 Learner (with + 1 GPU) requests a placement group with the bundles: + [{"cpu": 1}, {"gpu": 1, "cpu": 1}, {"cpu": 1}, {"cpu": 1}], where the + first bundle is for the local (main Algorithm) process, the second one + for the 1 Learner worker and the last 2 bundles are for the two + EnvRunners. These bundles can now be "placed" on the same or different + nodes depending on the value of `placement_strategy`: + "PACK": Packs bundles into as few nodes as possible. + "SPREAD": Places bundles across distinct nodes as even as possible. + "STRICT_PACK": Packs bundles into one node. The group is not allowed + to span multiple nodes. + "STRICT_SPREAD": Packs bundles across distinct nodes. + + Returns: + This updated AlgorithmConfig object. + """ + if num_cpus_per_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(num_cpus_per_worker)", + new="AlgorithmConfig.env_runners(num_cpus_per_env_runner)", + error=False, + ) + self.num_cpus_per_env_runner = num_cpus_per_worker + + if num_gpus_per_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(num_gpus_per_worker)", + new="AlgorithmConfig.env_runners(num_gpus_per_env_runner)", + error=False, + ) + self.num_gpus_per_env_runner = num_gpus_per_worker + + if custom_resources_per_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(custom_resources_per_worker)", + new="AlgorithmConfig.env_runners(custom_resources_per_env_runner)", + error=False, + ) + self.custom_resources_per_env_runner = custom_resources_per_worker + + if num_learner_workers != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(num_learner_workers)", + new="AlgorithmConfig.learners(num_learner)", + error=False, + ) + self.num_learners = num_learner_workers + + if num_cpus_per_learner_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(num_cpus_per_learner_worker)", + new="AlgorithmConfig.learners(num_cpus_per_learner)", + error=False, + ) + self.num_cpus_per_learner = num_cpus_per_learner_worker + + if num_gpus_per_learner_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(num_gpus_per_learner_worker)", + new="AlgorithmConfig.learners(num_gpus_per_learner)", + error=False, + ) + self.num_gpus_per_learner = num_gpus_per_learner_worker + + if local_gpu_idx != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(local_gpu_idx)", + new="AlgorithmConfig.learners(local_gpu_idx)", + error=False, + ) + self.local_gpu_idx = local_gpu_idx + + if num_cpus_for_local_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.resources(num_cpus_for_local_worker)", + new="AlgorithmConfig.resources(num_cpus_for_main_process)", + error=False, + ) + self.num_cpus_for_main_process = num_cpus_for_local_worker + + if num_cpus_for_main_process is not NotProvided: + self.num_cpus_for_main_process = num_cpus_for_main_process + if num_gpus is not NotProvided: + self.num_gpus = num_gpus + if _fake_gpus is not NotProvided: + self._fake_gpus = _fake_gpus + if placement_strategy is not NotProvided: + self.placement_strategy = placement_strategy + + return self + + def framework( + self, + framework: Optional[str] = NotProvided, + *, + eager_tracing: Optional[bool] = NotProvided, + eager_max_retraces: Optional[int] = NotProvided, + tf_session_args: Optional[Dict[str, Any]] = NotProvided, + local_tf_session_args: Optional[Dict[str, Any]] = NotProvided, + torch_compile_learner: Optional[bool] = NotProvided, + torch_compile_learner_what_to_compile: Optional[str] = NotProvided, + torch_compile_learner_dynamo_mode: Optional[str] = NotProvided, + torch_compile_learner_dynamo_backend: Optional[str] = NotProvided, + torch_compile_worker: Optional[bool] = NotProvided, + torch_compile_worker_dynamo_backend: Optional[str] = NotProvided, + torch_compile_worker_dynamo_mode: Optional[str] = NotProvided, + torch_ddp_kwargs: Optional[Dict[str, Any]] = NotProvided, + torch_skip_nan_gradients: Optional[bool] = NotProvided, + ) -> Self: + """Sets the config's DL framework settings. + + Args: + framework: torch: PyTorch; tf2: TensorFlow 2.x (eager execution or traced + if eager_tracing=True); tf: TensorFlow (static-graph); + eager_tracing: Enable tracing in eager mode. This greatly improves + performance (speedup ~2x), but makes it slightly harder to debug + since Python code won't be evaluated after the initial eager pass. + Only possible if framework=tf2. + eager_max_retraces: Maximum number of tf.function re-traces before a + runtime error is raised. This is to prevent unnoticed retraces of + methods inside the `..._eager_traced` Policy, which could slow down + execution by a factor of 4, without the user noticing what the root + cause for this slowdown could be. + Only necessary for framework=tf2. + Set to None to ignore the re-trace count and never throw an error. + tf_session_args: Configures TF for single-process operation by default. + local_tf_session_args: Override the following tf session args on the local + worker + torch_compile_learner: If True, forward_train methods on TorchRLModules + on the learner are compiled. If not specified, the default is to compile + forward train on the learner. + torch_compile_learner_what_to_compile: A TorchCompileWhatToCompile + mode specifying what to compile on the learner side if + torch_compile_learner is True. See TorchCompileWhatToCompile for + details and advice on its usage. + torch_compile_learner_dynamo_backend: The torch dynamo backend to use on + the learner. + torch_compile_learner_dynamo_mode: The torch dynamo mode to use on the + learner. + torch_compile_worker: If True, forward exploration and inference methods on + TorchRLModules on the workers are compiled. If not specified, + the default is to not compile forward methods on the workers because + retracing can be expensive. + torch_compile_worker_dynamo_backend: The torch dynamo backend to use on + the workers. + torch_compile_worker_dynamo_mode: The torch dynamo mode to use on the + workers. + torch_ddp_kwargs: The kwargs to pass into + `torch.nn.parallel.DistributedDataParallel` when using `num_learners + > 1`. This is specifically helpful when searching for unused parameters + that are not used in the backward pass. This can give hints for errors + in custom models where some parameters do not get touched in the + backward pass although they should. + torch_skip_nan_gradients: If updates with `nan` gradients should be entirely + skipped. This skips updates in the optimizer entirely if they contain + any `nan` gradient. This can help to avoid biasing moving-average based + optimizers - like Adam. This can help in training phases where policy + updates can be highly unstable such as during the early stages of + training or with highly exploratory policies. In such phases many + gradients might turn `nan` and setting them to zero could corrupt the + optimizer's internal state. The default is `False` and turns `nan` + gradients to zero. If many `nan` gradients are encountered consider (a) + monitoring gradients by setting `log_gradients` in `AlgorithmConfig` to + `True`, (b) use proper weight initialization (e.g. Xavier, Kaiming) via + the `model_config_dict` in `AlgorithmConfig.rl_module` and/or (c) + gradient clipping via `grad_clip` in `AlgorithmConfig.training`. + + Returns: + This updated AlgorithmConfig object. + """ + if framework is not NotProvided: + if framework == "tfe": + deprecation_warning( + old="AlgorithmConfig.framework('tfe')", + new="AlgorithmConfig.framework('tf2')", + error=True, + ) + self.framework_str = framework + if eager_tracing is not NotProvided: + self.eager_tracing = eager_tracing + if eager_max_retraces is not NotProvided: + self.eager_max_retraces = eager_max_retraces + if tf_session_args is not NotProvided: + self.tf_session_args = tf_session_args + if local_tf_session_args is not NotProvided: + self.local_tf_session_args = local_tf_session_args + + if torch_compile_learner is not NotProvided: + self.torch_compile_learner = torch_compile_learner + if torch_compile_learner_dynamo_backend is not NotProvided: + self.torch_compile_learner_dynamo_backend = ( + torch_compile_learner_dynamo_backend + ) + if torch_compile_learner_dynamo_mode is not NotProvided: + self.torch_compile_learner_dynamo_mode = torch_compile_learner_dynamo_mode + if torch_compile_learner_what_to_compile is not NotProvided: + self.torch_compile_learner_what_to_compile = ( + torch_compile_learner_what_to_compile + ) + if torch_compile_worker is not NotProvided: + self.torch_compile_worker = torch_compile_worker + if torch_compile_worker_dynamo_backend is not NotProvided: + self.torch_compile_worker_dynamo_backend = ( + torch_compile_worker_dynamo_backend + ) + if torch_compile_worker_dynamo_mode is not NotProvided: + self.torch_compile_worker_dynamo_mode = torch_compile_worker_dynamo_mode + if torch_ddp_kwargs is not NotProvided: + self.torch_ddp_kwargs = torch_ddp_kwargs + if torch_skip_nan_gradients is not NotProvided: + self.torch_skip_nan_gradients = torch_skip_nan_gradients + + return self + + def api_stack( + self, + enable_rl_module_and_learner: Optional[bool] = NotProvided, + enable_env_runner_and_connector_v2: Optional[bool] = NotProvided, + ) -> Self: + """Sets the config's API stack settings. + + Args: + enable_rl_module_and_learner: Enables the usage of `RLModule` (instead of + `ModelV2`) and Learner (instead of the training-related parts of + `Policy`). Must be used with `enable_env_runner_and_connector_v2=True`. + Together, these two settings activate the "new API stack" of RLlib. + enable_env_runner_and_connector_v2: Enables the usage of EnvRunners + (SingleAgentEnvRunner and MultiAgentEnvRunner) and ConnectorV2. + When setting this to True, `enable_rl_module_and_learner` must be True + as well. Together, these two settings activate the "new API stack" of + RLlib. + + Returns: + This updated AlgorithmConfig object. + """ + if enable_rl_module_and_learner is not NotProvided: + self.enable_rl_module_and_learner = enable_rl_module_and_learner + + if enable_rl_module_and_learner is True and self.exploration_config: + self._prior_exploration_config = self.exploration_config + self.exploration_config = {} + + elif enable_rl_module_and_learner is False and not self.exploration_config: + if self._prior_exploration_config is not None: + self.exploration_config = self._prior_exploration_config + self._prior_exploration_config = None + else: + logger.warning( + "config.enable_rl_module_and_learner was set to False, but no " + "prior exploration config was found to be restored." + ) + + if enable_env_runner_and_connector_v2 is not NotProvided: + self.enable_env_runner_and_connector_v2 = enable_env_runner_and_connector_v2 + + return self + + def environment( + self, + env: Optional[Union[str, EnvType]] = NotProvided, + *, + env_config: Optional[EnvConfigDict] = NotProvided, + observation_space: Optional[gym.Space] = NotProvided, + action_space: Optional[gym.Space] = NotProvided, + render_env: Optional[bool] = NotProvided, + clip_rewards: Optional[Union[bool, float]] = NotProvided, + normalize_actions: Optional[bool] = NotProvided, + clip_actions: Optional[bool] = NotProvided, + disable_env_checking: Optional[bool] = NotProvided, + is_atari: Optional[bool] = NotProvided, + action_mask_key: Optional[str] = NotProvided, + # Deprecated args. + env_task_fn=DEPRECATED_VALUE, + ) -> Self: + """Sets the config's RL-environment settings. + + Args: + env: The environment specifier. This can either be a tune-registered env, + via `tune.register_env([name], lambda env_ctx: [env object])`, + or a string specifier of an RLlib supported type. In the latter case, + RLlib tries to interpret the specifier as either an Farama-Foundation + gymnasium env, a PyBullet env, or a fully qualified classpath to an Env + class, e.g. "ray.rllib.examples.envs.classes.random_env.RandomEnv". + env_config: Arguments dict passed to the env creator as an EnvContext + object (which is a dict plus the properties: `num_env_runners`, + `worker_index`, `vector_index`, and `remote`). + observation_space: The observation space for the Policies of this Algorithm. + action_space: The action space for the Policies of this Algorithm. + render_env: If True, try to render the environment on the local worker or on + worker 1 (if num_env_runners > 0). For vectorized envs, this usually + means that only the first sub-environment is rendered. + In order for this to work, your env has to implement the + `render()` method which either: + a) handles window generation and rendering itself (returning True) or + b) returns a numpy uint8 image of shape [height x width x 3 (RGB)]. + clip_rewards: Whether to clip rewards during Policy's postprocessing. + None (default): Clip for Atari only (r=sign(r)). + True: r=sign(r): Fixed rewards -1.0, 1.0, or 0.0. + False: Never clip. + [float value]: Clip at -value and + value. + Tuple[value1, value2]: Clip at value1 and value2. + normalize_actions: If True, RLlib learns entirely inside a normalized + action space (0.0 centered with small stddev; only affecting Box + components). RLlib unsquashes actions (and clip, just in case) to the + bounds of the env's action space before sending actions back to the env. + clip_actions: If True, the RLlib default ModuleToEnv connector clips + actions according to the env's bounds (before sending them into the + `env.step()` call). + disable_env_checking: Disable RLlib's env checks after a gymnasium.Env + instance has been constructed in an EnvRunner. Note that the checks + include an `env.reset()` and `env.step()` (with a random action), which + might tinker with your env's logic and behavior and thus negatively + influence sample collection- and/or learning behavior. + is_atari: This config can be used to explicitly specify whether the env is + an Atari env or not. If not specified, RLlib tries to auto-detect + this. + action_mask_key: If observation is a dictionary, expect the value by + the key `action_mask_key` to contain a valid actions mask (`numpy.int8` + array of zeros and ones). Defaults to "action_mask". + + Returns: + This updated AlgorithmConfig object. + """ + if env_task_fn != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.environment(env_task_fn=..)", + error=True, + ) + if env is not NotProvided: + self.env = env + if env_config is not NotProvided: + deep_update(self.env_config, env_config, True) + if observation_space is not NotProvided: + self.observation_space = observation_space + if action_space is not NotProvided: + self.action_space = action_space + if render_env is not NotProvided: + self.render_env = render_env + if clip_rewards is not NotProvided: + self.clip_rewards = clip_rewards + if normalize_actions is not NotProvided: + self.normalize_actions = normalize_actions + if clip_actions is not NotProvided: + self.clip_actions = clip_actions + if disable_env_checking is not NotProvided: + self.disable_env_checking = disable_env_checking + if is_atari is not NotProvided: + self._is_atari = is_atari + if action_mask_key is not NotProvided: + self.action_mask_key = action_mask_key + + return self + + def env_runners( + self, + *, + env_runner_cls: Optional[type] = NotProvided, + num_env_runners: Optional[int] = NotProvided, + create_local_env_runner: Optional[bool] = NotProvided, + create_env_on_local_worker: Optional[bool] = NotProvided, + num_envs_per_env_runner: Optional[int] = NotProvided, + gym_env_vectorize_mode: Optional[Union[str, gym.VectorizeMode]] = NotProvided, + num_cpus_per_env_runner: Optional[int] = NotProvided, + num_gpus_per_env_runner: Optional[Union[float, int]] = NotProvided, + custom_resources_per_env_runner: Optional[dict] = NotProvided, + validate_env_runners_after_construction: Optional[bool] = NotProvided, + sample_timeout_s: Optional[float] = NotProvided, + max_requests_in_flight_per_env_runner: Optional[int] = NotProvided, + env_to_module_connector: Optional[ + Callable[[EnvType], Union["ConnectorV2", List["ConnectorV2"]]] + ] = NotProvided, + module_to_env_connector: Optional[ + Callable[[EnvType, "RLModule"], Union["ConnectorV2", List["ConnectorV2"]]] + ] = NotProvided, + add_default_connectors_to_env_to_module_pipeline: Optional[bool] = NotProvided, + add_default_connectors_to_module_to_env_pipeline: Optional[bool] = NotProvided, + episode_lookback_horizon: Optional[int] = NotProvided, + merge_env_runner_states: Optional[Union[str, bool]] = NotProvided, + broadcast_env_runner_states: Optional[bool] = NotProvided, + compress_observations: Optional[bool] = NotProvided, + rollout_fragment_length: Optional[Union[int, str]] = NotProvided, + batch_mode: Optional[str] = NotProvided, + explore: Optional[bool] = NotProvided, + episodes_to_numpy: Optional[bool] = NotProvided, + # @OldAPIStack settings. + use_worker_filter_stats: Optional[bool] = NotProvided, + update_worker_filter_stats: Optional[bool] = NotProvided, + exploration_config: Optional[dict] = NotProvided, # @OldAPIStack + sample_collector: Optional[Type[SampleCollector]] = NotProvided, # @OldAPIStack + remote_worker_envs: Optional[bool] = NotProvided, # @OldAPIStack + remote_env_batch_wait_ms: Optional[float] = NotProvided, # @OldAPIStack + preprocessor_pref: Optional[str] = NotProvided, # @OldAPIStack + observation_filter: Optional[str] = NotProvided, # @OldAPIStack + enable_tf1_exec_eagerly: Optional[bool] = NotProvided, # @OldAPIStack + sampler_perf_stats_ema_coef: Optional[float] = NotProvided, # @OldAPIStack + # Deprecated args. + num_rollout_workers=DEPRECATED_VALUE, + num_envs_per_worker=DEPRECATED_VALUE, + validate_workers_after_construction=DEPRECATED_VALUE, + ignore_worker_failures=DEPRECATED_VALUE, + recreate_failed_workers=DEPRECATED_VALUE, + restart_failed_sub_environments=DEPRECATED_VALUE, + num_consecutive_worker_failures_tolerance=DEPRECATED_VALUE, + worker_health_probe_timeout_s=DEPRECATED_VALUE, + worker_restore_timeout_s=DEPRECATED_VALUE, + synchronize_filter=DEPRECATED_VALUE, + enable_connectors=DEPRECATED_VALUE, + ) -> Self: + """Sets the rollout worker configuration. + + Args: + env_runner_cls: The EnvRunner class to use for environment rollouts (data + collection). + num_env_runners: Number of EnvRunner actors to create for parallel sampling. + Setting this to 0 forces sampling to be done in the local + EnvRunner (main process or the Algorithm's actor when using Tune). + num_envs_per_env_runner: Number of environments to step through + (vector-wise) per EnvRunner. This enables batching when computing + actions through RLModule inference, which can improve performance + for inference-bottlenecked workloads. + gym_env_vectorize_mode: The gymnasium vectorization mode for vector envs. + Must be a `gymnasium.VectorizeMode` (enum) value. + Default is SYNC. Set this to ASYNC to parallelize the individual sub + environments within the vector. This can speed up your EnvRunners + significantly when using heavier environments. Set this to + VECTOR_ENTRY_POINT in case your env creator, also known as + "gym entry point", already returns a gym.vector.VectorEnv and you + don't need RLlib to vectorize the environments for the runners. + num_cpus_per_env_runner: Number of CPUs to allocate per EnvRunner. + num_gpus_per_env_runner: Number of GPUs to allocate per EnvRunner. This can + be fractional. This is usually needed only if your env itself requires a + GPU (i.e., it is a GPU-intensive video game), or model inference is + unusually expensive. + custom_resources_per_env_runner: Any custom Ray resources to allocate per + EnvRunner. + sample_timeout_s: The timeout in seconds for calling `sample()` on remote + EnvRunner workers. Results (episode list) from workers that take longer + than this time are discarded. Only used by algorithms that sample + synchronously in turn with their update step (e.g., PPO or DQN). Not + relevant for any algos that sample asynchronously, such as APPO or + IMPALA. + max_requests_in_flight_per_env_runner: Max number of in-flight requests + to each EnvRunner (actor)). See the + `ray.rllib.utils.actor_manager.FaultTolerantActorManager` class for more + details. + Tuning these values is important when running experiments with + large sample batches, where there is the risk that the object store may + fill up, causing spilling of objects to disk. This can cause any + asynchronous requests to become very slow, making your experiment run + slowly as well. You can inspect the object store during your experiment + through a call to `ray memory` on your head node, and by using the Ray + dashboard. If you're seeing that the object store is filling up, + turn down the number of remote requests in flight or enable compression + or increase the object store memory through, for example: + `ray.init(object_store_memory=10 * 1024 * 1024 * 1024) # =10 GB` + sample_collector: For the old API stack only. The SampleCollector class to + be used to collect and retrieve environment-, model-, and sampler data. + Override the SampleCollector base class to implement your own + collection/buffering/retrieval logic. + create_local_env_runner: If True, create a local EnvRunner instance, besides + the `num_env_runners` remote EnvRunner actors. If `num_env_runners` is + 0, this setting is ignored and one local EnvRunner is created + regardless. + create_env_on_local_worker: When `num_env_runners` > 0, the driver + (local_worker; worker-idx=0) does not need an environment. This is + because it doesn't have to sample (done by remote_workers; + worker_indices > 0) nor evaluate (done by evaluation workers; + see below). + env_to_module_connector: A callable taking an Env as input arg and returning + an env-to-module ConnectorV2 (might be a pipeline) object. + module_to_env_connector: A callable taking an Env and an RLModule as input + args and returning a module-to-env ConnectorV2 (might be a pipeline) + object. + add_default_connectors_to_env_to_module_pipeline: If True (default), RLlib's + EnvRunners automatically add the default env-to-module ConnectorV2 + pieces to the EnvToModulePipeline. These automatically perform adding + observations and states (in case of stateful Module(s)), agent-to-module + mapping, batching, and conversion to tensor data. Only if you know + exactly what you are doing, you should set this setting to False. + Note that this setting is only relevant if the new API stack is used + (including the new EnvRunner classes). + add_default_connectors_to_module_to_env_pipeline: If True (default), RLlib's + EnvRunners automatically add the default module-to-env ConnectorV2 + pieces to the ModuleToEnvPipeline. These automatically perform removing + the additional time-rank (if applicable, in case of stateful + Module(s)), module-to-agent unmapping, un-batching (to lists), and + conversion from tensor data to numpy. Only if you know exactly what you + are doing, you should set this setting to False. + Note that this setting is only relevant if the new API stack is used + (including the new EnvRunner classes). + episode_lookback_horizon: The amount of data (in timesteps) to keep from the + preceeding episode chunk when a new chunk (for the same episode) is + generated to continue sampling at a later time. The larger this value, + the more an env-to-module connector can look back in time + and compile RLModule input data from this information. For example, if + your custom env-to-module connector (and your custom RLModule) requires + the previous 10 rewards as inputs, you must set this to at least 10. + merge_env_runner_states: True, if remote EnvRunner actor states should be + merged into central connector pipelines. Use "training_only" (default) + for only doing this for the training EnvRunners, NOT for the evaluation + EnvRunners. + broadcast_env_runner_states: True, if merged EnvRunner states (from the + central connector pipelines) should be broadcast back to all remote + EnvRunner actors. + use_worker_filter_stats: Whether to use the workers in the EnvRunnerGroup to + update the central filters (held by the local worker). If False, stats + from the workers aren't used and are discarded. + update_worker_filter_stats: Whether to push filter updates from the central + filters (held by the local worker) to the remote workers' filters. + Setting this to True might be useful within the evaluation config in + order to disable the usage of evaluation trajectories for synching + the central filter (used for training). + rollout_fragment_length: Divide episodes into fragments of this many steps + each during sampling. Trajectories of this size are collected from + EnvRunners and combined into a larger batch of `train_batch_size` + for learning. + For example, given rollout_fragment_length=100 and + train_batch_size=1000: + 1. RLlib collects 10 fragments of 100 steps each from rollout workers. + 2. These fragments are concatenated and we perform an epoch of SGD. + When using multiple envs per worker, the fragment size is multiplied by + `num_envs_per_env_runner`. This is since we are collecting steps from + multiple envs in parallel. For example, if num_envs_per_env_runner=5, + then EnvRunners return experiences in chunks of 5*100 = 500 steps. + The dataflow here can vary per algorithm. For example, PPO further + divides the train batch into minibatches for multi-epoch SGD. + Set `rollout_fragment_length` to "auto" to have RLlib compute an exact + value to match the given batch size. + batch_mode: How to build individual batches with the EnvRunner(s). Batches + coming from distributed EnvRunners are usually concat'd to form the + train batch. Note that "steps" below can mean different things (either + env- or agent-steps) and depends on the `count_steps_by` setting, + adjustable via `AlgorithmConfig.multi_agent(count_steps_by=..)`: + 1) "truncate_episodes": Each call to `EnvRunner.sample()` returns a + batch of at most `rollout_fragment_length * num_envs_per_env_runner` in + size. The batch is exactly `rollout_fragment_length * num_envs` + in size if postprocessing does not change batch sizes. Episodes + may be truncated in order to meet this size requirement. + This mode guarantees evenly sized batches, but increases + variance as the future return must now be estimated at truncation + boundaries. + 2) "complete_episodes": Each call to `EnvRunner.sample()` returns a + batch of at least `rollout_fragment_length * num_envs_per_env_runner` in + size. Episodes aren't truncated, but multiple episodes + may be packed within one batch to meet the (minimum) batch size. + Note that when `num_envs_per_env_runner > 1`, episode steps are + buffered until the episode completes, and hence batches may contain + significant amounts of off-policy data. + explore: Default exploration behavior, iff `explore=None` is passed into + compute_action(s). Set to False for no exploration behavior (e.g., + for evaluation). + episodes_to_numpy: Whether to numpy'ize episodes before + returning them from an EnvRunner. False by default. If True, EnvRunners + call `to_numpy()` on those episode (chunks) to be returned by + `EnvRunners.sample()`. + exploration_config: A dict specifying the Exploration object's config. + remote_worker_envs: If using num_envs_per_env_runner > 1, whether to create + those new envs in remote processes instead of in the same worker. + This adds overheads, but can make sense if your envs can take much + time to step / reset (e.g., for StarCraft). Use this cautiously; + overheads are significant. + remote_env_batch_wait_ms: Timeout that remote workers are waiting when + polling environments. 0 (continue when at least one env is ready) is + a reasonable default, but optimal value could be obtained by measuring + your environment step / reset and model inference perf. + validate_env_runners_after_construction: Whether to validate that each + created remote EnvRunner is healthy after its construction process. + preprocessor_pref: Whether to use "rllib" or "deepmind" preprocessors by + default. Set to None for using no preprocessor. In this case, the + model has to handle possibly complex observations from the + environment. + observation_filter: Element-wise observation filter, either "NoFilter" + or "MeanStdFilter". + compress_observations: Whether to LZ4 compress individual observations + in the SampleBatches collected during rollouts. + enable_tf1_exec_eagerly: Explicitly tells the rollout worker to enable + TF eager execution. This is useful for example when framework is + "torch", but a TF2 policy needs to be restored for evaluation or + league-based purposes. + sampler_perf_stats_ema_coef: If specified, perf stats are in EMAs. This + is the coeff of how much new data points contribute to the averages. + Default is None, which uses simple global average instead. + The EMA update rule is: updated = (1 - ema_coef) * old + ema_coef * new + + Returns: + This updated AlgorithmConfig object. + """ + if enable_connectors != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(enable_connectors=...)", + error=False, + ) + if num_rollout_workers != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(num_rollout_workers)", + new="AlgorithmConfig.env_runners(num_env_runners)", + error=True, + ) + if num_envs_per_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(num_envs_per_worker)", + new="AlgorithmConfig.env_runners(num_envs_per_env_runner)", + error=True, + ) + if validate_workers_after_construction != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(validate_workers_after_construction)", + new="AlgorithmConfig.env_runners(validate_env_runners_after_" + "construction)", + error=True, + ) + + if env_runner_cls is not NotProvided: + self.env_runner_cls = env_runner_cls + if num_env_runners is not NotProvided: + self.num_env_runners = num_env_runners + if num_envs_per_env_runner is not NotProvided: + if num_envs_per_env_runner <= 0: + raise ValueError( + f"`num_envs_per_env_runner` ({num_envs_per_env_runner}) must be " + "larger 0!" + ) + self.num_envs_per_env_runner = num_envs_per_env_runner + if gym_env_vectorize_mode is not NotProvided: + self.gym_env_vectorize_mode = gym_env_vectorize_mode + if num_cpus_per_env_runner is not NotProvided: + self.num_cpus_per_env_runner = num_cpus_per_env_runner + if num_gpus_per_env_runner is not NotProvided: + self.num_gpus_per_env_runner = num_gpus_per_env_runner + if custom_resources_per_env_runner is not NotProvided: + self.custom_resources_per_env_runner = custom_resources_per_env_runner + + if sample_timeout_s is not NotProvided: + self.sample_timeout_s = sample_timeout_s + if max_requests_in_flight_per_env_runner is not NotProvided: + self.max_requests_in_flight_per_env_runner = ( + max_requests_in_flight_per_env_runner + ) + if sample_collector is not NotProvided: + self.sample_collector = sample_collector + if create_local_env_runner is not NotProvided: + self.create_local_env_runner = create_local_env_runner + if create_env_on_local_worker is not NotProvided: + self.create_env_on_local_worker = create_env_on_local_worker + if env_to_module_connector is not NotProvided: + self._env_to_module_connector = env_to_module_connector + if module_to_env_connector is not NotProvided: + self._module_to_env_connector = module_to_env_connector + if add_default_connectors_to_env_to_module_pipeline is not NotProvided: + self.add_default_connectors_to_env_to_module_pipeline = ( + add_default_connectors_to_env_to_module_pipeline + ) + if add_default_connectors_to_module_to_env_pipeline is not NotProvided: + self.add_default_connectors_to_module_to_env_pipeline = ( + add_default_connectors_to_module_to_env_pipeline + ) + if episode_lookback_horizon is not NotProvided: + self.episode_lookback_horizon = episode_lookback_horizon + if merge_env_runner_states is not NotProvided: + self.merge_env_runner_states = merge_env_runner_states + if broadcast_env_runner_states is not NotProvided: + self.broadcast_env_runner_states = broadcast_env_runner_states + if use_worker_filter_stats is not NotProvided: + self.use_worker_filter_stats = use_worker_filter_stats + if update_worker_filter_stats is not NotProvided: + self.update_worker_filter_stats = update_worker_filter_stats + if rollout_fragment_length is not NotProvided: + if not ( + ( + isinstance(rollout_fragment_length, int) + and rollout_fragment_length > 0 + ) + or rollout_fragment_length == "auto" + ): + raise ValueError("`rollout_fragment_length` must be int >0 or 'auto'!") + self.rollout_fragment_length = rollout_fragment_length + if batch_mode is not NotProvided: + if batch_mode not in ["truncate_episodes", "complete_episodes"]: + raise ValueError( + f"`batch_mode` ({batch_mode}) must be one of [truncate_episodes|" + "complete_episodes]!" + ) + self.batch_mode = batch_mode + if explore is not NotProvided: + self.explore = explore + if episodes_to_numpy is not NotProvided: + self.episodes_to_numpy = episodes_to_numpy + + # @OldAPIStack + if exploration_config is not NotProvided: + # Override entire `exploration_config` if `type` key changes. + # Update, if `type` key remains the same or is not specified. + new_exploration_config = deep_update( + {"exploration_config": self.exploration_config}, + {"exploration_config": exploration_config}, + False, + ["exploration_config"], + ["exploration_config"], + ) + self.exploration_config = new_exploration_config["exploration_config"] + if remote_worker_envs is not NotProvided: + self.remote_worker_envs = remote_worker_envs + if remote_env_batch_wait_ms is not NotProvided: + self.remote_env_batch_wait_ms = remote_env_batch_wait_ms + if validate_env_runners_after_construction is not NotProvided: + self.validate_env_runners_after_construction = ( + validate_env_runners_after_construction + ) + if preprocessor_pref is not NotProvided: + self.preprocessor_pref = preprocessor_pref + if observation_filter is not NotProvided: + self.observation_filter = observation_filter + if synchronize_filter is not NotProvided: + self.synchronize_filters = synchronize_filter + if compress_observations is not NotProvided: + self.compress_observations = compress_observations + if enable_tf1_exec_eagerly is not NotProvided: + self.enable_tf1_exec_eagerly = enable_tf1_exec_eagerly + if sampler_perf_stats_ema_coef is not NotProvided: + self.sampler_perf_stats_ema_coef = sampler_perf_stats_ema_coef + + # Deprecated settings. + if synchronize_filter != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(synchronize_filter=..)", + new="AlgorithmConfig.env_runners(update_worker_filter_stats=..)", + error=True, + ) + if ignore_worker_failures != DEPRECATED_VALUE: + deprecation_warning( + old="ignore_worker_failures is deprecated, and will soon be a no-op", + error=True, + ) + if recreate_failed_workers != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(recreate_failed_workers=..)", + new="AlgorithmConfig.fault_tolerance(recreate_failed_workers=..)", + error=True, + ) + if restart_failed_sub_environments != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(restart_failed_sub_environments=..)", + new=( + "AlgorithmConfig.fault_tolerance(" + "restart_failed_sub_environments=..)" + ), + error=True, + ) + if num_consecutive_worker_failures_tolerance != DEPRECATED_VALUE: + deprecation_warning( + old=( + "AlgorithmConfig.env_runners(" + "num_consecutive_worker_failures_tolerance=..)" + ), + new=( + "AlgorithmConfig.fault_tolerance(" + "num_consecutive_worker_failures_tolerance=..)" + ), + error=True, + ) + if worker_health_probe_timeout_s != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(worker_health_probe_timeout_s=..)", + new="AlgorithmConfig.fault_tolerance(worker_health_probe_timeout_s=..)", + error=True, + ) + if worker_restore_timeout_s != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.env_runners(worker_restore_timeout_s=..)", + new="AlgorithmConfig.fault_tolerance(worker_restore_timeout_s=..)", + error=True, + ) + + return self + + def learners( + self, + *, + num_learners: Optional[int] = NotProvided, + num_cpus_per_learner: Optional[Union[str, float, int]] = NotProvided, + num_gpus_per_learner: Optional[Union[float, int]] = NotProvided, + num_aggregator_actors_per_learner: Optional[int] = NotProvided, + max_requests_in_flight_per_aggregator_actor: Optional[float] = NotProvided, + local_gpu_idx: Optional[int] = NotProvided, + max_requests_in_flight_per_learner: Optional[int] = NotProvided, + ) -> Self: + """Sets LearnerGroup and Learner worker related configurations. + + Args: + num_learners: Number of Learner workers used for updating the RLModule. + A value of 0 means training takes place on a local Learner on main + process CPUs or 1 GPU (determined by `num_gpus_per_learner`). + For multi-gpu training, you have to set `num_learners` to > 1 and set + `num_gpus_per_learner` accordingly (e.g., 4 GPUs total and model fits on + 1 GPU: `num_learners=4; num_gpus_per_learner=1` OR 4 GPUs total and + model requires 2 GPUs: `num_learners=2; num_gpus_per_learner=2`). + num_cpus_per_learner: Number of CPUs allocated per Learner worker. + If "auto" (default), use 1 if `num_gpus_per_learner=0`, otherwise 0. + Only necessary for custom processing pipeline inside each Learner + requiring multiple CPU cores. + If `num_learners=0`, RLlib creates only one local Learner instance and + the number of CPUs on the main process is + `max(num_cpus_per_learner, num_cpus_for_main_process)`. + num_gpus_per_learner: Number of GPUs allocated per Learner worker. If + `num_learners=0`, any value greater than 0 runs the + training on a single GPU on the main process, while a value of 0 runs + the training on main process CPUs. + num_aggregator_actors_per_learner: The number of aggregator actors per + Learner (if num_learners=0, one local learner is created). Must be at + least 1. Aggregator actors perform the task of a) converting episodes + into a train batch and b) move that train batch to the same GPU that + the corresponding learner is located on. Good values are 1 or 2, but + this strongly depends on your setup and `EnvRunner` throughput. + max_requests_in_flight_per_aggregator_actor: How many in-flight requests + are allowed per aggregator actor before new requests are dropped? + local_gpu_idx: If `num_gpus_per_learner` > 0, and + `num_learners` < 2, then RLlib uses this GPU index for training. This is + an index into the available + CUDA devices. For example if `os.environ["CUDA_VISIBLE_DEVICES"] = "1"` + and `local_gpu_idx=0`, RLlib uses the GPU with ID=1 on the node. + max_requests_in_flight_per_learner: Max number of in-flight requests + to each Learner (actor). You normally do not have to tune this setting + (default is 3), however, for asynchronous algorithms, this determines + the "queue" size for incoming batches (or lists of episodes) into each + Learner worker, thus also determining, how much off-policy'ness would be + acceptable. The off-policy'ness is the difference between the numbers of + updates a policy has undergone on the Learner vs the EnvRunners. + See the `ray.rllib.utils.actor_manager.FaultTolerantActorManager` class + for more details. + + Returns: + This updated AlgorithmConfig object. + """ + if num_learners is not NotProvided: + self.num_learners = num_learners + if num_cpus_per_learner is not NotProvided: + self.num_cpus_per_learner = num_cpus_per_learner + if num_gpus_per_learner is not NotProvided: + self.num_gpus_per_learner = num_gpus_per_learner + if num_aggregator_actors_per_learner is not NotProvided: + self.num_aggregator_actors_per_learner = num_aggregator_actors_per_learner + if max_requests_in_flight_per_aggregator_actor is not NotProvided: + self.max_requests_in_flight_per_aggregator_actor = ( + max_requests_in_flight_per_aggregator_actor + ) + if local_gpu_idx is not NotProvided: + self.local_gpu_idx = local_gpu_idx + if max_requests_in_flight_per_learner is not NotProvided: + self.max_requests_in_flight_per_learner = max_requests_in_flight_per_learner + + return self + + def training( + self, + *, + gamma: Optional[float] = NotProvided, + lr: Optional[LearningRateOrSchedule] = NotProvided, + grad_clip: Optional[float] = NotProvided, + grad_clip_by: Optional[str] = NotProvided, + train_batch_size: Optional[int] = NotProvided, + train_batch_size_per_learner: Optional[int] = NotProvided, + num_epochs: Optional[int] = NotProvided, + minibatch_size: Optional[int] = NotProvided, + shuffle_batch_per_epoch: Optional[bool] = NotProvided, + model: Optional[dict] = NotProvided, + optimizer: Optional[dict] = NotProvided, + # Deprecated args. + num_aggregator_actors_per_learner=DEPRECATED_VALUE, + max_requests_in_flight_per_aggregator_actor=DEPRECATED_VALUE, + num_sgd_iter=DEPRECATED_VALUE, + max_requests_in_flight_per_sampler_worker=DEPRECATED_VALUE, + # Moved to `learners()` method. + learner_class: Optional[Type["Learner"]] = NotProvided, + learner_connector: Optional[ + Callable[ + [gym.spaces.Space, gym.spaces.Space], + Union["ConnectorV2", List["ConnectorV2"]], + ] + ] = NotProvided, + add_default_connectors_to_learner_pipeline: Optional[bool] = NotProvided, + learner_config_dict: Optional[Dict[str, Any]] = NotProvided, + ) -> Self: + """Sets the training related configuration. + + Args: + gamma: Float specifying the discount factor of the Markov Decision process. + lr: The learning rate (float) or learning rate schedule in the format of + [[timestep, lr-value], [timestep, lr-value], ...] + In case of a schedule, intermediary timesteps are assigned to + linearly interpolated learning rate values. A schedule config's first + entry must start with timestep 0, i.e.: [[0, initial_value], [...]]. + Note: If you require a) more than one optimizer (per RLModule), + b) optimizer types that are not Adam, c) a learning rate schedule that + is not a linearly interpolated, piecewise schedule as described above, + or d) specifying c'tor arguments of the optimizer that are not the + learning rate (e.g. Adam's epsilon), then you must override your + Learner's `configure_optimizer_for_module()` method and handle + lr-scheduling yourself. + grad_clip: If None, no gradient clipping is applied. Otherwise, + depending on the setting of `grad_clip_by`, the (float) value of + `grad_clip` has the following effect: + If `grad_clip_by=value`: Clips all computed gradients individually + inside the interval [-`grad_clip`, +`grad_clip`]. + If `grad_clip_by=norm`, computes the L2-norm of each weight/bias + gradient tensor individually and then clip all gradients such that these + L2-norms do not exceed `grad_clip`. The L2-norm of a tensor is computed + via: `sqrt(SUM(w0^2, w1^2, ..., wn^2))` where w[i] are the elements of + the tensor (no matter what the shape of this tensor is). + If `grad_clip_by=global_norm`, computes the square of the L2-norm of + each weight/bias gradient tensor individually, sum up all these squared + L2-norms across all given gradient tensors (e.g. the entire module to + be updated), square root that overall sum, and then clip all gradients + such that this global L2-norm does not exceed the given value. + The global L2-norm over a list of tensors (e.g. W and V) is computed + via: + `sqrt[SUM(w0^2, w1^2, ..., wn^2) + SUM(v0^2, v1^2, ..., vm^2)]`, where + w[i] and v[j] are the elements of the tensors W and V (no matter what + the shapes of these tensors are). + grad_clip_by: See `grad_clip` for the effect of this setting on gradient + clipping. Allowed values are `value`, `norm`, and `global_norm`. + train_batch_size_per_learner: Train batch size per individual Learner + worker. This setting only applies to the new API stack. The number + of Learner workers can be set via `config.resources( + num_learners=...)`. The total effective batch size is then + `num_learners` x `train_batch_size_per_learner` and you can + access it with the property `AlgorithmConfig.total_train_batch_size`. + train_batch_size: Training batch size, if applicable. When on the new API + stack, this setting should no longer be used. Instead, use + `train_batch_size_per_learner` (in combination with + `num_learners`). + num_epochs: The number of complete passes over the entire train batch (per + Learner). Each pass might be further split into n minibatches (if + `minibatch_size` provided). + minibatch_size: The size of minibatches to use to further split the train + batch into. + shuffle_batch_per_epoch: Whether to shuffle the train batch once per epoch. + If the train batch has a time rank (axis=1), shuffling only takes + place along the batch axis to not disturb any intact (episode) + trajectories. + model: Arguments passed into the policy model. See models/catalog.py for a + full list of the available model options. + TODO: Provide ModelConfig objects instead of dicts. + optimizer: Arguments to pass to the policy optimizer. This setting is not + used when `enable_rl_module_and_learner=True`. + + Returns: + This updated AlgorithmConfig object. + """ + if learner_class is not NotProvided: + deprecation_warning( + old="config.training(learner_class=..)", + new="config.learners(learner_class=..)", + error=False, + ) + self._learner_class = learner_class + if learner_connector is not NotProvided: + deprecation_warning( + old="config.training(learner_connector=..)", + new="config.learners(learner_connector=..)", + error=False, + ) + self._learner_connector = learner_connector + if add_default_connectors_to_learner_pipeline is not NotProvided: + deprecation_warning( + old="config.training(add_default_connectors_to_learner_pipeline=..)", + new="config.learners(add_default_connectors_to_learner_pipeline=..)", + error=False, + ) + self.add_default_connectors_to_learner_pipeline = ( + add_default_connectors_to_learner_pipeline + ) + if learner_config_dict is not NotProvided: + deprecation_warning( + old="config.training(learner_config_dict=..)", + new="config.learners(learner_config_dict=..)", + error=False, + ) + self.learner_config_dict.update(learner_config_dict) + + if num_aggregator_actors_per_learner != DEPRECATED_VALUE: + deprecation_warning( + old="config.training(num_aggregator_actors_per_learner=..)", + new="config.learners(num_aggregator_actors_per_learner=..)", + error=False, + ) + self.num_aggregator_actors_per_learner = num_aggregator_actors_per_learner + if max_requests_in_flight_per_aggregator_actor != DEPRECATED_VALUE: + deprecation_warning( + old="config.training(max_requests_in_flight_per_aggregator_actor=..)", + new="config.learners(max_requests_in_flight_per_aggregator_actor=..)", + error=False, + ) + self.max_requests_in_flight_per_aggregator_actor = ( + max_requests_in_flight_per_aggregator_actor + ) + + if num_sgd_iter != DEPRECATED_VALUE: + deprecation_warning( + old="config.training(num_sgd_iter=..)", + new="config.training(num_epochs=..)", + error=False, + ) + num_epochs = num_sgd_iter + if max_requests_in_flight_per_sampler_worker != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.training(" + "max_requests_in_flight_per_sampler_worker=...)", + new="AlgorithmConfig.env_runners(" + "max_requests_in_flight_per_env_runner=...)", + error=False, + ) + self.env_runners( + max_requests_in_flight_per_env_runner=( + max_requests_in_flight_per_sampler_worker + ), + ) + + if gamma is not NotProvided: + self.gamma = gamma + if lr is not NotProvided: + self.lr = lr + if grad_clip is not NotProvided: + self.grad_clip = grad_clip + if grad_clip_by is not NotProvided: + if grad_clip_by not in ["value", "norm", "global_norm"]: + raise ValueError( + f"`grad_clip_by` ({grad_clip_by}) must be one of: 'value', 'norm', " + "or 'global_norm'!" + ) + self.grad_clip_by = grad_clip_by + if train_batch_size_per_learner is not NotProvided: + self._train_batch_size_per_learner = train_batch_size_per_learner + if train_batch_size is not NotProvided: + self.train_batch_size = train_batch_size + if num_epochs is not NotProvided: + self.num_epochs = num_epochs + if minibatch_size is not NotProvided: + self.minibatch_size = minibatch_size + if shuffle_batch_per_epoch is not NotProvided: + self.shuffle_batch_per_epoch = shuffle_batch_per_epoch + + if model is not NotProvided: + self.model.update(model) + if ( + model.get("_use_default_native_models", DEPRECATED_VALUE) + != DEPRECATED_VALUE + ): + deprecation_warning( + old="AlgorithmConfig.training(_use_default_native_models=True)", + help="_use_default_native_models is not supported " + "anymore. To get rid of this error, set `config.api_stack(" + "enable_rl_module_and_learner=True)`. Native models will " + "be better supported by the upcoming RLModule API.", + # Error out if user tries to enable this. + error=model["_use_default_native_models"], + ) + if optimizer is not NotProvided: + self.optimizer = merge_dicts(self.optimizer, optimizer) + + return self + + def callbacks( + self, + callbacks_class: Optional[ + Union[Type[RLlibCallback], List[Type[RLlibCallback]]] + ] = NotProvided, + *, + on_algorithm_init: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_train_result: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_evaluate_start: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_evaluate_end: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_evaluate_offline_start: Optional[ + Union[Callable, List[Callable]] + ] = NotProvided, + on_evaluate_offline_end: Optional[ + Union[Callable, List[Callable]] + ] = NotProvided, + on_env_runners_recreated: Optional[ + Union[Callable, List[Callable]] + ] = NotProvided, + on_offline_eval_runners_recreated: Optional[ + Union[Callable, List[Callable]] + ] = NotProvided, + on_checkpoint_loaded: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_environment_created: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_episode_created: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_episode_start: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_episode_step: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_episode_end: Optional[Union[Callable, List[Callable]]] = NotProvided, + on_sample_end: Optional[Union[Callable, List[Callable]]] = NotProvided, + ) -> Self: + """Sets the callbacks configuration. + + Args: + callbacks_class: RLlibCallback class, whose methods are called during + various phases of training and RL environment sample collection. + TODO (sven): Change the link to new rst callbacks page. + See the `RLlibCallback` class and + `examples/metrics/custom_metrics_and_callbacks.py` for more information. + on_algorithm_init: A callable or a list of callables. If a list, RLlib calls + the items in the same sequence. `on_algorithm_init` methods overridden + in `callbacks_class` take precedence and are called first. + See + :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_algorithm_init` # noqa + for more information. + on_evaluate_start: A callable or a list of callables. If a list, RLlib calls + the items in the same sequence. `on_evaluate_start` methods overridden + in `callbacks_class` take precedence and are called first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_evaluate_start` # noqa + for more information. + on_evaluate_end: A callable or a list of callables. If a list, RLlib calls + the items in the same sequence. `on_evaluate_end` methods overridden + in `callbacks_class` take precedence and are called first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_evaluate_end` # noqa + for more information. + on_env_runners_recreated: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_env_runners_recreated` + methods overridden in `callbacks_class` take precedence and are called + first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_env_runners_recreated` # noqa + for more information. + on_checkpoint_loaded: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_checkpoint_loaded` + methods overridden in `callbacks_class` take precedence and are called + first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_checkpoint_loaded` # noqa + for more information. + on_environment_created: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_environment_created` + methods overridden in `callbacks_class` take precedence and are called + first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_environment_created` # noqa + for more information. + on_episode_created: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_episode_created` methods + overridden in `callbacks_class` take precedence and are called first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_episode_created` # noqa + for more information. + on_episode_start: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_episode_start` methods + overridden in `callbacks_class` take precedence and are called first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_episode_start` # noqa + for more information. + on_episode_step: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_episode_step` methods + overridden in `callbacks_class` take precedence and are called first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_episode_step` # noqa + for more information. + on_episode_end: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_episode_end` methods + overridden in `callbacks_class` take precedence and are called first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_episode_end` # noqa + for more information. + on_sample_end: A callable or a list of callables. If a list, + RLlib calls the items in the same sequence. `on_sample_end` methods + overridden in `callbacks_class` take precedence and are called first. + See :py:meth:`~ray.rllib.callbacks.callbacks.RLlibCallback.on_sample_end` # noqa + for more information. + + Returns: + This updated AlgorithmConfig object. + """ + if callbacks_class is None: + callbacks_class = RLlibCallback + + if callbacks_class is not NotProvided: + # Check, whether given `callbacks` is a callable. + # TODO (sven): Once the old API stack is deprecated, this can also be None + # (which should then become the default value for this attribute). + to_check = force_list(callbacks_class) + if not all(callable(c) for c in to_check): + raise ValueError( + "`config.callbacks_class` must be a callable or list of callables that " + "returns a subclass of DefaultCallbacks, got " + f"{callbacks_class}!" + ) + self.callbacks_class = callbacks_class + if on_algorithm_init is not NotProvided: + self.callbacks_on_algorithm_init = on_algorithm_init + if on_train_result is not NotProvided: + self.callbacks_on_train_result = on_train_result + if on_evaluate_start is not NotProvided: + self.callbacks_on_evaluate_start = on_evaluate_start + if on_evaluate_end is not NotProvided: + self.callbacks_on_evaluate_end = on_evaluate_end + if on_evaluate_offline_start is not NotProvided: + self.callbacks_on_evaluate_offline_start = on_evaluate_offline_start + if on_evaluate_offline_end is not NotProvided: + self.callbacks_on_evaluate_offline_end = on_evaluate_offline_end + if on_env_runners_recreated is not NotProvided: + self.callbacks_on_env_runners_recreated = on_env_runners_recreated + if on_offline_eval_runners_recreated is not NotProvided: + self.callbacks_on_offline_eval_runners_recreated = ( + on_offline_eval_runners_recreated + ) + if on_checkpoint_loaded is not NotProvided: + self.callbacks_on_checkpoint_loaded = on_checkpoint_loaded + if on_environment_created is not NotProvided: + self.callbacks_on_environment_created = on_environment_created + if on_episode_created is not NotProvided: + self.callbacks_on_episode_created = on_episode_created + if on_episode_start is not NotProvided: + self.callbacks_on_episode_start = on_episode_start + if on_episode_step is not NotProvided: + self.callbacks_on_episode_step = on_episode_step + if on_episode_end is not NotProvided: + self.callbacks_on_episode_end = on_episode_end + if on_sample_end is not NotProvided: + self.callbacks_on_sample_end = on_sample_end + + return self + + def evaluation( + self, + *, + evaluation_interval: Optional[int] = NotProvided, + evaluation_duration: Optional[Union[int, str]] = NotProvided, + evaluation_duration_unit: Optional[str] = NotProvided, + evaluation_auto_duration_min_env_steps_per_sample: Optional[int] = NotProvided, + evaluation_auto_duration_max_env_steps_per_sample: Optional[int] = NotProvided, + evaluation_sample_timeout_s: Optional[float] = NotProvided, + evaluation_parallel_to_training: Optional[bool] = NotProvided, + evaluation_force_reset_envs_before_iteration: Optional[bool] = NotProvided, + evaluation_config: Optional[ + Union["AlgorithmConfig", PartialAlgorithmConfigDict] + ] = NotProvided, + off_policy_estimation_methods: Optional[Dict] = NotProvided, + ope_split_batch_by_episode: Optional[bool] = NotProvided, + evaluation_num_env_runners: Optional[int] = NotProvided, + custom_evaluation_function: Optional[Callable] = NotProvided, + # Offline evaluation. + offline_evaluation_interval: Optional[int] = NotProvided, + num_offline_eval_runners: Optional[int] = NotProvided, + offline_evaluation_type: Optional[Callable] = NotProvided, + offline_eval_runner_class: Optional[Callable] = NotProvided, + offline_loss_for_module_fn: Optional[Callable] = NotProvided, + offline_eval_batch_size_per_runner: Optional[int] = NotProvided, + dataset_num_iters_per_offline_eval_runner: Optional[int] = NotProvided, + offline_eval_rl_module_inference_only: Optional[bool] = NotProvided, + num_cpus_per_offline_eval_runner: Optional[int] = NotProvided, + num_gpus_per_offline_eval_runner: Optional[int] = NotProvided, + custom_resources_per_offline_eval_runner: Optional[ + Dict[str, Any] + ] = NotProvided, + offline_evaluation_timeout_s: Optional[float] = NotProvided, + max_requests_in_flight_per_offline_eval_runner: Optional[int] = NotProvided, + broadcast_offline_eval_runner_states: Optional[bool] = NotProvided, + validate_offline_eval_runners_after_construction: Optional[bool] = NotProvided, + restart_failed_offline_eval_runners: Optional[bool] = NotProvided, + ignore_offline_eval_runner_failures: Optional[bool] = NotProvided, + max_num_offline_eval_runner_restarts: Optional[int] = NotProvided, + offline_eval_runner_health_probe_timeout_s: Optional[float] = NotProvided, + offline_eval_runner_restore_timeout_s: Optional[float] = NotProvided, + # Deprecated args. + always_attach_evaluation_results=DEPRECATED_VALUE, + evaluation_num_workers=DEPRECATED_VALUE, + ) -> Self: + """Sets the config's evaluation settings. + + Args: + evaluation_interval: Evaluate with every `evaluation_interval` training + iterations. The evaluation stats are reported under the "evaluation" + metric key. Set to None (or 0) for no evaluation. + evaluation_duration: Duration for which to run evaluation each + `evaluation_interval`. The unit for the duration can be set via + `evaluation_duration_unit` to either "episodes" (default) or + "timesteps". If using multiple evaluation workers (EnvRunners) in the + `evaluation_num_env_runners > 1` setting, the amount of + episodes/timesteps to run are split amongst these. + A special value of "auto" can be used in case + `evaluation_parallel_to_training=True`. This is the recommended way when + trying to save as much time on evaluation as possible. The Algorithm + then runs as many timesteps via the evaluation workers as possible, + while not taking longer than the parallely running training step and + thus, never wasting any idle time on either training- or evaluation + workers. When using this setting (`evaluation_duration="auto"`), it is + strongly advised to set `evaluation_interval=1` and + `evaluation_force_reset_envs_before_iteration=True` at the same time. + evaluation_duration_unit: The unit, with which to count the evaluation + duration. Either "episodes" (default) or "timesteps". Note that this + setting is ignored if `evaluation_duration="auto"`. + evaluation_auto_duration_min_env_steps_per_sample: If `evaluation_duration` + is "auto" (in which case `evaluation_duration_unit` is always + "timesteps"), at least how many timesteps should be done per remote + `sample()` call. + evaluation_auto_duration_max_env_steps_per_sample: If `evaluation_duration` + is "auto" (in which case `evaluation_duration_unit` is always + "timesteps"), at most how many timesteps should be done per remote + `sample()` call. + evaluation_sample_timeout_s: The timeout (in seconds) for evaluation workers + to sample a complete episode in the case your config settings are: + `evaluation_duration != auto` and `evaluation_duration_unit=episode`. + After this time, the user receives a warning and instructions on how + to fix the issue. + evaluation_parallel_to_training: Whether to run evaluation in parallel to + the `Algorithm.training_step()` call, using threading. Default=False. + E.g. for evaluation_interval=1 -> In every call to `Algorithm.train()`, + the `Algorithm.training_step()` and `Algorithm.evaluate()` calls + run in parallel. Note that this setting - albeit extremely efficient b/c + it wastes no extra time for evaluation - causes the evaluation results + to lag one iteration behind the rest of the training results. This is + important when picking a good checkpoint. For example, if iteration 42 + reports a good evaluation `episode_return_mean`, be aware that these + results were achieved on the weights trained in iteration 41, so you + should probably pick the iteration 41 checkpoint instead. + evaluation_force_reset_envs_before_iteration: Whether all environments + should be force-reset (even if they are not done yet) right before + the evaluation step of the iteration begins. Setting this to True + (default) makes sure that the evaluation results aren't polluted with + episode statistics that were actually (at least partially) achieved with + an earlier set of weights. Note that this setting is only + supported on the new API stack w/ EnvRunners and ConnectorV2 + (`config.enable_rl_module_and_learner=True` AND + `config.enable_env_runner_and_connector_v2=True`). + evaluation_config: Typical usage is to pass extra args to evaluation env + creator and to disable exploration by computing deterministic actions. + IMPORTANT NOTE: Policy gradient algorithms are able to find the optimal + policy, even if this is a stochastic one. Setting "explore=False" here + results in the evaluation workers not using this optimal policy! + off_policy_estimation_methods: Specify how to evaluate the current policy, + along with any optional config parameters. This only has an effect when + reading offline experiences ("input" is not "sampler"). + Available keys: + {ope_method_name: {"type": ope_type, ...}} where `ope_method_name` + is a user-defined string to save the OPE results under, and + `ope_type` can be any subclass of OffPolicyEstimator, e.g. + ray.rllib.offline.estimators.is::ImportanceSampling + or your own custom subclass, or the full class path to the subclass. + You can also add additional config arguments to be passed to the + OffPolicyEstimator in the dict, e.g. + {"qreg_dr": {"type": DoublyRobust, "q_model_type": "qreg", "k": 5}} + ope_split_batch_by_episode: Whether to use SampleBatch.split_by_episode() to + split the input batch to episodes before estimating the ope metrics. In + case of bandits you should make this False to see improvements in ope + evaluation speed. In case of bandits, it is ok to not split by episode, + since each record is one timestep already. The default is True. + evaluation_num_env_runners: Number of parallel EnvRunners to use for + evaluation. Note that this is set to zero by default, which means + evaluation is run in the algorithm process (only if + `evaluation_interval` is not 0 or None). If you increase this, also + increases the Ray resource usage of the algorithm since evaluation + workers are created separately from those EnvRunners used to sample data + for training. + custom_evaluation_function: Customize the evaluation method. This must be a + function of signature (algo: Algorithm, eval_workers: EnvRunnerGroup) -> + (metrics: dict, env_steps: int, agent_steps: int) (metrics: dict if + `enable_env_runner_and_connector_v2=True`), where `env_steps` and + `agent_steps` define the number of sampled steps during the evaluation + iteration. See the Algorithm.evaluate() method to see the default + implementation. The Algorithm guarantees all eval workers have the + latest policy state before this function is called. + offline_evaluation_interval: Evaluate offline with every + `offline_evaluation_interval` training iterations. The offline evaluation + stats are reported under the "evaluation/offline_evaluation" metric key. Set + to None (or 0) for no offline evaluation. + num_offline_eval_runners: Number of OfflineEvaluationRunner actors to create + for parallel evaluation. Setting this to 0 forces sampling to be done in the + local OfflineEvaluationRunner (main process or the Algorithm's actor when + using Tune). + offline_evaluation_type: Type of offline evaluation to run. Either `"eval_loss"` + for evaluating the validation loss of the policy, `"is"` for importance + sampling, or `"pdis"` for per-decision importance sampling. If you want to + implement your own offline evaluation method write an `OfflineEvaluationRunner` + and use the `AlgorithmConfig.offline_eval_runner_class`. + offline_eval_runner_class: An `OfflineEvaluationRunner` class that implements + custom offline evaluation logic. + offline_loss_for_module_fn: A callable to compute the loss per `RLModule` in + offline evaluation. If not provided the training loss function ( + `Learner.compute_loss_for_module`) is used. The signature must be ( + runner: OfflineEvaluationRunner, module_id: ModuleID, config: AlgorithmConfig, + batch: Dict[str, Any], fwd_out: Dict[str, TensorType]). + offline_eval_batch_size_per_runner: Evaluation batch size per individual + OfflineEvaluationRunner worker. This setting only applies to the new API + stack. The number of OfflineEvaluationRunner workers can be set via + `config.evaluation(num_offline_eval_runners=...)`. The total effective batch + size is then `num_offline_eval_runners` x + `offline_eval_batch_size_per_runner`. + dataset_num_iters_per_offline_eval_runner: Number of batches to evaluate in each + OfflineEvaluationRunner during a single evaluation. If None, each learner runs a + complete epoch over its data block (the dataset is partitioned into + at least as many blocks as there are runners). The default is `1`. + offline_eval_rl_module_inference_only: If `True`, the module spec is used in an + inference-only setting (no-loss) and the RLModule can thus be built in + its light version (if available). For example, the `inference_only` + version of an RLModule might only contain the networks required for + computing actions, but misses additional target- or critic networks. + Also, if `True`, the module does NOT contain those (sub) RLModules that have + their `learner_only` flag set to True. + num_cpus_per_offline_eval_runner: Number of CPUs to allocate per + OfflineEvaluationRunner. + num_gpus_per_offline_eval_runner: Number of GPUs to allocate per + OfflineEvaluationRunner. This can be fractional. This is usually needed only if + your (custom) loss function itself requires a GPU (i.e., it contains GPU- + intensive computations), or model inference is unusually expensive. + custom_resources_per_eval_runner: Any custom Ray resources to allocate per + OfflineEvaluationRunner. + offline_evaluation_timeout_s: The timeout in seconds for calling `run()` on remote + OfflineEvaluationRunner workers. Results (episode list) from workers that take + longer than this time are discarded. + max_requests_in_flight_per_offline_eval_runner: Max number of in-flight requests + to each OfflineEvaluationRunner (actor)). See the + `ray.rllib.utils.actor_manager.FaultTolerantActorManager` class for more + details. + Tuning these values is important when running experiments with + large evaluation batches, where there is the risk that the object store may + fill up, causing spilling of objects to disk. This can cause any + asynchronous requests to become very slow, making your experiment run + slowly as well. You can inspect the object store during your experiment + through a call to `ray memory` on your head node, and by using the Ray + dashboard. If you're seeing that the object store is filling up, + turn down the number of remote requests in flight or enable compression + or increase the object store memory through, for example: + `ray.init(object_store_memory=10 * 1024 * 1024 * 1024) # =10 GB`. + broadcast_offline_eval_runner_states: True, if merged OfflineEvaluationRunner + states (from the central connector pipelines) should be broadcast back to + all remote OfflineEvaluationRunner actors. + validate_offline_eval_runners_after_construction: Whether to validate that each + created remote OfflineEvaluationRunner is healthy after its construction process. + restart_failed_offline_eval_runners: Whether - upon an OfflineEvaluationRunner + failure - RLlib tries to restart the lost OfflineEvaluationRunner(s) as an + identical copy of the failed one(s). You should set this to True when training + on SPOT instances that may preempt any time and/or if you need to evaluate always a + complete dataset b/c OfflineEvaluationRunner(s) evaluate through streaming split + iterators on disjoint batches. The new, recreated OfflineEvaluationRunner(s) only + differ from the failed one in their `self.recreated_worker=True` property value + and have the same `worker_index` as the original(s). If this setting is True, the + value of the `ignore_offline_eval_runner_failures` setting is ignored. + ignore_offline_eval_runner_failures: Whether to ignore any OfflineEvalautionRunner + failures and continue running with the remaining OfflineEvaluationRunners. This + setting is ignored, if `restart_failed_offline_eval_runners=True`. + max_num_offline_eval_runner_restarts: The maximum number of times any + OfflineEvaluationRunner is allowed to be restarted (if + `restart_failed_offline_eval_runners` is True). + offline_eval_runner_health_probe_timeout_s: Max amount of time in seconds, we should + spend waiting for OfflineEvaluationRunner health probe calls + (`OfflineEvaluationRunner.ping.remote()`) to respond. Health pings are very cheap, + however, we perform the health check via a blocking `ray.get()`, so the + default value should not be too large. + offline_eval_runner_restore_timeout_s: Max amount of time we should wait to restore + states on recovered OfflineEvaluationRunner actors. Default is 30 mins. + + Returns: + This updated AlgorithmConfig object. + """ + if always_attach_evaluation_results != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.evaluation(always_attach_evaluation_results=..)", + help="This setting is no longer needed, b/c Tune does not error " + "anymore (only warns) when a metrics key can't be found in the " + "results.", + error=True, + ) + if evaluation_num_workers != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.evaluation(evaluation_num_workers=..)", + new="AlgorithmConfig.evaluation(evaluation_num_env_runners=..)", + error=False, + ) + self.evaluation_num_env_runners = evaluation_num_workers + + if evaluation_interval is not NotProvided: + self.evaluation_interval = evaluation_interval + if evaluation_duration is not NotProvided: + self.evaluation_duration = evaluation_duration + if evaluation_duration_unit is not NotProvided: + self.evaluation_duration_unit = evaluation_duration_unit + if evaluation_auto_duration_min_env_steps_per_sample is not NotProvided: + self.evaluation_auto_duration_min_env_steps_per_sample = ( + evaluation_auto_duration_min_env_steps_per_sample + ) + if evaluation_auto_duration_max_env_steps_per_sample is not NotProvided: + self.evaluation_auto_duration_max_env_steps_per_sample = ( + evaluation_auto_duration_max_env_steps_per_sample + ) + if evaluation_sample_timeout_s is not NotProvided: + self.evaluation_sample_timeout_s = evaluation_sample_timeout_s + if evaluation_parallel_to_training is not NotProvided: + self.evaluation_parallel_to_training = evaluation_parallel_to_training + if evaluation_force_reset_envs_before_iteration is not NotProvided: + self.evaluation_force_reset_envs_before_iteration = ( + evaluation_force_reset_envs_before_iteration + ) + if evaluation_config is not NotProvided: + # If user really wants to set this to None, we should allow this here, + # instead of creating an empty dict. + if evaluation_config is None: + self.evaluation_config = None + # Update (don't replace) the existing overrides with the provided ones. + else: + from ray.rllib.algorithms.algorithm import Algorithm + + self.evaluation_config = deep_update( + self.evaluation_config or {}, + evaluation_config, + True, + Algorithm._allow_unknown_subkeys, + Algorithm._override_all_subkeys_if_type_changes, + Algorithm._override_all_key_list, + ) + if off_policy_estimation_methods is not NotProvided: + self.off_policy_estimation_methods = off_policy_estimation_methods + if evaluation_num_env_runners is not NotProvided: + self.evaluation_num_env_runners = evaluation_num_env_runners + if custom_evaluation_function is not NotProvided: + self.custom_evaluation_function = custom_evaluation_function + if ope_split_batch_by_episode is not NotProvided: + self.ope_split_batch_by_episode = ope_split_batch_by_episode + # Offline evaluation. + if offline_evaluation_interval is not NotProvided: + self.offline_evaluation_interval = offline_evaluation_interval + if num_offline_eval_runners is not NotProvided: + self.num_offline_eval_runners = num_offline_eval_runners + if offline_evaluation_type is not NotProvided: + self.offline_evaluation_type = offline_evaluation_type + if offline_eval_runner_class is not NotProvided: + self.offline_eval_runner_class = offline_eval_runner_class + if offline_loss_for_module_fn is not NotProvided: + self.offline_loss_for_module_fn = offline_loss_for_module_fn + if offline_eval_batch_size_per_runner is not NotProvided: + self.offline_eval_batch_size_per_runner = offline_eval_batch_size_per_runner + if dataset_num_iters_per_offline_eval_runner is not NotProvided: + self.dataset_num_iters_per_eval_runner = ( + dataset_num_iters_per_offline_eval_runner + ) + if offline_eval_rl_module_inference_only is not NotProvided: + self.offline_eval_rl_module_inference_only = ( + offline_eval_rl_module_inference_only + ) + if num_cpus_per_offline_eval_runner is not NotProvided: + self.num_cpus_per_offline_eval_runner = num_cpus_per_offline_eval_runner + if num_gpus_per_offline_eval_runner is not NotProvided: + self.num_gpus_per_offline_eval_runner = num_gpus_per_offline_eval_runner + if custom_resources_per_offline_eval_runner is not NotProvided: + self.custom_resources_per_offline_eval_runner = ( + custom_resources_per_offline_eval_runner + ) + if offline_evaluation_timeout_s is not NotProvided: + self.offline_evaluation_timeout_s = offline_evaluation_timeout_s + if max_requests_in_flight_per_offline_eval_runner is not NotProvided: + self.max_requests_in_flight_per_offline_eval_runner = ( + max_requests_in_flight_per_offline_eval_runner + ) + if broadcast_offline_eval_runner_states is not NotProvided: + self.broadcast_offline_eval_runner_states = ( + broadcast_offline_eval_runner_states + ) + if validate_offline_eval_runners_after_construction is not NotProvided: + self.validate_offline_eval_runners_after_construction = ( + validate_offline_eval_runners_after_construction + ) + if restart_failed_offline_eval_runners is not NotProvided: + self.restart_failed_offline_eval_runners = ( + restart_failed_offline_eval_runners + ) + if ignore_offline_eval_runner_failures is not NotProvided: + self.ignore_offline_eval_runner_failures = ( + ignore_offline_eval_runner_failures + ) + if max_num_offline_eval_runner_restarts is not NotProvided: + self.max_num_offline_eval_runner_restarts = ( + max_num_offline_eval_runner_restarts + ) + if offline_eval_runner_health_probe_timeout_s is not NotProvided: + self.offline_eval_runner_health_probe_timeout_s = ( + offline_eval_runner_health_probe_timeout_s + ) + if offline_eval_runner_restore_timeout_s is not NotProvided: + self.offline_eval_runner_restore_timeout_s = ( + offline_eval_runner_restore_timeout_s + ) + + return self + + def offline_data( + self, + *, + input_: Optional[Union[str, Callable[[IOContext], InputReader]]] = NotProvided, + offline_data_class: Optional[Type] = NotProvided, + input_read_method: Optional[Union[str, Callable]] = NotProvided, + input_read_method_kwargs: Optional[Dict] = NotProvided, + input_read_schema: Optional[Dict[str, str]] = NotProvided, + input_read_episodes: Optional[bool] = NotProvided, + input_read_sample_batches: Optional[bool] = NotProvided, + input_read_batch_size: Optional[int] = NotProvided, + input_filesystem: Optional[str] = NotProvided, + input_filesystem_kwargs: Optional[Dict] = NotProvided, + input_compress_columns: Optional[List[str]] = NotProvided, + materialize_data: Optional[bool] = NotProvided, + materialize_mapped_data: Optional[bool] = NotProvided, + map_batches_kwargs: Optional[Dict] = NotProvided, + iter_batches_kwargs: Optional[Dict] = NotProvided, + ignore_final_observation: Optional[bool] = NotProvided, + prelearner_class: Optional[Type] = NotProvided, + prelearner_buffer_class: Optional[Type] = NotProvided, + prelearner_buffer_kwargs: Optional[Dict] = NotProvided, + prelearner_module_synch_period: Optional[int] = NotProvided, + dataset_num_iters_per_learner: Optional[int] = NotProvided, + input_config: Optional[Dict] = NotProvided, + actions_in_input_normalized: Optional[bool] = NotProvided, + postprocess_inputs: Optional[bool] = NotProvided, + shuffle_buffer_size: Optional[int] = NotProvided, + output: Optional[str] = NotProvided, + output_config: Optional[Dict] = NotProvided, + output_compress_columns: Optional[List[str]] = NotProvided, + output_max_file_size: Optional[float] = NotProvided, + output_max_rows_per_file: Optional[int] = NotProvided, + output_write_remaining_data: Optional[bool] = NotProvided, + output_write_method: Optional[str] = NotProvided, + output_write_method_kwargs: Optional[Dict] = NotProvided, + output_filesystem: Optional[str] = NotProvided, + output_filesystem_kwargs: Optional[Dict] = NotProvided, + output_write_episodes: Optional[bool] = NotProvided, + offline_sampling: Optional[str] = NotProvided, + ) -> Self: + """Sets the config's offline data settings. + + Args: + input_: Specify how to generate experiences: + - "sampler": Generate experiences via online (env) simulation (default). + - A local directory or file glob expression (e.g., "/tmp/*.json"). + - A list of individual file paths/URIs (e.g., ["/tmp/1.json", + "s3://bucket/2.json"]). + - A dict with string keys and sampling probabilities as values (e.g., + {"sampler": 0.4, "/tmp/*.json": 0.4, "s3://bucket/expert.json": 0.2}). + - A callable that takes an `IOContext` object as only arg and returns a + `ray.rllib.offline.InputReader`. + - A string key that indexes a callable with + `tune.registry.register_input` + offline_data_class: An optional `OfflineData` class that is used to define + the offline data pipeline, including the dataset and the sampling + methodology. Override the `OfflineData` class and pass your derived + class here, if you need some primer transformations specific to your + data or your loss. Usually overriding the `OfflinePreLearner` and using + the resulting customization via `prelearner_class` suffices for most + cases. The default is `None` which uses the base `OfflineData` defined + in `ray.rllib.offline.offline_data.OfflineData`. + input_read_method: Read method for the `ray.data.Dataset` to read in the + offline data from `input_`. The default is `read_parquet` for Parquet + files. See https://docs.ray.io/en/latest/data/api/input_output.html for + more info about available read methods in `ray.data`. + input_read_method_kwargs: Keyword args for `input_read_method`. These + are passed by RLlib into the read method without checking. Use these + keyword args together with `map_batches_kwargs` and + `iter_batches_kwargs` to tune the performance of the data pipeline. + It is strongly recommended to rely on Ray Data's automatic read + performance tuning. + input_read_schema: Table schema for converting offline data to episodes. + This schema maps the offline data columns to + ray.rllib.core.columns.Columns: + `{Columns.OBS: 'o_t', Columns.ACTIONS: 'a_t', ...}`. Columns in + the data set that are not mapped via this schema are sorted into + episodes' `extra_model_outputs`. If no schema is passed in the default + schema used is `ray.rllib.offline.offline_data.SCHEMA`. If your data set + contains already the names in this schema, no `input_read_schema` is + needed. The same applies if the data is in RLlib's `EpisodeType` or its + old `SampleBatch` format. + input_read_episodes: Whether offline data is already stored in RLlib's + `EpisodeType` format, i.e. `ray.rllib.env.SingleAgentEpisode` (multi + -agent is planned but not supported, yet). Reading episodes directly + avoids additional transform steps and is usually faster and + therefore the recommended format when your application remains fully + inside of RLlib's schema. The other format is a columnar format and is + agnostic to the RL framework used. Use the latter format, if you are + unsure when to use the data or in which RL framework. The default is + to read column data, for example, `False`. `input_read_episodes`, and + `input_read_sample_batches` can't be `True` at the same time. See + also `output_write_episodes` to define the output data format when + recording. + input_read_sample_batches: Whether offline data is stored in RLlib's old + stack `SampleBatch` type. This is usually the case for older data + recorded with RLlib in JSON line format. Reading in `SampleBatch` + data needs extra transforms and might not concatenate episode chunks + contained in different `SampleBatch`es in the data. If possible avoid + to read `SampleBatch`es and convert them in a controlled form into + RLlib's `EpisodeType` (i.e. `SingleAgentEpisode`). The default is + `False`. `input_read_episodes`, and `input_read_sample_batches` can't + be `True` at the same time. + input_read_batch_size: Batch size to pull from the data set. This could + differ from the `train_batch_size_per_learner`, if a dataset holds + `EpisodeType` (i.e., `SingleAgentEpisode`) or `SampleBatch`, or any + other data type that contains multiple timesteps in a single row of + the dataset. In such cases a single batch of size + `train_batch_size_per_learner` will potentially pull a multiple of + `train_batch_size_per_learner` timesteps from the offline dataset. The + default is `None` in which the `train_batch_size_per_learner` is pulled. + input_filesystem: A cloud filesystem to handle access to cloud storage when + reading experiences. Can be either "gcs" for Google Cloud Storage, + "s3" for AWS S3 buckets, "abs" for Azure Blob Storage, or any + filesystem supported by PyArrow. In general the file path is sufficient + for accessing data from public or local storage systems. See + https://arrow.apache.org/docs/python/filesystems.html for details. + input_filesystem_kwargs: A dictionary holding the kwargs for the filesystem + given by `input_filesystem`. See `gcsfs.GCSFilesystem` for GCS, + `pyarrow.fs.S3FileSystem`, for S3, and `ablfs.AzureBlobFilesystem` for + ABS filesystem arguments. + input_compress_columns: What input columns are compressed with LZ4 in the + input data. If data is stored in RLlib's `SingleAgentEpisode` ( + `MultiAgentEpisode` not supported, yet). Note the providing + `rllib.core.columns.Columns.OBS` also tries to decompress + `rllib.core.columns.Columns.NEXT_OBS`. + materialize_data: Whether the raw data should be materialized in memory. + This boosts performance, but requires enough memory to avoid an OOM, so + make sure that your cluster has the resources available. For very large + data you might want to switch to streaming mode by setting this to + `False` (default). If your algorithm does not need the RLModule in the + Learner connector pipeline or all (learner) connectors are stateless + you should consider setting `materialize_mapped_data` to `True` + instead (and set `materialize_data` to `False`). If your data does not + fit into memory and your Learner connector pipeline requires an RLModule + or is stateful, set both `materialize_data` and + `materialize_mapped_data` to `False`. + materialize_mapped_data: Whether the data should be materialized after + running it through the Learner connector pipeline (i.e. after running + the `OfflinePreLearner`). This improves performance, but should only be + used in case the (learner) connector pipeline does not require an + RLModule and the (learner) connector pipeline is stateless. For example, + MARWIL's Learner connector pipeline requires the RLModule for value + function predictions and training batches would become stale after some + iterations causing learning degradation or divergence. Also ensure that + your cluster has enough memory available to avoid an OOM. If set to + `True` (True), make sure that `materialize_data` is set to `False` to + avoid materialization of two datasets. If your data does not fit into + memory and your Learner connector pipeline requires an RLModule or is + stateful, set both `materialize_data` and `materialize_mapped_data` to + `False`. + map_batches_kwargs: Keyword args for the `map_batches` method. These are + passed into the `ray.data.Dataset.map_batches` method when sampling + without checking. If no arguments passed in the default arguments + `{'concurrency': max(2, num_learners), 'zero_copy_batch': True}` is + used. Use these keyword args together with `input_read_method_kwargs` + and `iter_batches_kwargs` to tune the performance of the data pipeline. + iter_batches_kwargs: Keyword args for the `iter_batches` method. These are + passed into the `ray.data.Dataset.iter_batches` method when sampling + without checking. If no arguments are passed in, the default argument + `{'prefetch_batches': 2}` is used. Use these keyword args + together with `input_read_method_kwargs` and `map_batches_kwargs` to + tune the performance of the data pipeline. + ignore_final_observation: If the final observation in an episode chunk should + be ignored. This concerns mainly column-based data and instead of using a + user-provided `NEXT_OBS` sets final observations to zero. This should be + used with BC only, as in true Offline RL algorithms the final observation + is important. + prelearner_class: An optional `OfflinePreLearner` class that is used to + transform data batches in `ray.data.map_batches` used in the + `OfflineData` class to transform data from columns to batches that can + be used in the `Learner.update...()` methods. Override the + `OfflinePreLearner` class and pass your derived class in here, if you + need to make some further transformations specific for your data or + loss. The default is None which uses the base `OfflinePreLearner` + defined in `ray.rllib.offline.offline_prelearner`. + prelearner_buffer_class: An optional `EpisodeReplayBuffer` class that RLlib + uses to buffer experiences when data is in `EpisodeType` or + RLlib's previous `SampleBatch` type format. In this case, a single + data row may contain multiple timesteps and the buffer serves two + purposes: (a) to store intermediate data in memory, and (b) to ensure + that RLlib samples exactly `train_batch_size_per_learner` experiences + per batch. The default is RLlib's `EpisodeReplayBuffer`. + prelearner_buffer_kwargs: Optional keyword arguments for initializing the + `EpisodeReplayBuffer`. In most cases this value is simply the `capacity` + for the default buffer that RLlib uses (`EpisodeReplayBuffer`), but it + may differ if the `prelearner_buffer_class` uses a custom buffer. + prelearner_module_synch_period: The period (number of batches converted) + after which the `RLModule` held by the `PreLearner` should sync weights. + The `PreLearner` is used to preprocess batches for the learners. The + higher this value, the more off-policy the `PreLearner`'s module is. + Values too small force the `PreLearner` to sync more frequently + and thus might slow down the data pipeline. The default value chosen + by the `OfflinePreLearner` is 10. + dataset_num_iters_per_learner: Number of updates to run in each learner + during a single training iteration. If None, each learner runs a + complete epoch over its data block (the dataset is partitioned into + at least as many blocks as there are learners). The default is `None`. + This value must be set to `1`, if RLlib uses a single (local) learner. + input_config: Arguments that describe the settings for reading the input. + If input is "sample", this is the environment configuration, e.g. + `env_name` and `env_config`, etc. See `EnvContext` for more info. + If the input is "dataset", this contains e.g. `format`, `path`. + actions_in_input_normalized: True, if the actions in a given offline "input" + are already normalized (between -1.0 and 1.0). This is usually the case + when the offline file has been generated by another RLlib algorithm + (e.g. PPO or SAC), while "normalize_actions" was set to True. + postprocess_inputs: Whether to run postprocess_trajectory() on the + trajectory fragments from offline inputs. Note that postprocessing is + done using the *current* policy, not the *behavior* policy, which + is typically undesirable for on-policy algorithms. + shuffle_buffer_size: If positive, input batches are shuffled via a + sliding window buffer of this number of batches. Use this if the input + data is not in random enough order. Input is delayed until the shuffle + buffer is filled. + output: Specify where experiences should be saved: + - None: don't save any experiences + - "logdir" to save to the agent log dir + - a path/URI to save to a custom output directory (e.g., "s3://bckt/") + - a function that returns a rllib.offline.OutputWriter + output_config: Arguments accessible from the IOContext for configuring + custom output. + output_compress_columns: What sample batch columns to LZ4 compress in the + output data. Note that providing `rllib.core.columns.Columns.OBS` also + compresses `rllib.core.columns.Columns.NEXT_OBS`. + output_max_file_size: Max output file size (in bytes) before rolling over + to a new file. + output_max_rows_per_file: Max output row numbers before rolling over to a + new file. + output_write_remaining_data: Determines whether any remaining data in the + recording buffers should be stored to disk. It is only applicable if + `output_max_rows_per_file` is defined. When sampling data, it is + buffered until the threshold specified by `output_max_rows_per_file` + is reached. Only complete multiples of `output_max_rows_per_file` are + written to disk, while any leftover data remains in the buffers. If a + recording session is stopped, residual data may still reside in these + buffers. Setting `output_write_remaining_data` to `True` ensures this + data is flushed to disk. By default, this attribute is set to `False`. + output_write_method: Write method for the `ray.data.Dataset` to write the + offline data to `output`. The default is `read_parquet` for Parquet + files. See https://docs.ray.io/en/latest/data/api/input_output.html for + more info about available read methods in `ray.data`. + output_write_method_kwargs: `kwargs` for the `output_write_method`. These + are passed into the write method without checking. + output_filesystem: A cloud filesystem to handle access to cloud storage when + writing experiences. Should be either "gcs" for Google Cloud Storage, + "s3" for AWS S3 buckets, or "abs" for Azure Blob Storage. + output_filesystem_kwargs: A dictionary holding the kwargs for the filesystem + given by `output_filesystem`. See `gcsfs.GCSFilesystem` for GCS, + `pyarrow.fs.S3FileSystem`, for S3, and `ablfs.AzureBlobFilesystem` for + ABS filesystem arguments. + output_write_episodes: If RLlib should record data in its RLlib's + `EpisodeType` format (that is, `SingleAgentEpisode` objects). Use this + format, if you need RLlib to order data in time and directly group by + episodes for example to train stateful modules or if you plan to use + recordings exclusively in RLlib. Otherwise RLlib records data in tabular + (columnar) format. Default is `True`. + offline_sampling: Whether sampling for the Algorithm happens via + reading from offline data. If True, EnvRunners don't limit the number + of collected batches within the same `sample()` call based on + the number of sub-environments within the worker (no sub-environments + present). + + Returns: + This updated AlgorithmConfig object. + """ + if input_ is not NotProvided: + self.input_ = input_ + if offline_data_class is not NotProvided: + self.offline_data_class = offline_data_class + if input_read_method is not NotProvided: + self.input_read_method = input_read_method + if input_read_method_kwargs is not NotProvided: + self.input_read_method_kwargs = input_read_method_kwargs + if input_read_schema is not NotProvided: + self.input_read_schema = input_read_schema + if input_read_episodes is not NotProvided: + self.input_read_episodes = input_read_episodes + if input_read_sample_batches is not NotProvided: + self.input_read_sample_batches = input_read_sample_batches + if input_read_batch_size is not NotProvided: + self.input_read_batch_size = input_read_batch_size + if input_filesystem is not NotProvided: + self.input_filesystem = input_filesystem + if input_filesystem_kwargs is not NotProvided: + self.input_filesystem_kwargs = input_filesystem_kwargs + if input_compress_columns is not NotProvided: + self.input_compress_columns = input_compress_columns + if materialize_data is not NotProvided: + self.materialize_data = materialize_data + if materialize_mapped_data is not NotProvided: + self.materialize_mapped_data = materialize_mapped_data + if map_batches_kwargs is not NotProvided: + self.map_batches_kwargs = map_batches_kwargs + if iter_batches_kwargs is not NotProvided: + self.iter_batches_kwargs = iter_batches_kwargs + if ignore_final_observation is not NotProvided: + self.ignore_final_observation = ignore_final_observation + if prelearner_class is not NotProvided: + self.prelearner_class = prelearner_class + if prelearner_buffer_class is not NotProvided: + self.prelearner_buffer_class = prelearner_buffer_class + if prelearner_buffer_kwargs is not NotProvided: + self.prelearner_buffer_kwargs = prelearner_buffer_kwargs + if prelearner_module_synch_period is not NotProvided: + self.prelearner_module_synch_period = prelearner_module_synch_period + if dataset_num_iters_per_learner is not NotProvided: + self.dataset_num_iters_per_learner = dataset_num_iters_per_learner + if input_config is not NotProvided: + if not isinstance(input_config, dict): + raise ValueError( + f"input_config must be a dict, got {type(input_config)}." + ) + # TODO (Kourosh) Once we use a complete separation between rollout worker + # and input dataset reader we can remove this. + # For now Error out if user attempts to set these parameters. + msg = "{} should not be set in the input_config. RLlib uses {} instead." + if input_config.get("num_cpus_per_read_task") is not None: + raise ValueError( + msg.format( + "num_cpus_per_read_task", + "config.env_runners(num_cpus_per_env_runner=..)", + ) + ) + if input_config.get("parallelism") is not None: + if self.in_evaluation: + raise ValueError( + msg.format( + "parallelism", + "config.evaluation(evaluation_num_env_runners=..)", + ) + ) + else: + raise ValueError( + msg.format( + "parallelism", "config.env_runners(num_env_runners=..)" + ) + ) + self.input_config = input_config + if actions_in_input_normalized is not NotProvided: + self.actions_in_input_normalized = actions_in_input_normalized + if postprocess_inputs is not NotProvided: + self.postprocess_inputs = postprocess_inputs + if shuffle_buffer_size is not NotProvided: + self.shuffle_buffer_size = shuffle_buffer_size + # TODO (simon): Enable storing to general log-directory. + if output is not NotProvided: + self.output = output + if output_config is not NotProvided: + self.output_config = output_config + if output_compress_columns is not NotProvided: + self.output_compress_columns = output_compress_columns + if output_max_file_size is not NotProvided: + self.output_max_file_size = output_max_file_size + if output_max_rows_per_file is not NotProvided: + self.output_max_rows_per_file = output_max_rows_per_file + if output_write_remaining_data is not NotProvided: + self.output_write_remaining_data = output_write_remaining_data + if output_write_method is not NotProvided: + self.output_write_method = output_write_method + if output_write_method_kwargs is not NotProvided: + self.output_write_method_kwargs = output_write_method_kwargs + if output_filesystem is not NotProvided: + self.output_filesystem = output_filesystem + if output_filesystem_kwargs is not NotProvided: + self.output_filesystem_kwargs = output_filesystem_kwargs + if output_write_episodes is not NotProvided: + self.output_write_episodes = output_write_episodes + if offline_sampling is not NotProvided: + self.offline_sampling = offline_sampling + + return self + + def multi_agent( + self, + *, + policies: Optional[ + Union[MultiAgentPolicyConfigDict, Collection[PolicyID]] + ] = NotProvided, + policy_map_capacity: Optional[int] = NotProvided, + policy_mapping_fn: Optional[ + Callable[[AgentID, "EpisodeType"], PolicyID] + ] = NotProvided, + policies_to_train: Optional[ + Union[Collection[PolicyID], Callable[[PolicyID, SampleBatchType], bool]] + ] = NotProvided, + policy_states_are_swappable: Optional[bool] = NotProvided, + observation_fn: Optional[Callable] = NotProvided, + count_steps_by: Optional[str] = NotProvided, + # Deprecated args: + algorithm_config_overrides_per_module=DEPRECATED_VALUE, + replay_mode=DEPRECATED_VALUE, + # Now done via Ray object store, which has its own cloud-supported + # spillover mechanism. + policy_map_cache=DEPRECATED_VALUE, + ) -> Self: + """Sets the config's multi-agent settings. + + Validates the new multi-agent settings and translates everything into + a unified multi-agent setup format. For example a `policies` list or set + of IDs is properly converted into a dict mapping these IDs to PolicySpecs. + + Args: + policies: Map of type MultiAgentPolicyConfigDict from policy ids to either + 4-tuples of (policy_cls, obs_space, act_space, config) or PolicySpecs. + These tuples or PolicySpecs define the class of the policy, the + observation- and action spaces of the policies, and any extra config. + policy_map_capacity: Keep this many policies in the "policy_map" (before + writing least-recently used ones to disk/S3). + policy_mapping_fn: Function mapping agent ids to policy ids. The signature + is: `(agent_id, episode, worker, **kwargs) -> PolicyID`. + policies_to_train: Determines those policies that should be updated. + Options are: + - None, for training all policies. + - An iterable of PolicyIDs that should be trained. + - A callable, taking a PolicyID and a SampleBatch or MultiAgentBatch + and returning a bool (indicating whether the given policy is trainable + or not, given the particular batch). This allows you to have a policy + trained only on certain data (e.g. when playing against a certain + opponent). + policy_states_are_swappable: Whether all Policy objects in this map can be + "swapped out" via a simple `state = A.get_state(); B.set_state(state)`, + where `A` and `B` are policy instances in this map. You should set + this to True for significantly speeding up the PolicyMap's cache lookup + times, iff your policies all share the same neural network + architecture and optimizer types. If True, the PolicyMap doesn't + have to garbage collect old, least recently used policies, but instead + keeps them in memory and simply override their state with the state of + the most recently accessed one. + For example, in a league-based training setup, you might have 100s of + the same policies in your map (playing against each other in various + combinations), but all of them share the same state structure + (are "swappable"). + observation_fn: Optional function that can be used to enhance the local + agent observations to include more state. See + rllib/evaluation/observation_function.py for more info. + count_steps_by: Which metric to use as the "batch size" when building a + MultiAgentBatch. The two supported values are: + "env_steps": Count each time the env is "stepped" (no matter how many + multi-agent actions are passed/how many multi-agent observations + have been returned in the previous step). + "agent_steps": Count each individual agent step as one step. + + Returns: + This updated AlgorithmConfig object. + """ + if policies is not NotProvided: + # Make sure our Policy IDs are ok (this should work whether `policies` + # is a dict or just any Sequence). + for pid in policies: + validate_module_id(pid, error=True) + + # Collection: Convert to dict. + if isinstance(policies, (set, tuple, list)): + policies = {p: PolicySpec() for p in policies} + # Validate each policy spec in a given dict. + if isinstance(policies, dict): + for pid, spec in policies.items(): + # If not a PolicySpec object, values must be lists/tuples of len 4. + if not isinstance(spec, PolicySpec): + if not isinstance(spec, (list, tuple)) or len(spec) != 4: + raise ValueError( + "Policy specs must be tuples/lists of " + "(cls or None, obs_space, action_space, config), " + f"got {spec} for PolicyID={pid}" + ) + # TODO: Switch from dict to AlgorithmConfigOverride, once available. + # Config not a dict. + elif ( + not isinstance(spec.config, (AlgorithmConfig, dict)) + and spec.config is not None + ): + raise ValueError( + f"Multi-agent policy config for {pid} must be a dict or " + f"AlgorithmConfig object, but got {type(spec.config)}!" + ) + self.policies = policies + else: + raise ValueError( + "`policies` must be dict mapping PolicyID to PolicySpec OR a " + "set/tuple/list of PolicyIDs!" + ) + + if algorithm_config_overrides_per_module != DEPRECATED_VALUE: + deprecation_warning(old="", error=False) + self.rl_module( + algorithm_config_overrides_per_module=( + algorithm_config_overrides_per_module + ) + ) + + if policy_map_capacity is not NotProvided: + self.policy_map_capacity = policy_map_capacity + + if policy_mapping_fn is not NotProvided: + # Create `policy_mapping_fn` from a config dict. + # Helpful if users would like to specify custom callable classes in + # yaml files. + if isinstance(policy_mapping_fn, dict): + policy_mapping_fn = from_config(policy_mapping_fn) + self.policy_mapping_fn = policy_mapping_fn + + if observation_fn is not NotProvided: + self.observation_fn = observation_fn + + if policy_map_cache != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.multi_agent(policy_map_cache=..)", + error=True, + ) + + if replay_mode != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.multi_agent(replay_mode=..)", + new="AlgorithmConfig.training(" + "replay_buffer_config={'replay_mode': ..})", + error=True, + ) + + if count_steps_by is not NotProvided: + if count_steps_by not in ["env_steps", "agent_steps"]: + raise ValueError( + "config.multi_agent(count_steps_by=..) must be one of " + f"[env_steps|agent_steps], not {count_steps_by}!" + ) + self.count_steps_by = count_steps_by + + if policies_to_train is not NotProvided: + assert ( + isinstance(policies_to_train, (list, set, tuple)) + or callable(policies_to_train) + or policies_to_train is None + ), ( + "ERROR: `policies_to_train` must be a [list|set|tuple] or a " + "callable taking PolicyID and SampleBatch and returning " + "True|False (trainable or not?) or None (for always training all " + "policies)." + ) + # Check `policies_to_train` for invalid entries. + if isinstance(policies_to_train, (list, set, tuple)): + if len(policies_to_train) == 0: + logger.warning( + "`config.multi_agent(policies_to_train=..)` is empty! " + "Make sure - if you would like to learn at least one policy - " + "to add its ID to that list." + ) + self.policies_to_train = policies_to_train + + if policy_states_are_swappable is not NotProvided: + self.policy_states_are_swappable = policy_states_are_swappable + + return self + + def reporting( + self, + *, + keep_per_episode_custom_metrics: Optional[bool] = NotProvided, + metrics_episode_collection_timeout_s: Optional[float] = NotProvided, + metrics_num_episodes_for_smoothing: Optional[int] = NotProvided, + min_time_s_per_iteration: Optional[float] = NotProvided, + min_train_timesteps_per_iteration: Optional[int] = NotProvided, + min_sample_timesteps_per_iteration: Optional[int] = NotProvided, + log_gradients: Optional[bool] = NotProvided, + ) -> Self: + """Sets the config's reporting settings. + + Args: + keep_per_episode_custom_metrics: Store raw custom metrics without + calculating max, min, mean + metrics_episode_collection_timeout_s: Wait for metric batches for at most + this many seconds. Those that have not returned in time are collected + in the next train iteration. + metrics_num_episodes_for_smoothing: Smooth rollout metrics over this many + episodes, if possible. + In case rollouts (sample collection) just started, there may be fewer + than this many episodes in the buffer and we'll compute metrics + over this smaller number of available episodes. + In case there are more than this many episodes collected in a single + training iteration, use all of these episodes for metrics computation, + meaning don't ever cut any "excess" episodes. + Set this to 1 to disable smoothing and to always report only the most + recently collected episode's return. + min_time_s_per_iteration: Minimum time (in sec) to accumulate within a + single `Algorithm.train()` call. This value does not affect learning, + only the number of times `Algorithm.training_step()` is called by + `Algorithm.train()`. If - after one such step attempt, the time taken + has not reached `min_time_s_per_iteration`, performs n more + `Algorithm.training_step()` calls until the minimum time has been + consumed. Set to 0 or None for no minimum time. + min_train_timesteps_per_iteration: Minimum training timesteps to accumulate + within a single `train()` call. This value does not affect learning, + only the number of times `Algorithm.training_step()` is called by + `Algorithm.train()`. If - after one such step attempt, the training + timestep count has not been reached, performs n more + `training_step()` calls until the minimum timesteps have been + executed. Set to 0 or None for no minimum timesteps. + min_sample_timesteps_per_iteration: Minimum env sampling timesteps to + accumulate within a single `train()` call. This value does not affect + learning, only the number of times `Algorithm.training_step()` is + called by `Algorithm.train()`. If - after one such step attempt, the env + sampling timestep count has not been reached, performs n more + `training_step()` calls until the minimum timesteps have been + executed. Set to 0 or None for no minimum timesteps. + log_gradients: Log gradients to results. If this is `True` the global norm + of the gradients dictionariy for each optimizer is logged to results. + The default is `False`. + + Returns: + This updated AlgorithmConfig object. + """ + if keep_per_episode_custom_metrics is not NotProvided: + self.keep_per_episode_custom_metrics = keep_per_episode_custom_metrics + if metrics_episode_collection_timeout_s is not NotProvided: + self.metrics_episode_collection_timeout_s = ( + metrics_episode_collection_timeout_s + ) + if metrics_num_episodes_for_smoothing is not NotProvided: + self.metrics_num_episodes_for_smoothing = metrics_num_episodes_for_smoothing + if min_time_s_per_iteration is not NotProvided: + self.min_time_s_per_iteration = min_time_s_per_iteration + if min_train_timesteps_per_iteration is not NotProvided: + self.min_train_timesteps_per_iteration = min_train_timesteps_per_iteration + if min_sample_timesteps_per_iteration is not NotProvided: + self.min_sample_timesteps_per_iteration = min_sample_timesteps_per_iteration + if log_gradients is not NotProvided: + self.log_gradients = log_gradients + + return self + + def checkpointing( + self, + export_native_model_files: Optional[bool] = NotProvided, + checkpoint_trainable_policies_only: Optional[bool] = NotProvided, + ) -> Self: + """Sets the config's checkpointing settings. + + Args: + export_native_model_files: Whether an individual Policy- + or the Algorithm's checkpoints also contain (tf or torch) native + model files. These could be used to restore just the NN models + from these files w/o requiring RLlib. These files are generated + by calling the tf- or torch- built-in saving utility methods on + the actual models. + checkpoint_trainable_policies_only: Whether to only add Policies to the + Algorithm checkpoint (in sub-directory "policies/") that are trainable + according to the `is_trainable_policy` callable of the local worker. + + Returns: + This updated AlgorithmConfig object. + """ + + if export_native_model_files is not NotProvided: + self.export_native_model_files = export_native_model_files + if checkpoint_trainable_policies_only is not NotProvided: + self.checkpoint_trainable_policies_only = checkpoint_trainable_policies_only + + return self + + def debugging( + self, + *, + logger_creator: Optional[Callable[[], Logger]] = NotProvided, + logger_config: Optional[dict] = NotProvided, + log_level: Optional[str] = NotProvided, + log_sys_usage: Optional[bool] = NotProvided, + fake_sampler: Optional[bool] = NotProvided, + seed: Optional[int] = NotProvided, + ) -> Self: + """Sets the config's debugging settings. + + Args: + logger_creator: Callable that creates a ray.tune.Logger + object. If unspecified, a default logger is created. + logger_config: Define logger-specific configuration to be used inside Logger + Default value None allows overwriting with nested dicts. + log_level: Set the ray.rllib.* log level for the agent process and its + workers. Should be one of DEBUG, INFO, WARN, or ERROR. The DEBUG level + also periodically prints out summaries of relevant internal dataflow + (this is also printed out once at startup at the INFO level). + log_sys_usage: Log system resource metrics to results. This requires + `psutil` to be installed for sys stats, and `gputil` for GPU metrics. + fake_sampler: Use fake (infinite speed) sampler. For testing only. + seed: This argument, in conjunction with worker_index, sets the random + seed of each worker, so that identically configured trials have + identical results. This makes experiments reproducible. + + Returns: + This updated AlgorithmConfig object. + """ + if logger_creator is not NotProvided: + self.logger_creator = logger_creator + if logger_config is not NotProvided: + self.logger_config = logger_config + if log_level is not NotProvided: + self.log_level = log_level + if log_sys_usage is not NotProvided: + self.log_sys_usage = log_sys_usage + if fake_sampler is not NotProvided: + self.fake_sampler = fake_sampler + if seed is not NotProvided: + self.seed = seed + + return self + + def fault_tolerance( + self, + *, + restart_failed_env_runners: Optional[bool] = NotProvided, + ignore_env_runner_failures: Optional[bool] = NotProvided, + max_num_env_runner_restarts: Optional[int] = NotProvided, + delay_between_env_runner_restarts_s: Optional[float] = NotProvided, + restart_failed_sub_environments: Optional[bool] = NotProvided, + num_consecutive_env_runner_failures_tolerance: Optional[int] = NotProvided, + env_runner_health_probe_timeout_s: Optional[float] = NotProvided, + env_runner_restore_timeout_s: Optional[float] = NotProvided, + # Deprecated args. + recreate_failed_env_runners=DEPRECATED_VALUE, + ignore_worker_failures=DEPRECATED_VALUE, + recreate_failed_workers=DEPRECATED_VALUE, + max_num_worker_restarts=DEPRECATED_VALUE, + delay_between_worker_restarts_s=DEPRECATED_VALUE, + num_consecutive_worker_failures_tolerance=DEPRECATED_VALUE, + worker_health_probe_timeout_s=DEPRECATED_VALUE, + worker_restore_timeout_s=DEPRECATED_VALUE, + ) -> Self: + """Sets the config's fault tolerance settings. + + Args: + restart_failed_env_runners: Whether - upon an EnvRunner failure - RLlib + tries to restart the lost EnvRunner(s) as an identical copy of the + failed one(s). You should set this to True when training on SPOT + instances that may preempt any time. The new, recreated EnvRunner(s) + only differ from the failed one in their `self.recreated_worker=True` + property value and have the same `worker_index` as the original(s). + If this setting is True, the value of the `ignore_env_runner_failures` + setting is ignored. + ignore_env_runner_failures: Whether to ignore any EnvRunner failures + and continue running with the remaining EnvRunners. This setting is + ignored, if `restart_failed_env_runners=True`. + max_num_env_runner_restarts: The maximum number of times any EnvRunner + is allowed to be restarted (if `restart_failed_env_runners` is True). + delay_between_env_runner_restarts_s: The delay (in seconds) between two + consecutive EnvRunner restarts (if `restart_failed_env_runners` is + True). + restart_failed_sub_environments: If True and any sub-environment (within + a vectorized env) throws any error during env stepping, the + EnvRunner tries to restart the faulty sub-environment. This is done + without disturbing the other (still intact) sub-environment and without + the EnvRunner crashing. You can raise + `ray.rllib.env.env_runner.StepFailedRecreateEnvError` from your + environment's `step` method to not log the error. + num_consecutive_env_runner_failures_tolerance: The number of consecutive + times an EnvRunner failure (also for evaluation) is tolerated before + finally crashing the Algorithm. Only useful if either + `ignore_env_runner_failures` or `restart_failed_env_runners` is True. + Note that for `restart_failed_sub_environments` and sub-environment + failures, the EnvRunner itself is NOT affected and won't throw any + errors as the flawed sub-environment is silently restarted under the + hood. + env_runner_health_probe_timeout_s: Max amount of time in seconds, we should + spend waiting for EnvRunner health probe calls + (`EnvRunner.ping.remote()`) to respond. Health pings are very cheap, + however, we perform the health check via a blocking `ray.get()`, so the + default value should not be too large. + env_runner_restore_timeout_s: Max amount of time we should wait to restore + states on recovered EnvRunner actors. Default is 30 mins. + + Returns: + This updated AlgorithmConfig object. + """ + if recreate_failed_env_runners != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(recreate_failed_env_runners)", + new="AlgorithmConfig.fault_tolerance(restart_failed_env_runners)", + error=True, + ) + if ignore_worker_failures != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(ignore_worker_failures)", + new="AlgorithmConfig.fault_tolerance(ignore_env_runner_failures)", + error=True, + ) + if recreate_failed_workers != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(recreate_failed_workers)", + new="AlgorithmConfig.fault_tolerance(restart_failed_env_runners)", + error=True, + ) + if max_num_worker_restarts != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(max_num_worker_restarts)", + new="AlgorithmConfig.fault_tolerance(max_num_env_runner_restarts)", + error=True, + ) + if delay_between_worker_restarts_s != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(delay_between_worker_restarts_s)", + new="AlgorithmConfig.fault_tolerance(delay_between_env_runner_" + "restarts_s)", + error=True, + ) + if num_consecutive_worker_failures_tolerance != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(num_consecutive_worker_" + "failures_tolerance)", + new="AlgorithmConfig.fault_tolerance(num_consecutive_env_runner_" + "failures_tolerance)", + error=True, + ) + if worker_health_probe_timeout_s != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(worker_health_probe_timeout_s)", + new="AlgorithmConfig.fault_tolerance(" + "env_runner_health_probe_timeout_s)", + error=True, + ) + if worker_restore_timeout_s != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.fault_tolerance(worker_restore_timeout_s)", + new="AlgorithmConfig.fault_tolerance(env_runner_restore_timeout_s)", + error=True, + ) + + if ignore_env_runner_failures is not NotProvided: + self.ignore_env_runner_failures = ignore_env_runner_failures + if restart_failed_env_runners is not NotProvided: + self.restart_failed_env_runners = restart_failed_env_runners + if max_num_env_runner_restarts is not NotProvided: + self.max_num_env_runner_restarts = max_num_env_runner_restarts + if delay_between_env_runner_restarts_s is not NotProvided: + self.delay_between_env_runner_restarts_s = ( + delay_between_env_runner_restarts_s + ) + if restart_failed_sub_environments is not NotProvided: + self.restart_failed_sub_environments = restart_failed_sub_environments + if num_consecutive_env_runner_failures_tolerance is not NotProvided: + self.num_consecutive_env_runner_failures_tolerance = ( + num_consecutive_env_runner_failures_tolerance + ) + if env_runner_health_probe_timeout_s is not NotProvided: + self.env_runner_health_probe_timeout_s = env_runner_health_probe_timeout_s + if env_runner_restore_timeout_s is not NotProvided: + self.env_runner_restore_timeout_s = env_runner_restore_timeout_s + + return self + + def rl_module( + self, + *, + model_config: Optional[Union[Dict[str, Any], DefaultModelConfig]] = NotProvided, + rl_module_spec: Optional[RLModuleSpecType] = NotProvided, + algorithm_config_overrides_per_module: Optional[ + Dict[ModuleID, PartialAlgorithmConfigDict] + ] = NotProvided, + # Deprecated arg. + model_config_dict=DEPRECATED_VALUE, + _enable_rl_module_api=DEPRECATED_VALUE, + ) -> Self: + """Sets the config's RLModule settings. + + Args: + model_config: The DefaultModelConfig object (or a config dictionary) passed + as `model_config` arg into each RLModule's constructor. This is used + for all RLModules, if not otherwise specified through `rl_module_spec`. + rl_module_spec: The RLModule spec to use for this config. It can be either + a RLModuleSpec or a MultiRLModuleSpec. If the + observation_space, action_space, catalog_class, or the model config is + not specified it is inferred from the env and other parts of the + algorithm config object. + algorithm_config_overrides_per_module: Only used if + `enable_rl_module_and_learner=True`. + A mapping from ModuleIDs to per-module AlgorithmConfig override dicts, + which apply certain settings, + e.g. the learning rate, from the main AlgorithmConfig only to this + particular module (within a MultiRLModule). + You can create override dicts by using the `AlgorithmConfig.overrides` + utility. For example, to override your learning rate and (PPO) lambda + setting just for a single RLModule with your MultiRLModule, do: + config.multi_agent(algorithm_config_overrides_per_module={ + "module_1": PPOConfig.overrides(lr=0.0002, lambda_=0.75), + }) + + Returns: + This updated AlgorithmConfig object. + """ + if _enable_rl_module_api != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.rl_module(_enable_rl_module_api=..)", + new="AlgorithmConfig.api_stack(enable_rl_module_and_learner=..)", + error=True, + ) + if model_config_dict != DEPRECATED_VALUE: + deprecation_warning( + old="AlgorithmConfig.rl_module(model_config_dict=..)", + new="AlgorithmConfig.rl_module(model_config=..)", + error=False, + ) + model_config = model_config_dict + + if model_config is not NotProvided: + self._model_config = model_config + if rl_module_spec is not NotProvided: + self._rl_module_spec = rl_module_spec + if algorithm_config_overrides_per_module is not NotProvided: + if not isinstance(algorithm_config_overrides_per_module, dict): + raise ValueError( + "`algorithm_config_overrides_per_module` must be a dict mapping " + "module IDs to config override dicts! You provided " + f"{algorithm_config_overrides_per_module}." + ) + self.algorithm_config_overrides_per_module.update( + algorithm_config_overrides_per_module + ) + + return self + + def experimental( + self, + *, + _validate_config: Optional[bool] = True, + _use_msgpack_checkpoints: Optional[bool] = NotProvided, + _torch_grad_scaler_class: Optional[Type] = NotProvided, + _torch_lr_scheduler_classes: Optional[ + Union[List[Type], Dict[ModuleID, List[Type]]] + ] = NotProvided, + _tf_policy_handles_more_than_one_loss: Optional[bool] = NotProvided, + _disable_preprocessor_api: Optional[bool] = NotProvided, + _disable_action_flattening: Optional[bool] = NotProvided, + _disable_initialize_loss_from_dummy_batch: Optional[bool] = NotProvided, + ) -> Self: + """Sets the config's experimental settings. + + Args: + _validate_config: Whether to run `validate()` on this config. True by + default. If False, ignores any calls to `self.validate()`. + _use_msgpack_checkpoints: Create state files in all checkpoints through + msgpack rather than pickle. + _torch_grad_scaler_class: Class to use for torch loss scaling (and gradient + unscaling). The class must implement the following methods to be + compatible with a `TorchLearner`. These methods/APIs match exactly those + of torch's own `torch.amp.GradScaler` (see here for more details + https://pytorch.org/docs/stable/amp.html#gradient-scaling): + `scale([loss])` to scale the loss by some factor. + `get_scale()` to get the current scale factor value. + `step([optimizer])` to unscale the grads (divide by the scale factor) + and step the given optimizer. + `update()` to update the scaler after an optimizer step (for example to + adjust the scale factor). + _torch_lr_scheduler_classes: A list of `torch.lr_scheduler.LRScheduler` + (see here for more details + https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate) + classes or a dictionary mapping module IDs to such a list of respective + scheduler classes. Multiple scheduler classes can be applied in sequence + and are stepped in the same sequence as defined here. Note, most + learning rate schedulers need arguments to be configured, that is, you + might have to partially initialize the schedulers in the list(s) using + `functools.partial`. + _tf_policy_handles_more_than_one_loss: Experimental flag. + If True, TFPolicy handles more than one loss or optimizer. + Set this to True, if you would like to return more than + one loss term from your `loss_fn` and an equal number of optimizers + from your `optimizer_fn`. + _disable_preprocessor_api: Experimental flag. + If True, no (observation) preprocessor is created and + observations arrive in model as they are returned by the env. + _disable_action_flattening: Experimental flag. + If True, RLlib doesn't flatten the policy-computed actions into + a single tensor (for storage in SampleCollectors/output files/etc..), + but leave (possibly nested) actions as-is. Disabling flattening affects: + - SampleCollectors: Have to store possibly nested action structs. + - Models that have the previous action(s) as part of their input. + - Algorithms reading from offline files (incl. action information). + + Returns: + This updated AlgorithmConfig object. + """ + if _validate_config is not NotProvided: + self._validate_config = _validate_config + if _use_msgpack_checkpoints is not NotProvided: + self._use_msgpack_checkpoints = _use_msgpack_checkpoints + if _tf_policy_handles_more_than_one_loss is not NotProvided: + self._tf_policy_handles_more_than_one_loss = ( + _tf_policy_handles_more_than_one_loss + ) + if _disable_preprocessor_api is not NotProvided: + self._disable_preprocessor_api = _disable_preprocessor_api + if _disable_action_flattening is not NotProvided: + self._disable_action_flattening = _disable_action_flattening + if _disable_initialize_loss_from_dummy_batch is not NotProvided: + self._disable_initialize_loss_from_dummy_batch = ( + _disable_initialize_loss_from_dummy_batch + ) + if _torch_grad_scaler_class is not NotProvided: + self._torch_grad_scaler_class = _torch_grad_scaler_class + if _torch_lr_scheduler_classes is not NotProvided: + self._torch_lr_scheduler_classes = _torch_lr_scheduler_classes + + return self + + @property + def is_atari(self) -> bool: + """True if if specified env is an Atari env.""" + + # Not yet determined, try to figure this out. + if self._is_atari is None: + # Atari envs are usually specified via a string like "PongNoFrameskip-v4" + # or "ale_py:ALE/Breakout-v5". + # We do NOT attempt to auto-detect Atari env for other specified types like + # a callable, to avoid running heavy logics in validate(). + # For these cases, users can explicitly set `environment(atari=True)`. + if type(self.env) is not str: + return False + try: + env = gym.make(self.env) + # Any gymnasium error -> Cannot be an Atari env. + except gym.error.Error: + return False + + self._is_atari = is_atari(env) + # Clean up env's resources, if any. + env.close() + + return self._is_atari + + @property + def is_multi_agent(self) -> bool: + """Returns whether this config specifies a multi-agent setup. + + Returns: + True, if a) >1 policies defined OR b) 1 policy defined, but its ID is NOT + DEFAULT_POLICY_ID. + """ + return len(self.policies) > 1 or DEFAULT_POLICY_ID not in self.policies + + @property + def learner_class(self) -> Type["Learner"]: + """Returns the Learner sub-class to use by this Algorithm. + + Either + a) User sets a specific learner class via calling `.training(learner_class=...)` + b) User leaves learner class unset (None) and the AlgorithmConfig itself + figures out the actual learner class by calling its own + `.get_default_learner_class()` method. + """ + return self._learner_class or self.get_default_learner_class() + + @property + def model_config(self): + """Defines the model configuration used. + + This method combines the auto configuration `self _model_config_auto_includes` + defined by an algorithm with the user-defined configuration in + `self._model_config`.This configuration dictionary is used to + configure the `RLModule` in the new stack and the `ModelV2` in the old + stack. + + Returns: + A dictionary with the model configuration. + """ + return self._model_config_auto_includes | ( + self._model_config + if isinstance(self._model_config, dict) + else dataclasses.asdict(self._model_config) + ) + + @property + def rl_module_spec(self): + default_rl_module_spec = self.get_default_rl_module_spec() + _check_rl_module_spec(default_rl_module_spec) + + # `self._rl_module_spec` has been user defined (via call to `self.rl_module()`). + if self._rl_module_spec is not None: + # Merge provided RL Module spec class with defaults. + _check_rl_module_spec(self._rl_module_spec) + # Merge given spec with default one (in case items are missing, such as + # spaces, module class, etc.) + if isinstance(self._rl_module_spec, RLModuleSpec): + if isinstance(default_rl_module_spec, RLModuleSpec): + default_rl_module_spec.update(self._rl_module_spec) + return default_rl_module_spec + elif isinstance(default_rl_module_spec, MultiRLModuleSpec): + raise ValueError( + "Cannot merge MultiRLModuleSpec with RLModuleSpec!" + ) + else: + multi_rl_module_spec = copy.deepcopy(self._rl_module_spec) + multi_rl_module_spec.update(default_rl_module_spec) + return multi_rl_module_spec + + # `self._rl_module_spec` has not been user defined -> return default one. + else: + return default_rl_module_spec + + @property + def train_batch_size_per_learner(self) -> int: + # If not set explicitly, try to infer the value. + if self._train_batch_size_per_learner is None: + return self.train_batch_size // (self.num_learners or 1) + return self._train_batch_size_per_learner + + @train_batch_size_per_learner.setter + def train_batch_size_per_learner(self, value: int) -> None: + self._train_batch_size_per_learner = value + + @property + def total_train_batch_size(self) -> int: + """Returns the effective total train batch size. + + New API stack: `train_batch_size_per_learner` * [effective num Learners]. + + @OldAPIStack: User never touches `train_batch_size_per_learner` or + `num_learners`) -> `train_batch_size`. + """ + return self.train_batch_size_per_learner * (self.num_learners or 1) + + # TODO: Make rollout_fragment_length as read-only property and replace the current + # self.rollout_fragment_length a private variable. + def get_rollout_fragment_length(self, worker_index: int = 0) -> int: + """Automatically infers a proper rollout_fragment_length setting if "auto". + + Uses the simple formula: + `rollout_fragment_length` = `total_train_batch_size` / + (`num_envs_per_env_runner` * `num_env_runners`) + + If result is a fraction AND `worker_index` is provided, makes + those workers add additional timesteps, such that the overall batch size (across + the workers) adds up to exactly the `total_train_batch_size`. + + Returns: + The user-provided `rollout_fragment_length` or a computed one (if user + provided value is "auto"), making sure `total_train_batch_size` is reached + exactly in each iteration. + """ + if self.rollout_fragment_length == "auto": + # Example: + # 2 workers, 2 envs per worker, 2000 train batch size: + # -> 2000 / 4 -> 500 + # 4 workers, 3 envs per worker, 2500 train batch size: + # -> 2500 / 12 -> 208.333 -> diff=4 (208 * 12 = 2496) + # -> worker 1, 2: 209, workers 3, 4: 208 + # 2 workers, 20 envs per worker, 512 train batch size: + # -> 512 / 40 -> 12.8 -> diff=32 (12 * 40 = 480) + # -> worker 1: 13, workers 2: 12 + rollout_fragment_length = self.total_train_batch_size / ( + self.num_envs_per_env_runner * (self.num_env_runners or 1) + ) + if int(rollout_fragment_length) != rollout_fragment_length: + diff = self.total_train_batch_size - int( + rollout_fragment_length + ) * self.num_envs_per_env_runner * (self.num_env_runners or 1) + if ((worker_index - 1) * self.num_envs_per_env_runner) >= diff: + return int(rollout_fragment_length) + else: + return int(rollout_fragment_length) + 1 + return int(rollout_fragment_length) + else: + return self.rollout_fragment_length + + # TODO: Make evaluation_config as read-only property and replace the current + # self.evaluation_config a private variable. + def get_evaluation_config_object( + self, + ) -> Optional["AlgorithmConfig"]: + """Creates a full AlgorithmConfig object from `self.evaluation_config`. + + Returns: + A fully valid AlgorithmConfig object that can be used for the evaluation + EnvRunnerGroup. If `self` is already an evaluation config object, return + None. + """ + if self.in_evaluation: + assert self.evaluation_config is None + return None + + evaluation_config = self.evaluation_config + # Already an AlgorithmConfig -> copy and use as-is. + if isinstance(evaluation_config, AlgorithmConfig): + eval_config_obj = evaluation_config.copy(copy_frozen=False) + # Create unfrozen copy of self to be used as the to-be-returned eval + # AlgorithmConfig. + else: + eval_config_obj = self.copy(copy_frozen=False) + # Update with evaluation override settings: + eval_config_obj.update_from_dict(evaluation_config or {}) + + # Switch on the `in_evaluation` flag and remove `evaluation_config` + # (set to None). + eval_config_obj.in_evaluation = True + eval_config_obj.evaluation_config = None + + # Force-set the `num_env_runners` setting to `self.evaluation_num_env_runners`. + # Actually, the `self.evaluation_num_env_runners` is merely a convenience + # attribute and might be set instead through: + # `config.evaluation(evaluation_config={"num_env_runners": ...})` + eval_config_obj.num_env_runners = self.evaluation_num_env_runners + + # NOTE: The following if-block is only relevant for the old API stack. + # For the new API stack (EnvRunners), the evaluation methods of Algorithm + # explicitly tell each EnvRunner on each sample call, how many timesteps + # of episodes to collect. + # Evaluation duration unit: episodes. + # Switch on `complete_episode` rollouts. Also, make sure + # rollout fragments are short so we never have more than one + # episode in one rollout. + if self.evaluation_duration_unit == "episodes": + eval_config_obj.batch_mode = "complete_episodes" + eval_config_obj.rollout_fragment_length = 1 + # Evaluation duration unit: timesteps. + # - Set `batch_mode=truncate_episodes` so we don't perform rollouts + # strictly along episode borders. + # Set `rollout_fragment_length` such that desired steps are divided + # equally amongst workers or - in "auto" duration mode - set it + # to a reasonably small number (10), such that a single `sample()` + # call doesn't take too much time and we can stop evaluation as soon + # as possible after the train step is completed. + else: + eval_config_obj.batch_mode = "truncate_episodes" + eval_config_obj.rollout_fragment_length = ( + # Set to a moderately small (but not too small) value in order + # to a) not overshoot too much the parallelly running `training_step` + # but also to b) avoid too many `sample()` remote calls. + # 100 seems like a good middle ground. + 100 + if self.evaluation_duration == "auto" + else int( + math.ceil( + self.evaluation_duration + / (self.evaluation_num_env_runners or 1) + ) + ) + ) + + return eval_config_obj + + def validate_train_batch_size_vs_rollout_fragment_length(self) -> None: + """Detects mismatches for `train_batch_size` vs `rollout_fragment_length`. + + Only applicable for algorithms, whose train_batch_size should be directly + dependent on rollout_fragment_length (synchronous sampling, on-policy PG algos). + + If rollout_fragment_length != "auto", makes sure that the product of + `rollout_fragment_length` x `num_env_runners` x `num_envs_per_env_runner` + roughly (10%) matches the provided `train_batch_size`. Otherwise, errors with + asking the user to set rollout_fragment_length to `auto` or to a matching + value. + + Raises: + ValueError: If there is a mismatch between user provided + `rollout_fragment_length` and `total_train_batch_size`. + """ + if self.rollout_fragment_length != "auto" and not self.in_evaluation: + min_batch_size = ( + max(self.num_env_runners, 1) + * self.num_envs_per_env_runner + * self.rollout_fragment_length + ) + batch_size = min_batch_size + while batch_size < self.total_train_batch_size: + batch_size += min_batch_size + if batch_size - self.total_train_batch_size > ( + 0.1 * self.total_train_batch_size + ) or batch_size - min_batch_size - self.total_train_batch_size > ( + 0.1 * self.total_train_batch_size + ): + suggested_rollout_fragment_length = self.total_train_batch_size // ( + self.num_envs_per_env_runner * (self.num_env_runners or 1) + ) + self._value_error( + "Your desired `total_train_batch_size` " + f"({self.total_train_batch_size}={self.num_learners} " + f"learners x {self.train_batch_size_per_learner}) " + "or a value 10% off of that cannot be achieved with your other " + f"settings (num_env_runners={self.num_env_runners}; " + f"num_envs_per_env_runner={self.num_envs_per_env_runner}; " + f"rollout_fragment_length={self.rollout_fragment_length})! " + "Try setting `rollout_fragment_length` to 'auto' OR to a value of " + f"{suggested_rollout_fragment_length}." + ) + + def get_torch_compile_worker_config(self): + """Returns the TorchCompileConfig to use on workers.""" + + from ray.rllib.core.rl_module.torch.torch_compile_config import ( + TorchCompileConfig, + ) + + return TorchCompileConfig( + torch_dynamo_backend=self.torch_compile_worker_dynamo_backend, + torch_dynamo_mode=self.torch_compile_worker_dynamo_mode, + ) + + def get_default_rl_module_spec(self) -> RLModuleSpecType: + """Returns the RLModule spec to use for this algorithm. + + Override this method in the subclass to return the RLModule spec, given + the input framework. + + Returns: + The RLModuleSpec (or MultiRLModuleSpec) to + use for this algorithm's RLModule. + """ + raise NotImplementedError + + def get_default_learner_class(self) -> Union[Type["Learner"], str]: + """Returns the Learner class to use for this algorithm. + + Override this method in the sub-class to return the Learner class type given + the input framework. + + Returns: + The Learner class to use for this algorithm either as a class type or as + a string (e.g. "ray.rllib.algorithms.ppo.ppo_learner.PPOLearner"). + """ + raise NotImplementedError + + def get_rl_module_spec( + self, + env: Optional[EnvType] = None, + spaces: Optional[Dict[str, Tuple[gym.Space, gym.Space]]] = None, + inference_only: Optional[bool] = None, + ) -> RLModuleSpec: + """Returns the RLModuleSpec based on the given env/spaces and this config. + + Args: + env: An optional environment instance, from which to infer the observation- + and action spaces for the RLModule. If not provided, tries to infer + from `spaces`, otherwise from `self.observation_space` and + `self.action_space`. Raises an error, if no information on spaces can be + inferred. + spaces: Optional dict mapping ModuleIDs to 2-tuples of observation- and + action space that should be used for the respective RLModule. + These spaces are usually provided by an already instantiated remote + EnvRunner (call `EnvRunner.get_spaces()` to receive this dict). If not + provided, RLlib tries to infer this from `env`, if provided, otherwise + from `self.observation_space` and `self.action_space`. Raises an error, + if no information on spaces can be inferred. + inference_only: If `True`, the returned module spec is used in an + inference-only setting (sampling) and the RLModule can thus be built in + its light version (if available). For example, the `inference_only` + version of an RLModule might only contain the networks required for + computing actions, but misses additional target- or critic networks. + + Returns: + A new RLModuleSpec instance that can be used to build an RLModule. + """ + rl_module_spec = copy.deepcopy(self.rl_module_spec) + + # If a MultiRLModuleSpec -> Reduce to single-agent (and assert that + # all non DEFAULT_MODULE_IDs are `learner_only` (so they are not built on + # EnvRunner). + if isinstance(rl_module_spec, MultiRLModuleSpec): + error = False + if DEFAULT_MODULE_ID not in rl_module_spec: + error = True + if inference_only: + for mid, spec in rl_module_spec.rl_module_specs.items(): + if mid != DEFAULT_MODULE_ID: + if not spec.learner_only: + error = True + elif len(rl_module_spec) > 1: + error = True + if error: + raise ValueError( + "When calling `AlgorithmConfig.get_rl_module_spec()`, the " + "configuration must contain the `DEFAULT_MODULE_ID` key and all " + "other keys' specs must have the setting `learner_only=True`! If " + "you are using a more complex setup, call " + "`AlgorithmConfig.get_multi_rl_module_spec(...)` instead." + ) + rl_module_spec = rl_module_spec[DEFAULT_MODULE_ID] + + if rl_module_spec.observation_space is None: + if spaces is not None: + rl_module_spec.observation_space = spaces[DEFAULT_MODULE_ID][0] + elif env is not None and isinstance(env, gym.Env): + rl_module_spec.observation_space = getattr( + env, "single_observation_space", env.observation_space + ) + else: + rl_module_spec.observation_space = self.observation_space + + if rl_module_spec.action_space is None: + if spaces is not None: + rl_module_spec.action_space = spaces[DEFAULT_MODULE_ID][1] + elif env is not None and isinstance(env, gym.Env): + rl_module_spec.action_space = getattr( + env, "single_action_space", env.action_space + ) + else: + rl_module_spec.action_space = self.action_space + + # If module_config_dict is not defined, set to our generic one. + if rl_module_spec.model_config is None: + rl_module_spec.model_config = self.model_config + + if inference_only is not None: + rl_module_spec.inference_only = inference_only + + return rl_module_spec + + def get_multi_rl_module_spec( + self, + *, + env: Optional[EnvType] = None, + spaces: Optional[Dict[PolicyID, Tuple[gym.Space, gym.Space]]] = None, + inference_only: bool = False, + # @HybridAPIStack + policy_dict: Optional[Dict[str, PolicySpec]] = None, + single_agent_rl_module_spec: Optional[RLModuleSpec] = None, + ) -> MultiRLModuleSpec: + """Returns the MultiRLModuleSpec based on the given env/spaces. + + Args: + env: An optional environment instance, from which to infer the different + spaces for the individual RLModules. If not provided, tries to infer + from `spaces`, otherwise from `self.observation_space` and + `self.action_space`. Raises an error, if no information on spaces can be + inferred. + spaces: Optional dict mapping ModuleIDs to 2-tuples of observation- and + action space that should be used for the respective RLModule. + These spaces are usually provided by an already instantiated remote + EnvRunner (call `EnvRunner.get_spaces()`). If not provided, tries + to infer from `env`, otherwise from `self.observation_space` and + `self.action_space`. Raises an error, if no information on spaces can be + inferred. + inference_only: If `True`, the returned module spec is used in an + inference-only setting (sampling) and the RLModule can thus be built in + its light version (if available). For example, the `inference_only` + version of an RLModule might only contain the networks required for + computing actions, but misses additional target- or critic networks. + Also, if `True`, the returned spec does NOT contain those (sub) + RLModuleSpecs that have their `learner_only` flag set to True. + + Returns: + A new MultiRLModuleSpec instance that can be used to build a MultiRLModule. + """ + # TODO (Kourosh,sven): When we replace policy entirely there is no need for + # this function to map policy_dict to multi_rl_module_specs anymore. The module + # spec is directly given by the user or inferred from env and spaces. + if policy_dict is None: + policy_dict, _ = self.get_multi_agent_setup(env=env, spaces=spaces) + + # TODO (Kourosh): Raise an error if the config is not frozen + # If the module is single-agent convert it to multi-agent spec + + # The default RLModuleSpec (might be multi-agent or single-agent). + default_rl_module_spec = self.get_default_rl_module_spec() + # The currently configured RLModuleSpec (might be multi-agent or single-agent). + # If None, use the default one. + current_rl_module_spec = self._rl_module_spec or default_rl_module_spec + + # Algorithm is currently setup as a single-agent one. + if isinstance(current_rl_module_spec, RLModuleSpec): + # Use either the provided `single_agent_rl_module_spec` (a + # RLModuleSpec), the currently configured one of this + # AlgorithmConfig object, or the default one. + single_agent_rl_module_spec = ( + single_agent_rl_module_spec or current_rl_module_spec + ) + single_agent_rl_module_spec.inference_only = inference_only + # Now construct the proper MultiRLModuleSpec. + multi_rl_module_spec = MultiRLModuleSpec( + rl_module_specs={ + k: copy.deepcopy(single_agent_rl_module_spec) + for k in policy_dict.keys() + }, + ) + + # Algorithm is currently setup as a multi-agent one. + else: + # The user currently has a MultiAgentSpec setup (either via + # self._rl_module_spec or the default spec of this AlgorithmConfig). + assert isinstance(current_rl_module_spec, MultiRLModuleSpec) + + # Default is single-agent but the user has provided a multi-agent spec + # so the use-case is multi-agent. + if isinstance(default_rl_module_spec, RLModuleSpec): + # The individual (single-agent) module specs are defined by the user + # in the currently setup MultiRLModuleSpec -> Use that + # RLModuleSpec. + if isinstance(current_rl_module_spec.rl_module_specs, RLModuleSpec): + single_agent_spec = single_agent_rl_module_spec or ( + current_rl_module_spec.rl_module_specs + ) + single_agent_spec.inference_only = inference_only + module_specs = { + k: copy.deepcopy(single_agent_spec) for k in policy_dict.keys() + } + + # The individual (single-agent) module specs have not been configured + # via this AlgorithmConfig object -> Use provided single-agent spec or + # the default spec (which is also a RLModuleSpec in this + # case). + else: + single_agent_spec = ( + single_agent_rl_module_spec or default_rl_module_spec + ) + single_agent_spec.inference_only = inference_only + module_specs = { + k: copy.deepcopy( + current_rl_module_spec.rl_module_specs.get( + k, single_agent_spec + ) + ) + for k in ( + policy_dict | current_rl_module_spec.rl_module_specs + ).keys() + } + + # Now construct the proper MultiRLModuleSpec. + # We need to infer the multi-agent class from `current_rl_module_spec` + # and fill in the module_specs dict. + multi_rl_module_spec = current_rl_module_spec.__class__( + multi_rl_module_class=current_rl_module_spec.multi_rl_module_class, + rl_module_specs=module_specs, + modules_to_load=current_rl_module_spec.modules_to_load, + load_state_path=current_rl_module_spec.load_state_path, + model_config=current_rl_module_spec.model_config, + ) + + # Default is multi-agent and user wants to override it -> Don't use the + # default. + else: + # User provided an override RLModuleSpec -> Use this to + # construct the individual RLModules within the MultiRLModuleSpec. + if single_agent_rl_module_spec is not None: + pass + # User has NOT provided an override RLModuleSpec. + else: + # But the currently setup multi-agent spec has a SingleAgentRLModule + # spec defined -> Use that to construct the individual RLModules + # within the MultiRLModuleSpec. + if isinstance(current_rl_module_spec.rl_module_specs, RLModuleSpec): + # The individual module specs are not given, it is given as one + # RLModuleSpec to be re-used for all + single_agent_rl_module_spec = ( + current_rl_module_spec.rl_module_specs + ) + # The currently set up multi-agent spec has NO + # RLModuleSpec in it -> Error (there is no way we can + # infer this information from anywhere at this point). + else: + raise ValueError( + "We have a MultiRLModuleSpec " + f"({current_rl_module_spec}), but no " + "`RLModuleSpec`s to compile the individual " + "RLModules' specs! Use " + "`AlgorithmConfig.get_multi_rl_module_spec(" + "policy_dict=.., rl_module_spec=..)`." + ) + + single_agent_rl_module_spec.inference_only = inference_only + + # Now construct the proper MultiRLModuleSpec. + multi_rl_module_spec = current_rl_module_spec.__class__( + multi_rl_module_class=current_rl_module_spec.multi_rl_module_class, + rl_module_specs={ + k: copy.deepcopy(single_agent_rl_module_spec) + for k in policy_dict.keys() + }, + modules_to_load=current_rl_module_spec.modules_to_load, + load_state_path=current_rl_module_spec.load_state_path, + model_config=current_rl_module_spec.model_config, + ) + + # Fill in the missing values from the specs that we already have. By combining + # PolicySpecs and the default RLModuleSpec. + for module_id in policy_dict | multi_rl_module_spec.rl_module_specs: + + # Remove/skip `learner_only=True` RLModules if `inference_only` is True. + module_spec = multi_rl_module_spec.rl_module_specs[module_id] + if inference_only and module_spec.learner_only: + multi_rl_module_spec.remove_modules(module_id) + continue + + if module_spec.module_class is None: + if isinstance(default_rl_module_spec, RLModuleSpec): + module_spec.module_class = default_rl_module_spec.module_class + elif isinstance(default_rl_module_spec.rl_module_specs, RLModuleSpec): + module_class = default_rl_module_spec.rl_module_specs.module_class + # This should be already checked in validate() but we check it + # again here just in case + if module_class is None: + raise ValueError( + "The default rl_module spec cannot have an empty " + "module_class under its RLModuleSpec." + ) + module_spec.module_class = module_class + elif module_id in default_rl_module_spec.rl_module_specs: + module_spec.module_class = default_rl_module_spec.rl_module_specs[ + module_id + ].module_class + else: + raise ValueError( + f"Module class for module {module_id} cannot be inferred. " + f"It is neither provided in the rl_module_spec that " + "is passed in nor in the default module spec used in " + "the algorithm." + ) + if module_spec.catalog_class is None: + if isinstance(default_rl_module_spec, RLModuleSpec): + module_spec.catalog_class = default_rl_module_spec.catalog_class + elif isinstance(default_rl_module_spec.rl_module_specs, RLModuleSpec): + catalog_class = default_rl_module_spec.rl_module_specs.catalog_class + module_spec.catalog_class = catalog_class + elif module_id in default_rl_module_spec.rl_module_specs: + module_spec.catalog_class = default_rl_module_spec.rl_module_specs[ + module_id + ].catalog_class + else: + raise ValueError( + f"Catalog class for module {module_id} cannot be inferred. " + f"It is neither provided in the rl_module_spec that " + "is passed in nor in the default module spec used in " + "the algorithm." + ) + # TODO (sven): Find a good way to pack module specific parameters from + # the algorithms into the `model_config_dict`. + if ( + module_spec.observation_space is None + or module_spec.action_space is None + ): + policy_spec = policy_dict.get( + module_id, policy_dict.get(DEFAULT_MODULE_ID) + ) + if policy_spec is not None: + if module_spec.observation_space is None: + module_spec.observation_space = policy_spec.observation_space + if module_spec.action_space is None: + module_spec.action_space = policy_spec.action_space + # In case the `RLModuleSpec` does not have a model config dict, we use the + # the one defined by the auto keys and the `model_config_dict` arguments in + # `self.rl_module()`. + if module_spec.model_config is None: + module_spec.model_config = self.model_config + # Otherwise we combine the two dictionaries where settings from the + # `RLModuleSpec` have higher priority. + else: + module_spec.model_config = ( + self.model_config | module_spec._get_model_config() + ) + + return multi_rl_module_spec + + def __setattr__(self, key, value): + """Gatekeeper in case we are in frozen state and need to error.""" + + # If we are frozen, do not allow to set any attributes anymore. + if hasattr(self, "_is_frozen") and self._is_frozen: + # TODO: Remove `simple_optimizer` entirely. + # Remove need to set `worker_index` in RolloutWorker's c'tor. + if key not in ["simple_optimizer", "worker_index", "_is_frozen"]: + raise AttributeError( + f"Cannot set attribute ({key}) of an already frozen " + "AlgorithmConfig!" + ) + # Backward compatibility for checkpoints taken with wheels, in which + # `self.rl_module_spec` was still settable (now it's a property). + if key == "rl_module_spec": + key = "_rl_module_spec" + + super().__setattr__(key, value) + + def __getitem__(self, item): + """Shim method to still support accessing properties by key lookup. + + This way, an AlgorithmConfig object can still be used as if a dict, e.g. + by Ray Tune. + + Examples: + .. testcode:: + + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + config = AlgorithmConfig() + print(config["lr"]) + + .. testoutput:: + + 0.001 + """ + # TODO: Uncomment this once all algorithms use AlgorithmConfigs under the + # hood (as well as Ray Tune). + # if log_once("algo_config_getitem"): + # logger.warning( + # "AlgorithmConfig objects should NOT be used as dict! " + # f"Try accessing `{item}` directly as a property." + # ) + # In case user accesses "old" keys, e.g. "num_workers", which need to + # be translated to their correct property names. + item = self._translate_special_keys(item) + return getattr(self, item) + + def __setitem__(self, key, value): + # TODO: Remove comments once all methods/functions only support + # AlgorithmConfigs and there is no more ambiguity anywhere in the code + # on whether an AlgorithmConfig is used or an old python config dict. + # raise AttributeError( + # "AlgorithmConfig objects should not have their values set like dicts" + # f"(`config['{key}'] = {value}`), " + # f"but via setting their properties directly (config.{prop} = {value})." + # ) + if key == "multiagent": + raise AttributeError( + "Cannot set `multiagent` key in an AlgorithmConfig!\nTry setting " + "the multi-agent components of your AlgorithmConfig object via the " + "`multi_agent()` method and its arguments.\nE.g. `config.multi_agent(" + "policies=.., policy_mapping_fn.., policies_to_train=..)`." + ) + super().__setattr__(key, value) + + def __contains__(self, item) -> bool: + """Shim method to help pretend we are a dict.""" + prop = self._translate_special_keys(item, warn_deprecated=False) + return hasattr(self, prop) + + def get(self, key, default=None): + """Shim method to help pretend we are a dict.""" + prop = self._translate_special_keys(key, warn_deprecated=False) + return getattr(self, prop, default) + + def pop(self, key, default=None): + """Shim method to help pretend we are a dict.""" + return self.get(key, default) + + def keys(self): + """Shim method to help pretend we are a dict.""" + return self.to_dict().keys() + + def values(self): + """Shim method to help pretend we are a dict.""" + return self.to_dict().values() + + def items(self): + """Shim method to help pretend we are a dict.""" + return self.to_dict().items() + + @property + def _model_config_auto_includes(self) -> Dict[str, Any]: + """Defines which `AlgorithmConfig` settings/properties should be + auto-included into `self.model_config`. + + The dictionary in this property contains the default configuration of an + algorithm. Together with the `self._model`, this method is used to + define the configuration sent to the `RLModule`. + + Returns: + A dictionary with the automatically included properties/settings of this + `AlgorithmConfig` object into `self.model_config`. + """ + return {} + + # ----------------------------------------------------------- + # Various validation methods for different types of settings. + # ----------------------------------------------------------- + def _value_error(self, errmsg) -> None: + msg = errmsg + ( + "\nTo suppress all validation errors, set " + "`config.experimental(_validate_config=False)` at your own risk." + ) + if self._validate_config: + raise ValueError(msg) + else: + logger.warning(errmsg) + + def _validate_env_runner_settings(self) -> None: + allowed_vectorize_modes = set( + list(gym.VectorizeMode) + [mode.value for mode in gym.VectorizeMode] + ) + if self.gym_env_vectorize_mode not in allowed_vectorize_modes: + self._value_error( + f"`gym_env_vectorize_mode` ({self.gym_env_vectorize_mode}) " + "must be a member of `gymnasium.VectorizeMode`! " + f"Allowed values are {allowed_vectorize_modes}." + ) + + def _validate_callbacks_settings(self) -> None: + """Validates callbacks settings.""" + # Old API stack: + # - self.callbacks_cls must be a subclass of RLlibCallback. + # - All self.callbacks_... attributes must be None. + if not self.enable_env_runner_and_connector_v2: + if ( + self.callbacks_on_environment_created is not None + or self.callbacks_on_algorithm_init is not None + or self.callbacks_on_train_result is not None + or self.callbacks_on_evaluate_start is not None + or self.callbacks_on_evaluate_end is not None + or self.callbacks_on_sample_end is not None + or self.callbacks_on_environment_created is not None + or self.callbacks_on_episode_created is not None + or self.callbacks_on_episode_start is not None + or self.callbacks_on_episode_step is not None + or self.callbacks_on_episode_end is not None + or self.callbacks_on_checkpoint_loaded is not None + or self.callbacks_on_env_runners_recreated is not None + or self.callbacks_on_offline_eval_runners_recreated is not None + ): + self._value_error( + "Config settings `config.callbacks(on_....=lambda ..)` aren't " + "supported on the old API stack! Switch to the new API stack " + "through `config.api_stack(enable_env_runner_and_connector_v2=True," + " enable_rl_module_and_learner=True)`." + ) + + def _validate_framework_settings(self) -> None: + """Validates framework settings and checks whether framework is installed.""" + _tf1, _tf, _tfv = None, None, None + _torch = None + if self.framework_str not in {"tf", "tf2"} and self.framework_str != "torch": + return + elif self.framework_str in {"tf", "tf2"}: + _tf1, _tf, _tfv = try_import_tf() + else: + _torch, _ = try_import_torch() + + # Can not use "tf" with learner API. + if self.framework_str == "tf" and self.enable_rl_module_and_learner: + self._value_error( + "Cannot use `framework=tf` with the new API stack! Either switch to tf2" + " via `config.framework('tf2')` OR disable the new API stack via " + "`config.api_stack(enable_rl_module_and_learner=False)`." + ) + + # Check if torch framework supports torch.compile. + if ( + _torch is not None + and self.framework_str == "torch" + and version.parse(_torch.__version__) < TORCH_COMPILE_REQUIRED_VERSION + and (self.torch_compile_learner or self.torch_compile_worker) + ): + self._value_error("torch.compile is only supported from torch 2.0.0") + + # Make sure the Learner's torch-what-to-compile setting is supported. + if self.torch_compile_learner: + from ray.rllib.core.learner.torch.torch_learner import ( + TorchCompileWhatToCompile, + ) + + if self.torch_compile_learner_what_to_compile not in [ + TorchCompileWhatToCompile.FORWARD_TRAIN, + TorchCompileWhatToCompile.COMPLETE_UPDATE, + ]: + self._value_error( + f"`config.torch_compile_learner_what_to_compile` must be one of [" + f"TorchCompileWhatToCompile.forward_train, " + f"TorchCompileWhatToCompile.complete_update] but is" + f" {self.torch_compile_learner_what_to_compile}" + ) + + self._check_if_correct_nn_framework_installed(_tf1, _tf, _torch) + self._resolve_tf_settings(_tf1, _tfv) + + def _validate_resources_settings(self): + """Checks, whether resources related settings make sense.""" + pass + + def _validate_multi_agent_settings(self): + """Checks, whether multi-agent related settings make sense.""" + + # Check `policies_to_train` for invalid entries. + if isinstance(self.policies_to_train, (list, set, tuple)): + for pid in self.policies_to_train: + if pid not in self.policies: + self._value_error( + "`config.multi_agent(policies_to_train=..)` contains " + f"policy ID ({pid}) that was not defined in " + f"`config.multi_agent(policies=..)`!" + ) + + def _validate_evaluation_settings(self): + """Checks, whether evaluation related settings make sense.""" + + # Async evaluation has been deprecated. Use "simple" parallel mode instead + # (which is also async): + # `config.evaluation(evaluation_parallel_to_training=True)`. + if self.enable_async_evaluation is True: + self._value_error( + "`enable_async_evaluation` has been deprecated (you should set this to " + "False)! Use `config.evaluation(evaluation_parallel_to_training=True)` " + "instead." + ) + + # If `evaluation_num_env_runners` > 0, warn if `evaluation_interval` is 0 or + # None. + if self.evaluation_num_env_runners > 0 and not self.evaluation_interval: + logger.warning( + f"You have specified {self.evaluation_num_env_runners} " + "evaluation workers, but your `evaluation_interval` is 0 or None! " + "Therefore, evaluation doesn't occur automatically with each" + " call to `Algorithm.train()`. Instead, you have to call " + "`Algorithm.evaluate()` manually in order to trigger an " + "evaluation run." + ) + # If `evaluation_num_env_runners=0` and + # `evaluation_parallel_to_training=True`, warn that you need + # at least one remote eval worker for parallel training and + # evaluation, and set `evaluation_parallel_to_training` to False. + if ( + self.evaluation_num_env_runners == 0 + and self.num_offline_eval_runners == 0 + and self.evaluation_parallel_to_training + ): + self._value_error( + "`evaluation_parallel_to_training` can only be done if " + "`evaluation_num_env_runners` > 0! Try setting " + "`config.evaluation_parallel_to_training` to False." + ) + + # If `evaluation_duration=auto`, error if + # `evaluation_parallel_to_training=False`. + if self.evaluation_duration == "auto": + if not self.evaluation_parallel_to_training: + self._value_error( + "`evaluation_duration=auto` not supported for " + "`evaluation_parallel_to_training=False`!" + ) + elif self.evaluation_duration_unit == "episodes": + logger.warning( + "When using `config.evaluation_duration='auto'`, the sampling unit " + "used is always 'timesteps'! You have set " + "`config.evaluation_duration_unit='episodes'`, which is ignored." + ) + + # Make sure, `evaluation_duration` is an int otherwise. + elif ( + not isinstance(self.evaluation_duration, int) + or self.evaluation_duration <= 0 + ): + self._value_error( + f"`evaluation_duration` ({self.evaluation_duration}) must be an " + f"int and >0!" + ) + + def _validate_input_settings(self): + """Checks, whether input related settings make sense.""" + + if self.input_ == "sampler" and self.off_policy_estimation_methods: + self._value_error( + "Off-policy estimation methods can only be used if the input is a " + "dataset. We currently do not support applying off_policy_estimation_" + "method on a sampler input." + ) + + if self.input_ == "dataset": + # If you need to read a Ray dataset set the parallelism and + # num_cpus_per_read_task from rollout worker settings + self.input_config["num_cpus_per_read_task"] = self.num_cpus_per_env_runner + if self.in_evaluation: + # If using dataset for evaluation, the parallelism gets set to + # evaluation_num_env_runners for backward compatibility and num_cpus + # gets set to num_cpus_per_env_runner from rollout worker. User only + # needs to set evaluation_num_env_runners. + self.input_config["parallelism"] = self.evaluation_num_env_runners or 1 + else: + # If using dataset for training, the parallelism and num_cpus gets set + # based on rollout worker parameters. This is for backwards + # compatibility for now. User only needs to set num_env_runners. + self.input_config["parallelism"] = self.num_env_runners or 1 + + def _validate_new_api_stack_settings(self): + """Checks, whether settings related to the new API stack make sense.""" + + # Old API stack checks. + if not self.enable_rl_module_and_learner: + # Throw a warning if the user has used `self.rl_module(rl_module_spec=...)` + # but has not enabled the new API stack at the same time. + if self._rl_module_spec is not None: + logger.warning( + "You have setup a RLModuleSpec (via calling " + "`config.rl_module(...)`), but have not enabled the new API stack. " + "To enable it, call `config.api_stack(enable_rl_module_and_learner=" + "True)`." + ) + # Throw a warning if the user has used `self.training(learner_class=...)` + # but has not enabled the new API stack at the same time. + if self._learner_class is not None: + logger.warning( + "You specified a custom Learner class (via " + f"`AlgorithmConfig.training(learner_class={self._learner_class})`, " + f"but have the new API stack disabled. You need to enable it via " + "`AlgorithmConfig.api_stack(enable_rl_module_and_learner=True)`." + ) + # User is using the new EnvRunners, but forgot to switch on + # `enable_rl_module_and_learner`. + if self.enable_env_runner_and_connector_v2: + self._value_error( + "You are using the new API stack EnvRunners (SingleAgentEnvRunner " + "or MultiAgentEnvRunner), but have forgotten to switch on the new " + "API stack! Try setting " + "`config.api_stack(enable_rl_module_and_learner=True)`." + ) + # Early out. The rest of this method is only for + # `enable_rl_module_and_learner=True`. + return + + # Warn about new API stack on by default. + if log_once(f"{self.algo_class.__name__}_on_new_api_stack"): + logger.warning( + f"You are running {self.algo_class.__name__} on the new API stack! " + "This is the new default behavior for this algorithm. If you don't " + "want to use the new API stack, set `config.api_stack(" + "enable_rl_module_and_learner=False," + "enable_env_runner_and_connector_v2=False)`. For a detailed migration " + "guide, see here: https://docs.ray.io/en/master/rllib/new-api-stack-migration-guide.html" # noqa + ) + + # Disabled hybrid API stack. Now, both `enable_rl_module_and_learner` and + # `enable_env_runner_and_connector_v2` must be True or both False. + if not self.enable_env_runner_and_connector_v2: + self._value_error( + "Setting `enable_rl_module_and_learner` to True and " + "`enable_env_runner_and_connector_v2` to False ('hybrid API stack'" + ") is not longer supported! Set both to True (new API stack) or both " + "to False (old API stack), instead." + ) + + # For those users that accidentally use the new API stack (because it's the + # default now for many algos), we need to make sure they are warned. + try: + tree.assert_same_structure(self.model, MODEL_DEFAULTS) + # Create copies excluding the specified key + check( + {k: v for k, v in self.model.items() if k != "vf_share_layers"}, + {k: v for k, v in MODEL_DEFAULTS.items() if k != "vf_share_layers"}, + ) + except Exception: + logger.warning( + "You configured a custom `model` config (probably through calling " + "config.training(model=..), whereas your config uses the new API " + "stack! In order to switch off the new API stack, set in your config: " + "`config.api_stack(enable_rl_module_and_learner=False, " + "enable_env_runner_and_connector_v2=False)`. If you DO want to use " + "the new API stack, configure your model, instead, through: " + "`config.rl_module(model_config={..})`." + ) + + # LR-schedule checking. + Scheduler.validate( + fixed_value_or_schedule=self.lr, + setting_name="lr", + description="learning rate", + ) + + # This is not compatible with RLModules, which all have a method + # `forward_exploration` to specify custom exploration behavior. + if self.exploration_config: + self._value_error( + "When the RLModule API is enabled, exploration_config can not be " + "set. If you want to implement custom exploration behaviour, " + "please modify the `forward_exploration` method of the " + "RLModule at hand. On configs that have a default exploration " + "config, this must be done via " + "`config.exploration_config={}`." + ) + + not_compatible_w_rlm_msg = ( + "Cannot use `{}` option with the new API stack (RLModule and " + "Learner APIs)! `{}` is part of the ModelV2 API and Policy API," + " which are not compatible with the new API stack. You can either " + "deactivate the new stack via `config.api_stack( " + "enable_rl_module_and_learner=False)`," + "or use the new stack (incl. RLModule API) and implement your " + "custom model as an RLModule." + ) + + if self.model["custom_model"] is not None: + self._value_error( + not_compatible_w_rlm_msg.format("custom_model", "custom_model") + ) + + if self.model["custom_model_config"] != {}: + self._value_error( + not_compatible_w_rlm_msg.format( + "custom_model_config", "custom_model_config" + ) + ) + + # TODO (sven): Once everything is on the new API stack, we won't need this method + # anymore. + def _validate_to_be_deprecated_settings(self): + # `render_env` is deprecated on new API stack. + if self.enable_env_runner_and_connector_v2 and self.render_env is not False: + deprecation_warning( + old="AlgorithmConfig.render_env", + help="The `render_env` setting is not supported on the new API stack! " + "In order to log videos to WandB (or other loggers), take a look at " + "this example here: " + "https://github.com/ray-project/ray/blob/master/rllib/examples/envs/env_rendering_and_recording.py", # noqa + ) + + if self.preprocessor_pref not in ["rllib", "deepmind", None]: + self._value_error( + "`config.preprocessor_pref` must be either 'rllib', 'deepmind' or None!" + ) + + # Check model config. + # If no preprocessing, propagate into model's config as well + # (so model knows whether inputs are preprocessed or not). + if self._disable_preprocessor_api is True: + self.model["_disable_preprocessor_api"] = True + # If no action flattening, propagate into model's config as well + # (so model knows whether action inputs are already flattened or not). + if self._disable_action_flattening is True: + self.model["_disable_action_flattening"] = True + if self.model.get("custom_preprocessor"): + deprecation_warning( + old="AlgorithmConfig.training(model={'custom_preprocessor': ...})", + help="Custom preprocessors are deprecated, " + "since they sometimes conflict with the built-in " + "preprocessors for handling complex observation spaces. " + "Please use wrapper classes around your environment " + "instead.", + error=True, + ) + + # Multi-GPU settings. + if self.simple_optimizer is True: + pass + # Multi-GPU setting: Must use MultiGPUTrainOneStep. + elif not self.enable_rl_module_and_learner and self.num_gpus > 1: + # TODO: AlphaStar uses >1 GPUs differently (1 per policy actor), so this is + # ok for tf2 here. + # Remove this hacky check, once we have fully moved to the Learner API. + if self.framework_str == "tf2" and type(self).__name__ != "AlphaStar": + self._value_error( + "`num_gpus` > 1 not supported yet for " + f"framework={self.framework_str}!" + ) + elif self.simple_optimizer is True: + self._value_error( + "Cannot use `simple_optimizer` if `num_gpus` > 1! " + "Consider not setting `simple_optimizer` in your config." + ) + self.simple_optimizer = False + # Auto-setting: Use simple-optimizer for tf-eager or multiagent, + # otherwise: MultiGPUTrainOneStep (if supported by the algo's execution + # plan). + elif self.simple_optimizer == DEPRECATED_VALUE: + # tf-eager: Must use simple optimizer. + if self.framework_str not in ["tf", "torch"]: + self.simple_optimizer = True + # Multi-agent case: Try using MultiGPU optimizer (only + # if all policies used are DynamicTFPolicies or TorchPolicies). + elif self.is_multi_agent: + from ray.rllib.policy.dynamic_tf_policy import DynamicTFPolicy + from ray.rllib.policy.torch_policy import TorchPolicy + + default_policy_cls = None + if self.algo_class: + default_policy_cls = self.algo_class.get_default_policy_class(self) + + policies = self.policies + policy_specs = ( + [ + PolicySpec(*spec) if isinstance(spec, (tuple, list)) else spec + for spec in policies.values() + ] + if isinstance(policies, dict) + else [PolicySpec() for _ in policies] + ) + + if any( + (spec.policy_class or default_policy_cls) is None + or not issubclass( + spec.policy_class or default_policy_cls, + (DynamicTFPolicy, TorchPolicy), + ) + for spec in policy_specs + ): + self.simple_optimizer = True + else: + self.simple_optimizer = False + else: + self.simple_optimizer = False + + # User manually set simple-optimizer to False -> Error if tf-eager. + elif self.simple_optimizer is False: + if self.framework_str == "tf2": + self._value_error( + "`simple_optimizer=False` not supported for " + f"config.framework({self.framework_str})!" + ) + + def _validate_offline_settings(self): + # If a user does not have an environment and cannot run evaluation, + # or does not want to run evaluation, she needs to provide at least + # action and observation spaces. Note, we require here the spaces, + # i.e. a user cannot provide an environment instead because we do + # not want to create the environment to receive spaces. + if ( + self.is_offline + and not self.is_online + and ( + not (self.evaluation_num_env_runners > 0 or self.evaluation_interval) + and (self.action_space is None or self.observation_space is None) + ) + ): + self._value_error( + "If no evaluation should be run, `action_space` and " + "`observation_space` must be provided." + ) + + if self.ignore_final_observation and self.algo_class.__name__ != "BC": + logger.warning( + "`ignore_final_observation=True` (zeros-out truncation observations), " + "but the algorithm isn't `BC`. It is recommended to use this " + "setting only with `BC`, b/c other RL algorithms rely on truncation-" + "observations due to value function estimates." + ) + + from ray.rllib.offline.offline_data import OfflineData + from ray.rllib.offline.offline_prelearner import OfflinePreLearner + + if self.offline_data_class and not issubclass( + self.offline_data_class, OfflineData + ): + self._value_error( + "Unknown `offline_data_class`. OfflineData class needs to inherit " + "from `OfflineData` class." + ) + if self.prelearner_class and not issubclass( + self.prelearner_class, OfflinePreLearner + ): + self._value_error( + "Unknown `prelearner_class`. PreLearner class needs to inherit " + "from `OfflinePreLearner` class." + ) + + from ray.rllib.utils.replay_buffers.episode_replay_buffer import ( + EpisodeReplayBuffer, + ) + + if self.prelearner_buffer_class and not issubclass( + self.prelearner_buffer_class, EpisodeReplayBuffer + ): + self._value_error( + "Unknown `prelearner_buffer_class`. The buffer class for the " + "prelearner needs to inherit from `EpisodeReplayBuffer`. " + "Specifically it needs to store and sample lists of " + "`Single-/MultiAgentEpisode`s." + ) + + if self.input_read_batch_size and not ( + self.input_read_episodes or self.input_read_sample_batches + ): + self._value_error( + "Setting `input_read_batch_size` is only allowed in case of a " + "dataset that holds either `EpisodeType` or `BatchType` data (i.e. " + "rows that contains multiple timesteps), but neither " + "`input_read_episodes` nor `input_read_sample_batches` is set to " + "`True`." + ) + + if ( + self.output + and self.output_write_episodes + and self.batch_mode != "complete_episodes" + ): + self._value_error( + "When recording episodes only complete episodes should be " + "recorded (i.e. `batch_mode=='complete_episodes'`). Otherwise " + "recorded episodes cannot be read in for training." + ) + + # Offline evaluation. + from ray.rllib.offline.offline_policy_evaluation_runner import ( + OfflinePolicyEvaluationTypes, + ) + + offline_eval_types = list(OfflinePolicyEvaluationTypes) + if ( + self.offline_evaluation_type + and self.offline_evaluation_type != "eval_loss" + and self.offline_evaluation_type not in OfflinePolicyEvaluationTypes + ): + self._value_error( + f"Unknown offline evaluation type: {self.offline_evaluation_type}." + "Available types of offline evaluation are either `'eval_loss' to evaluate " + f"the training loss on a validation dataset or {offline_eval_types}." + ) + + from ray.rllib.offline.offline_evaluation_runner import OfflineEvaluationRunner + + if self.offline_eval_runner_class and not issubclass( + self.offline_eval_runner_class, OfflineEvaluationRunner + ): + self._value_error( + "Unknown `offline_eval_runner_class`. OfflineEvaluationRunner class needs to inherit " + "from `OfflineEvaluationRunner` class." + ) + + @property + def is_online(self) -> bool: + """Defines if this config is for online RL. + + Note, a config can be for on- and offline training at the same time. + """ + return self._is_online + + @property + def is_offline(self) -> bool: + """Defines, if this config is for offline RL.""" + return ( + # Does the user provide any input path/class? + bool(self.input_) + # Is it a real string path or list of such paths. + and ( + isinstance(self.input_, str) + or (isinstance(self.input_, list) and isinstance(self.input_[0], str)) + ) + # Could be old stack - which is considered very differently. + and self.input_ != "sampler" + and self.enable_rl_module_and_learner + ) + + @staticmethod + def _serialize_dict(config): + # Serialize classes to classpaths: + if "callbacks_class" in config: + config["callbacks"] = config.pop("callbacks_class") + if "class" in config: + config["class"] = serialize_type(config["class"]) + config["callbacks"] = serialize_type(config["callbacks"]) + config["sample_collector"] = serialize_type(config["sample_collector"]) + if isinstance(config["env"], type): + config["env"] = serialize_type(config["env"]) + if "replay_buffer_config" in config and ( + isinstance(config["replay_buffer_config"].get("type"), type) + ): + config["replay_buffer_config"]["type"] = serialize_type( + config["replay_buffer_config"]["type"] + ) + if isinstance(config["exploration_config"].get("type"), type): + config["exploration_config"]["type"] = serialize_type( + config["exploration_config"]["type"] + ) + if isinstance(config["model"].get("custom_model"), type): + config["model"]["custom_model"] = serialize_type( + config["model"]["custom_model"] + ) + + # List'ify `policies`, iff a set or tuple (these types are not JSON'able). + ma_config = config.get("multiagent") + if ma_config is not None: + if isinstance(ma_config.get("policies"), (set, tuple)): + ma_config["policies"] = list(ma_config["policies"]) + # Do NOT serialize functions/lambdas. + if ma_config.get("policy_mapping_fn"): + ma_config["policy_mapping_fn"] = NOT_SERIALIZABLE + if ma_config.get("policies_to_train"): + ma_config["policies_to_train"] = NOT_SERIALIZABLE + # However, if these "multiagent" settings have been provided directly + # on the top-level (as they should), we override the settings under + # "multiagent". Note that the "multiagent" key should no longer be used anyways. + if isinstance(config.get("policies"), (set, tuple)): + config["policies"] = list(config["policies"]) + # Do NOT serialize functions/lambdas. + if config.get("policy_mapping_fn"): + config["policy_mapping_fn"] = NOT_SERIALIZABLE + if config.get("policies_to_train"): + config["policies_to_train"] = NOT_SERIALIZABLE + + return config + + @staticmethod + def _translate_special_keys(key: str, warn_deprecated: bool = True) -> str: + # Handle special key (str) -> `AlgorithmConfig.[some_property]` cases. + if key == "callbacks": + key = "callbacks_class" + elif key == "create_env_on_driver": + key = "create_env_on_local_worker" + elif key == "custom_eval_function": + key = "custom_evaluation_function" + elif key == "framework": + key = "framework_str" + elif key == "input": + key = "input_" + elif key == "lambda": + key = "lambda_" + elif key == "num_cpus_for_driver": + key = "num_cpus_for_main_process" + elif key == "num_workers": + key = "num_env_runners" + + # Deprecated keys. + if warn_deprecated: + if key == "collect_metrics_timeout": + deprecation_warning( + old="collect_metrics_timeout", + new="metrics_episode_collection_timeout_s", + error=True, + ) + elif key == "metrics_smoothing_episodes": + deprecation_warning( + old="config.metrics_smoothing_episodes", + new="config.metrics_num_episodes_for_smoothing", + error=True, + ) + elif key == "min_iter_time_s": + deprecation_warning( + old="config.min_iter_time_s", + new="config.min_time_s_per_iteration", + error=True, + ) + elif key == "min_time_s_per_reporting": + deprecation_warning( + old="config.min_time_s_per_reporting", + new="config.min_time_s_per_iteration", + error=True, + ) + elif key == "min_sample_timesteps_per_reporting": + deprecation_warning( + old="config.min_sample_timesteps_per_reporting", + new="config.min_sample_timesteps_per_iteration", + error=True, + ) + elif key == "min_train_timesteps_per_reporting": + deprecation_warning( + old="config.min_train_timesteps_per_reporting", + new="config.min_train_timesteps_per_iteration", + error=True, + ) + elif key == "timesteps_per_iteration": + deprecation_warning( + old="config.timesteps_per_iteration", + new="`config.min_sample_timesteps_per_iteration` OR " + "`config.min_train_timesteps_per_iteration`", + error=True, + ) + elif key == "evaluation_num_episodes": + deprecation_warning( + old="config.evaluation_num_episodes", + new="`config.evaluation_duration` and " + "`config.evaluation_duration_unit=episodes`", + error=True, + ) + + return key + + def _check_if_correct_nn_framework_installed(self, _tf1, _tf, _torch): + """Check if tf/torch experiment is running and tf/torch installed.""" + if self.framework_str in {"tf", "tf2"}: + if not (_tf1 or _tf): + raise ImportError( + ( + "TensorFlow was specified as the framework to use (via `config." + "framework([tf|tf2])`)! However, no installation was " + "found. You can install TensorFlow via `pip install tensorflow`" + ) + ) + elif self.framework_str == "torch": + if not _torch: + raise ImportError( + ( + "PyTorch was specified as the framework to use (via `config." + "framework('torch')`)! However, no installation was found. You " + "can install PyTorch via `pip install torch`." + ) + ) + + def _resolve_tf_settings(self, _tf1, _tfv): + """Check and resolve tf settings.""" + if _tf1 and self.framework_str == "tf2": + if self.framework_str == "tf2" and _tfv < 2: + raise ValueError( + "You configured `framework`=tf2, but your installed " + "pip tf-version is < 2.0! Make sure your TensorFlow " + "version is >= 2.x." + ) + if not _tf1.executing_eagerly(): + _tf1.enable_eager_execution() + # Recommend setting tracing to True for speedups. + logger.info( + f"Executing eagerly (framework='{self.framework_str}')," + f" with eager_tracing={self.eager_tracing}. For " + "production workloads, make sure to set eager_tracing=True" + " in order to match the speed of tf-static-graph " + "(framework='tf'). For debugging purposes, " + "`eager_tracing=False` is the best choice." + ) + # Tf-static-graph (framework=tf): Recommend upgrading to tf2 and + # enabling eager tracing for similar speed. + elif _tf1 and self.framework_str == "tf": + logger.info( + "Your framework setting is 'tf', meaning you are using " + "static-graph mode. Set framework='tf2' to enable eager " + "execution with tf2.x. You may also then want to set " + "eager_tracing=True in order to reach similar execution " + "speed as with static-graph mode." + ) + + @OldAPIStack + def get_multi_agent_setup( + self, + *, + policies: Optional[MultiAgentPolicyConfigDict] = None, + env: Optional[EnvType] = None, + spaces: Optional[Dict[PolicyID, Tuple[gym.Space, gym.Space]]] = None, + default_policy_class: Optional[Type[Policy]] = None, + ) -> Tuple[MultiAgentPolicyConfigDict, Callable[[PolicyID, SampleBatchType], bool]]: + r"""Compiles complete multi-agent config (dict) from the information in `self`. + + Infers the observation- and action spaces, the policy classes, and the policy's + configs. The returned `MultiAgentPolicyConfigDict` is fully unified and strictly + maps PolicyIDs to complete PolicySpec objects (with all their fields not-None). + + Examples: + .. testcode:: + + import gymnasium as gym + from ray.rllib.algorithms.ppo import PPOConfig + config = ( + PPOConfig() + .environment("CartPole-v1") + .framework("torch") + .multi_agent(policies={"pol1", "pol2"}, policies_to_train=["pol1"]) + ) + policy_dict, is_policy_to_train = config.get_multi_agent_setup( + env=gym.make("CartPole-v1")) + is_policy_to_train("pol1") + is_policy_to_train("pol2") + + Args: + policies: An optional multi-agent `policies` dict, mapping policy IDs + to PolicySpec objects. If not provided uses `self.policies` + instead. Note that the `policy_class`, `observation_space`, and + `action_space` properties in these PolicySpecs may be None and must + therefore be inferred here. + env: An optional env instance, from which to infer the different spaces for + the different policies. If not provided, tries to infer from + `spaces`. Otherwise from `self.observation_space` and + `self.action_space`. Raises an error, if no information on spaces can be + infered. + spaces: Optional dict mapping policy IDs to tuples of 1) observation space + and 2) action space that should be used for the respective policy. + These spaces were usually provided by an already instantiated remote + EnvRunner. Note that if the `env` argument is provided, tries to + infer spaces from `env` first. + default_policy_class: The Policy class to use should a PolicySpec have its + policy_class property set to None. + + Returns: + A tuple consisting of 1) a MultiAgentPolicyConfigDict and 2) a + `is_policy_to_train(PolicyID, SampleBatchType) -> bool` callable. + + Raises: + ValueError: In case, no spaces can be infered for the policy/ies. + ValueError: In case, two agents in the env map to the same PolicyID + (according to `self.policy_mapping_fn`), but have different action- or + observation spaces according to the infered space information. + """ + policies = copy.deepcopy(policies or self.policies) + + # Policies given as set/list/tuple (of PolicyIDs) -> Setup each policy + # automatically via empty PolicySpec (makes RLlib infer observation- and + # action spaces as well as the Policy's class). + if isinstance(policies, (set, list, tuple)): + policies = {pid: PolicySpec() for pid in policies} + + # Try extracting spaces from env or from given spaces dict. + env_obs_space = None + env_act_space = None + + # Env is a ray.remote: Get spaces via its (automatically added) + # `_get_spaces()` method. + if isinstance(env, ray.actor.ActorHandle): + env_obs_space, env_act_space = ray.get(env._get_spaces.remote()) + # Normal env (gym.Env or MultiAgentEnv): These should have the + # `observation_space` and `action_space` properties. + elif env is not None: + # `env` is a gymnasium.vector.Env. + if hasattr(env, "single_observation_space") and isinstance( + env.single_observation_space, gym.Space + ): + env_obs_space = env.single_observation_space + # `env` is a gymnasium.Env. + elif hasattr(env, "observation_space") and isinstance( + env.observation_space, gym.Space + ): + env_obs_space = env.observation_space + + # `env` is a gymnasium.vector.Env. + if hasattr(env, "single_action_space") and isinstance( + env.single_action_space, gym.Space + ): + env_act_space = env.single_action_space + # `env` is a gymnasium.Env. + elif hasattr(env, "action_space") and isinstance( + env.action_space, gym.Space + ): + env_act_space = env.action_space + + # Last resort: Try getting the env's spaces from the spaces + # dict's special __env__ key. + if spaces is not None: + if env_obs_space is None: + env_obs_space = spaces.get(INPUT_ENV_SPACES, [None])[0] + if env_act_space is None: + env_act_space = spaces.get(INPUT_ENV_SPACES, [None, None])[1] + + # Check each defined policy ID and unify its spec. + for pid, policy_spec in policies.copy().items(): + # Convert to PolicySpec if plain list/tuple. + if not isinstance(policy_spec, PolicySpec): + policies[pid] = policy_spec = PolicySpec(*policy_spec) + + # Infer policy classes for policies dict, if not provided (None). + if policy_spec.policy_class is None and default_policy_class is not None: + policies[pid].policy_class = default_policy_class + + # Infer observation space. + if policy_spec.observation_space is None: + env_unwrapped = env.unwrapped if hasattr(env, "unwrapped") else env + # Module's space is provided -> Use it as-is. + if spaces is not None and pid in spaces: + obs_space = spaces[pid][0] + # MultiAgentEnv -> Check, whether agents have different spaces. + elif isinstance(env_unwrapped, MultiAgentEnv): + obs_space = None + mapping_fn = self.policy_mapping_fn + aids = list( + env_unwrapped.possible_agents + if hasattr(env_unwrapped, "possible_agents") + and env_unwrapped.possible_agents + else env_unwrapped.get_agent_ids() + ) + if len(aids) == 0: + one_obs_space = env_unwrapped.observation_space + else: + one_obs_space = env_unwrapped.get_observation_space(aids[0]) + # If all obs spaces are the same, just use the first space. + if all( + env_unwrapped.get_observation_space(aid) == one_obs_space + for aid in aids + ): + obs_space = one_obs_space + # Need to reverse-map spaces (for the different agents) to certain + # policy IDs. We have to compare the ModuleID with all possible + # AgentIDs and find the agent ID that matches. + elif mapping_fn: + for aid in aids: + # Match: Assign spaces for this agentID to the PolicyID. + if mapping_fn(aid, None, worker=None) == pid: + # Make sure, different agents that map to the same + # policy don't have different spaces. + if ( + obs_space is not None + and env_unwrapped.get_observation_space(aid) + != obs_space + ): + raise ValueError( + "Two agents in your environment map to the " + "same policyID (as per your `policy_mapping" + "_fn`), however, these agents also have " + "different observation spaces!" + ) + obs_space = env_unwrapped.get_observation_space(aid) + # Just use env's obs space as-is. + elif env_obs_space is not None: + obs_space = env_obs_space + # Space given directly in config. + elif self.observation_space: + obs_space = self.observation_space + else: + raise ValueError( + "`observation_space` not provided in PolicySpec for " + f"{pid} and env does not have an observation space OR " + "no spaces received from other workers' env(s) OR no " + "`observation_space` specified in config!" + ) + + policies[pid].observation_space = obs_space + + # Infer action space. + if policy_spec.action_space is None: + env_unwrapped = env.unwrapped if hasattr(env, "unwrapped") else env + # Module's space is provided -> Use it as-is. + if spaces is not None and pid in spaces: + act_space = spaces[pid][1] + # MultiAgentEnv -> Check, whether agents have different spaces. + elif isinstance(env_unwrapped, MultiAgentEnv): + act_space = None + mapping_fn = self.policy_mapping_fn + aids = list( + env_unwrapped.possible_agents + if hasattr(env_unwrapped, "possible_agents") + and env_unwrapped.possible_agents + else env_unwrapped.get_agent_ids() + ) + if len(aids) == 0: + one_act_space = env_unwrapped.action_space + else: + one_act_space = env_unwrapped.get_action_space(aids[0]) + # If all obs spaces are the same, just use the first space. + if all( + env_unwrapped.get_action_space(aid) == one_act_space + for aid in aids + ): + act_space = one_act_space + # Need to reverse-map spaces (for the different agents) to certain + # policy IDs. We have to compare the ModuleID with all possible + # AgentIDs and find the agent ID that matches. + elif mapping_fn: + for aid in aids: + # Match: Assign spaces for this AgentID to the PolicyID. + if mapping_fn(aid, None, worker=None) == pid: + # Make sure, different agents that map to the same + # policy don't have different spaces. + if ( + act_space is not None + and env_unwrapped.get_action_space(aid) != act_space + ): + raise ValueError( + "Two agents in your environment map to the " + "same policyID (as per your `policy_mapping" + "_fn`), however, these agents also have " + "different action spaces!" + ) + act_space = env_unwrapped.get_action_space(aid) + # Just use env's action space as-is. + elif env_act_space is not None: + act_space = env_act_space + elif self.action_space: + act_space = self.action_space + else: + raise ValueError( + "`action_space` not provided in PolicySpec for " + f"{pid} and env does not have an action space OR " + "no spaces received from other workers' env(s) OR no " + "`action_space` specified in config!" + ) + policies[pid].action_space = act_space + + # Create entire AlgorithmConfig object from the provided override. + # If None, use {} as override. + if not isinstance(policies[pid].config, AlgorithmConfig): + assert policies[pid].config is None or isinstance( + policies[pid].config, dict + ) + policies[pid].config = self.copy(copy_frozen=False).update_from_dict( + policies[pid].config or {} + ) + + # If collection given, construct a simple default callable returning True + # if the PolicyID is found in the list/set of IDs. + if self.policies_to_train is not None and not callable(self.policies_to_train): + pols = set(self.policies_to_train) + + def is_policy_to_train(pid, batch=None): + return pid in pols + + else: + is_policy_to_train = self.policies_to_train + + return policies, is_policy_to_train + + @Deprecated(new="AlgorithmConfig.build_algo", error=False) + def build(self, *args, **kwargs): + return self.build_algo(*args, **kwargs) + + @Deprecated(new="AlgorithmConfig.get_multi_rl_module_spec()", error=True) + def get_marl_module_spec(self, *args, **kwargs): + pass + + @Deprecated(new="AlgorithmConfig.env_runners(..)", error=True) + def rollouts(self, *args, **kwargs): + pass + + @Deprecated(new="AlgorithmConfig.env_runners(..)", error=True) + def exploration(self, *args, **kwargs): + pass + + @property + @Deprecated( + new="AlgorithmConfig.fault_tolerance(restart_failed_env_runners=..)", + error=True, + ) + def recreate_failed_env_runners(self): + pass + + @recreate_failed_env_runners.setter + def recreate_failed_env_runners(self, value): + deprecation_warning( + old="AlgorithmConfig.recreate_failed_env_runners", + new="AlgorithmConfig.restart_failed_env_runners", + error=True, + ) + + @property + @Deprecated(new="AlgorithmConfig._enable_new_api_stack", error=True) + def _enable_new_api_stack(self): + pass + + @_enable_new_api_stack.setter + def _enable_new_api_stack(self, value): + deprecation_warning( + old="AlgorithmConfig._enable_new_api_stack", + new="AlgorithmConfig.enable_rl_module_and_learner", + error=True, + ) + + @property + @Deprecated(new="AlgorithmConfig.enable_env_runner_and_connector_v2", error=True) + def uses_new_env_runners(self): + pass + + @property + @Deprecated(new="AlgorithmConfig.num_env_runners", error=True) + def num_rollout_workers(self): + pass + + @num_rollout_workers.setter + def num_rollout_workers(self, value): + deprecation_warning( + old="AlgorithmConfig.num_rollout_workers", + new="AlgorithmConfig.num_env_runners", + error=True, + ) + + @property + @Deprecated(new="AlgorithmConfig.evaluation_num_workers", error=True) + def evaluation_num_workers(self): + pass + + @evaluation_num_workers.setter + def evaluation_num_workers(self, value): + deprecation_warning( + old="AlgorithmConfig.evaluation_num_workers", + new="AlgorithmConfig.evaluation_num_env_runners", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.num_envs_per_env_runner", error=True) + def num_envs_per_worker(self): + pass + + @num_envs_per_worker.setter + def num_envs_per_worker(self, value): + deprecation_warning( + old="AlgorithmConfig.num_envs_per_worker", + new="AlgorithmConfig.num_envs_per_env_runner", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.ignore_env_runner_failures", error=True) + def ignore_worker_failures(self): + pass + + @ignore_worker_failures.setter + def ignore_worker_failures(self, value): + deprecation_warning( + old="AlgorithmConfig.ignore_worker_failures", + new="AlgorithmConfig.ignore_env_runner_failures", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.restart_failed_env_runners", error=True) + def recreate_failed_workers(self): + pass + + @recreate_failed_workers.setter + def recreate_failed_workers(self, value): + deprecation_warning( + old="AlgorithmConfig.recreate_failed_workers", + new="AlgorithmConfig.restart_failed_env_runners", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.max_num_env_runner_restarts", error=True) + def max_num_worker_restarts(self): + pass + + @max_num_worker_restarts.setter + def max_num_worker_restarts(self, value): + deprecation_warning( + old="AlgorithmConfig.max_num_worker_restarts", + new="AlgorithmConfig.max_num_env_runner_restarts", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.delay_between_env_runner_restarts_s", error=True) + def delay_between_worker_restarts_s(self): + pass + + @delay_between_worker_restarts_s.setter + def delay_between_worker_restarts_s(self, value): + deprecation_warning( + old="AlgorithmConfig.delay_between_worker_restarts_s", + new="AlgorithmConfig.delay_between_env_runner_restarts_s", + error=True, + ) + pass + + @property + @Deprecated( + new="AlgorithmConfig.num_consecutive_env_runner_failures_tolerance", error=True + ) + def num_consecutive_worker_failures_tolerance(self): + pass + + @num_consecutive_worker_failures_tolerance.setter + def num_consecutive_worker_failures_tolerance(self, value): + deprecation_warning( + old="AlgorithmConfig.num_consecutive_worker_failures_tolerance", + new="AlgorithmConfig.num_consecutive_env_runner_failures_tolerance", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.env_runner_health_probe_timeout_s", error=True) + def worker_health_probe_timeout_s(self): + pass + + @worker_health_probe_timeout_s.setter + def worker_health_probe_timeout_s(self, value): + deprecation_warning( + old="AlgorithmConfig.worker_health_probe_timeout_s", + new="AlgorithmConfig.env_runner_health_probe_timeout_s", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.env_runner_restore_timeout_s", error=True) + def worker_restore_timeout_s(self): + pass + + @worker_restore_timeout_s.setter + def worker_restore_timeout_s(self, value): + deprecation_warning( + old="AlgorithmConfig.worker_restore_timeout_s", + new="AlgorithmConfig.env_runner_restore_timeout_s", + error=True, + ) + pass + + @property + @Deprecated( + new="AlgorithmConfig.validate_env_runners_after_construction", + error=True, + ) + def validate_workers_after_construction(self): + pass + + @validate_workers_after_construction.setter + def validate_workers_after_construction(self, value): + deprecation_warning( + old="AlgorithmConfig.validate_workers_after_construction", + new="AlgorithmConfig.validate_env_runners_after_construction", + error=True, + ) + pass + + # Cleanups from `resources()`. + @property + @Deprecated(new="AlgorithmConfig.num_cpus_per_env_runner", error=True) + def num_cpus_per_worker(self): + pass + + @num_cpus_per_worker.setter + def num_cpus_per_worker(self, value): + deprecation_warning( + old="AlgorithmConfig.num_cpus_per_worker", + new="AlgorithmConfig.num_cpus_per_env_runner", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.num_gpus_per_env_runner", error=True) + def num_gpus_per_worker(self): + pass + + @num_gpus_per_worker.setter + def num_gpus_per_worker(self, value): + deprecation_warning( + old="AlgorithmConfig.num_gpus_per_worker", + new="AlgorithmConfig.num_gpus_per_env_runner", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.custom_resources_per_env_runner", error=True) + def custom_resources_per_worker(self): + pass + + @custom_resources_per_worker.setter + def custom_resources_per_worker(self, value): + deprecation_warning( + old="AlgorithmConfig.custom_resources_per_worker", + new="AlgorithmConfig.custom_resources_per_env_runner", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.num_learners", error=True) + def num_learner_workers(self): + pass + + @num_learner_workers.setter + def num_learner_workers(self, value): + deprecation_warning( + old="AlgorithmConfig.num_learner_workers", + new="AlgorithmConfig.num_learners", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.num_cpus_per_learner", error=True) + def num_cpus_per_learner_worker(self): + pass + + @num_cpus_per_learner_worker.setter + def num_cpus_per_learner_worker(self, value): + deprecation_warning( + old="AlgorithmConfig.num_cpus_per_learner_worker", + new="AlgorithmConfig.num_cpus_per_learner", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.num_gpus_per_learner", error=True) + def num_gpus_per_learner_worker(self): + pass + + @num_gpus_per_learner_worker.setter + def num_gpus_per_learner_worker(self, value): + deprecation_warning( + old="AlgorithmConfig.num_gpus_per_learner_worker", + new="AlgorithmConfig.num_gpus_per_learner", + error=True, + ) + pass + + @property + @Deprecated(new="AlgorithmConfig.num_cpus_for_local_worker", error=True) + def num_cpus_for_local_worker(self): + pass + + @num_cpus_for_local_worker.setter + def num_cpus_for_local_worker(self, value): + deprecation_warning( + old="AlgorithmConfig.num_cpus_for_local_worker", + new="AlgorithmConfig.num_cpus_for_main_process", + error=True, + ) + pass + + +class DifferentiableAlgorithmConfig(AlgorithmConfig): + """An RLlib DifferentiableAlgorithmConfig builds a Meta algorithm from a given + configuration + + .. testcode:: + + from ray.rllib.algorithm.algorithm_config import DifferentiableAlgorithmConfig + # Construct a generic config for an algorithm that needs differentiable Learners. + config = ( + DifferentiableAlgorithmConfig() + .training(lr=3e-4) + .environment(env="CartPole-v1") + .learners( + differentiable_learner_configs=[ + DifferentiableLearnerConfig( + DifferentiableTorchLearner, + lr=1e-4, + ) + ] + ) + ) + # Similar to `AlgorithmConfig` the config using differentiable Learners can be + # used to build a respective `Algorithm`. + algo = config.build() + + + """ + + # A list of `DifferentiableLearnerConfig` instances that define differentiable + # `Learner`'s. Note, each of them needs to implement the `DifferentiableLearner` + # API. + differentiable_learner_configs: List[DifferentiableLearnerConfig] + + def __init__(self, algo_class=None): + """Initializes the DifferentiableLearnerConfig instance. + + Args: + algo_class: An optional Algorithm class that this config class belongs to. + Used (if provided) to build a respective Algorithm instance from this + config. + """ + # Initialize the `AlgorithmConfig` first. + super().__init__(algo_class=algo_class) + + # Initialize the list of differentiable learner configs to an empty list, which + # defines the default, i.e. the `MetaLearner` will have no nested updates. + self.differentiable_learner_configs: List[DifferentiableLearnerConfig] = [] + + def learners( + self, + *, + learner_class: Optional[Type["Learner"]] = NotProvided, + learner_connector: Optional[ + Callable[["RLModule"], Union["ConnectorV2", List["ConnectorV2"]]] + ] = NotProvided, + add_default_connectors_to_learner_pipeline: Optional[bool] = NotProvided, + learner_config_dict: Optional[Dict[str, Any]] = NotProvided, + differentiable_learner_configs: List[DifferentiableLearnerConfig] = NotProvided, + **kwargs, + ) -> "DifferentiableAlgorithmConfig": + """Sets the configurations for differentiable learners. + + Args: + learner_class: The `Learner` class to use for (distributed) updating of the + RLModule. Only used when `enable_rl_module_and_learner=True`. + learner_connector: A callable taking an env observation space and an env + action space as inputs and returning a learner ConnectorV2 (might be + a pipeline) object. + add_default_connectors_to_learner_pipeline: If True (default), RLlib's + Learners automatically add the default Learner ConnectorV2 + pieces to the LearnerPipeline. These automatically perform: + a) adding observations from episodes to the train batch, if this has not + already been done by a user-provided connector piece + b) if RLModule is stateful, add a time rank to the train batch, zero-pad + the data, and add the correct state inputs, if this has not already been + done by a user-provided connector piece. + c) add all other information (actions, rewards, terminateds, etc..) to + the train batch, if this has not already been done by a user-provided + connector piece. + Only if you know exactly what you are doing, you + should set this setting to False. + Note that this setting is only relevant if the new API stack is used + (including the new EnvRunner classes). + learner_config_dict: A dict to insert any settings accessible from within + the Learner instance. This should only be used in connection with custom + Learner subclasses and in case the user doesn't want to write an extra + `AlgorithmConfig` subclass just to add a few settings to the base Algo's + own config class. + differentiable_learner_configs: A list of `DifferentiableLearnerConfig` instances + defining the `DifferentiableLearner` classes used for the nested updates in + `Algorithm`'s learner. + """ + super().learners(**kwargs) + + if learner_class is not NotProvided: + self._learner_class = learner_class + if learner_connector is not NotProvided: + self._learner_connector = learner_connector + if add_default_connectors_to_learner_pipeline is not NotProvided: + self.add_default_connectors_to_learner_pipeline = ( + add_default_connectors_to_learner_pipeline + ) + if learner_config_dict is not NotProvided: + self.learner_config_dict.update(learner_config_dict) + if differentiable_learner_configs is not NotProvided: + self.differentiable_learner_configs = differentiable_learner_configs + + return self + + def validate(self): + """Validates all values in this config.""" + + # First, call the `validate` method of super. + super().validate() + + # TODO (simon): Maybe moving this to a private method? + # Ensure that the default learner class is derived from `TorchMetaLearner`. + from ray.rllib.core.learner.torch.torch_meta_learner import TorchMetaLearner + + if not issubclass(self.get_default_learner_class(), TorchMetaLearner): + self._value_error( + "`get_default_learner_class` must return a `MetaLearner` class " + f"or sublass but got {self.get_default_learner_class()}." + ) + # Make sure that the differentiable learner configs are contained in a list. + if not isinstance(self.differentiable_learner_configs, list): + self._value_error( + "`differentiable_learner_configs` must be a list of " + "`DifferentiableLearnerConfig` instances, but is " + f"{type(self.differentiable_learner_configs)}." + ) + # In addition, check, if all configurations are wrapped in a + # `DifferentiableLearnerConfig`. + elif not all( + isinstance(learner_cfg, DifferentiableLearnerConfig) + for learner_cfg in self.differentiable_learner_configs + ): + self._value_error( + "`differentiable_learner_configs` must be a list of " + "`DifferentiableLearnerConfig` instances, but at least " + "one instance is not a `DifferentiableLearnerConfig`." + ) + + def get_default_learner_class(self) -> Union[Type["TorchMetaLearner"], str]: + """Returns the `MetaLearner` class to use for this algorithm. + + Override this method in the sub-class to return the `MetaLearner`. + + Returns: + The `MetaLearner` class to use for this algorithm either as a class + type or as a string. (e.g. "ray.rllib.core.learner.torch.torch_meta_learner.TorchMetaLearner") + """ + return NotImplemented + + def get_differentiable_learner_classes( + self, + ) -> List[Union[Type["DifferentiableLearner"], str]]: + """Returns the `DifferentiableLearner` classes to use for this algorithm. + + Override this method in the sub-class to return the `DifferentiableLearner`. + + Returns: + The `DifferentiableLearner` class to use for this algorithm either as a class + type or as a string. (e.g. + "ray.rllib.core.learner.torch.torch_meta_learner.TorchDifferentiableLearner"). + """ + return NotImplemented + + def get_differentiable_learner_configs(self) -> List[DifferentiableLearnerConfig]: + """Returns the `DifferentiableLearnerConfigs` for all `DifferentiableLearner`s. + + Override this method in the sub-class to return the `DifferentiableLearnerConfig`s. + + Returns: + The `DifferentiableLearnerConfig` instances to use for this algorithm. + """ + return self.differentiable_learner_configs + + +class TorchCompileWhatToCompile(str, Enum): + """Enumerates schemes of what parts of the TorchLearner can be compiled. + + This can be either the entire update step of the learner or only the forward + methods (and therein the forward_train method) of the RLModule. + + .. note:: + - torch.compiled code can become slow on graph breaks or even raise + errors on unsupported operations. Empirically, compiling + `forward_train` should introduce little graph breaks, raise no + errors but result in a speedup comparable to compiling the + complete update. + - Using `complete_update` is experimental and may result in errors. + """ + + # Compile the entire update step of the learner. + # This includes the forward pass of the RLModule, the loss computation, and the + # optimizer step. + COMPLETE_UPDATE = "complete_update" + # Only compile the forward methods (and therein the forward_train method) of the + # RLModule. + FORWARD_TRAIN = "forward_train" diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fe4eb9b9a10fb5482dce8bbb4b89589207cbc84f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__init__.py @@ -0,0 +1,12 @@ +from ray.rllib.algorithms.appo.appo import APPO, APPOConfig +from ray.rllib.algorithms.appo.appo_tf_policy import APPOTF1Policy, APPOTF2Policy +from ray.rllib.algorithms.appo.appo_torch_policy import APPOTorchPolicy + +__all__ = [ + "APPO", + "APPOConfig", + # @OldAPIStack + "APPOTF1Policy", + "APPOTF2Policy", + "APPOTorchPolicy", +] diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c9f146ecb271d34908a61e3f6a30624a45589cd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e404198ed7cbbabb63def261dc2b31b7d8467e0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5d7550fe88a4a810cc9d14771245740454a63e09 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d3f412bd850f3b1e7e09871f49141a5a86be83b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_tf_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_tf_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f7fba93e95a28acac161bd9bdc8d04e3ec62ba70 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_tf_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_torch_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_torch_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85a20f00aa6ada06e53252bc04f7c664934c46fa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/appo_torch_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/default_appo_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/default_appo_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..637ef07ec25117291ccbee8f8049c37525bc5141 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/default_appo_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b903ed7e185eb28704ed9c579bf50ec2baf2932a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo.py new file mode 100644 index 0000000000000000000000000000000000000000..d023913ac0022ed77691a6ddf1711173f0e2412d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo.py @@ -0,0 +1,435 @@ +"""Asynchronous Proximal Policy Optimization (APPO) + +The algorithm is described in [1] (under the name of "IMPACT"): + +Detailed documentation: +https://docs.ray.io/en/master/rllib-algorithms.html#appo + +[1] IMPACT: Importance Weighted Asynchronous Architectures with Clipped Target Networks. +Luo et al. 2020 +https://arxiv.org/pdf/1912.00167 +""" + +import logging +from typing import Optional, Type + +from typing_extensions import Self + +from ray._common.deprecation import DEPRECATED_VALUE, deprecation_warning +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided +from ray.rllib.algorithms.impala.impala import IMPALA, IMPALAConfig +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.policy.policy import Policy +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics import ( + LAST_TARGET_UPDATE_TS, + LEARNER_STATS_KEY, + NUM_AGENT_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED, + NUM_TARGET_UPDATES, +) + +logger = logging.getLogger(__name__) + + +LEARNER_RESULTS_KL_KEY = "mean_kl_loss" +LEARNER_RESULTS_CURR_KL_COEFF_KEY = "curr_kl_coeff" +OLD_ACTION_DIST_KEY = "old_action_dist" + + +class APPOConfig(IMPALAConfig): + """Defines a configuration class from which an APPO Algorithm can be built. + + .. testcode:: + + from ray.rllib.algorithms.appo import APPOConfig + config = ( + APPOConfig() + .training(lr=0.01, grad_clip=30.0, train_batch_size_per_learner=50) + ) + config = config.learners(num_learners=1) + config = config.env_runners(num_env_runners=1) + config = config.environment("CartPole-v1") + + # Build an Algorithm object from the config and run 1 training iteration. + algo = config.build() + algo.train() + del algo + + .. testcode:: + + from ray.rllib.algorithms.appo import APPOConfig + from ray import tune + + config = APPOConfig() + # Update the config object. + config = config.training(lr=tune.grid_search([0.001,])) + # Set the config object's env. + config = config.environment(env="CartPole-v1") + # Use to_dict() to get the old-style python config dict when running with tune. + tune.Tuner( + "APPO", + run_config=tune.RunConfig( + stop={"training_iteration": 1}, + verbose=0, + ), + param_space=config.to_dict(), + + ).fit() + + .. testoutput:: + :hide: + + ... + """ + + def __init__(self, algo_class=None): + """Initializes a APPOConfig instance.""" + self.exploration_config = { + # The Exploration class to use. In the simplest case, this is the name + # (str) of any class present in the `rllib.utils.exploration` package. + # You can also provide the python class directly or the full location + # of your class (e.g. "ray.rllib.utils.exploration.epsilon_greedy. + # EpsilonGreedy"). + "type": "StochasticSampling", + # Add constructor kwargs here (if any). + } + + super().__init__(algo_class=algo_class or APPO) + + # fmt: off + # __sphinx_doc_begin__ + # APPO specific settings: + self.vtrace = True + self.use_gae = True + self.lambda_ = 1.0 + self.clip_param = 0.4 + self.use_kl_loss = False + self.kl_coeff = 1.0 + self.kl_target = 0.01 + self.target_worker_clipping = 2.0 + + # Circular replay buffer settings. + # Used in [1] for discrete action tasks: + # `circular_buffer_num_batches=4` and `circular_buffer_iterations_per_batch=2` + # For cont. action tasks: + # `circular_buffer_num_batches=16` and `circular_buffer_iterations_per_batch=20` + self.circular_buffer_num_batches = 4 + self.circular_buffer_iterations_per_batch = 2 + + # Override some of IMPALAConfig's default values with APPO-specific values. + self.num_env_runners = 2 + self.target_network_update_freq = 2 + self.broadcast_interval = 1 + self.grad_clip = 40.0 + # Note: Only when using enable_rl_module_and_learner=True can the clipping mode + # be configured by the user. On the old API stack, RLlib will always clip by + # global_norm, no matter the value of `grad_clip_by`. + self.grad_clip_by = "global_norm" + + self.opt_type = "adam" + self.lr = 0.0005 + self.decay = 0.99 + self.momentum = 0.0 + self.epsilon = 0.1 + self.vf_loss_coeff = 0.5 + self.entropy_coeff = 0.01 + self.tau = 1.0 + # __sphinx_doc_end__ + # fmt: on + + self.lr_schedule = None # @OldAPIStack + self.entropy_coeff_schedule = None # @OldAPIStack + self.num_gpus = 0 # @OldAPIStack + self.num_multi_gpu_tower_stacks = 1 # @OldAPIStack + self.minibatch_buffer_size = 1 # @OldAPIStack + self.replay_proportion = 0.0 # @OldAPIStack + self.replay_buffer_num_slots = 100 # @OldAPIStack + self.learner_queue_size = 16 # @OldAPIStack + self.learner_queue_timeout = 300 # @OldAPIStack + + # Deprecated keys. + self.target_update_frequency = DEPRECATED_VALUE + self.use_critic = DEPRECATED_VALUE + + @override(IMPALAConfig) + def training( + self, + *, + vtrace: Optional[bool] = NotProvided, + use_gae: Optional[bool] = NotProvided, + lambda_: Optional[float] = NotProvided, + clip_param: Optional[float] = NotProvided, + use_kl_loss: Optional[bool] = NotProvided, + kl_coeff: Optional[float] = NotProvided, + kl_target: Optional[float] = NotProvided, + target_network_update_freq: Optional[int] = NotProvided, + tau: Optional[float] = NotProvided, + target_worker_clipping: Optional[float] = NotProvided, + circular_buffer_num_batches: Optional[int] = NotProvided, + circular_buffer_iterations_per_batch: Optional[int] = NotProvided, + # Deprecated keys. + target_update_frequency=DEPRECATED_VALUE, + use_critic=DEPRECATED_VALUE, + **kwargs, + ) -> Self: + """Sets the training related configuration. + + Args: + vtrace: Whether to use V-trace weighted advantages. If false, PPO GAE + advantages will be used instead. + use_gae: If true, use the Generalized Advantage Estimator (GAE) + with a value function, see https://arxiv.org/pdf/1506.02438.pdf. + Only applies if vtrace=False. + lambda_: GAE (lambda) parameter. + clip_param: PPO surrogate slipping parameter. + use_kl_loss: Whether to use the KL-term in the loss function. + kl_coeff: Coefficient for weighting the KL-loss term. + kl_target: Target term for the KL-term to reach (via adjusting the + `kl_coeff` automatically). + target_network_update_freq: NOTE: This parameter is only applicable on + the new API stack. The frequency with which to update the target + policy network from the main trained policy network. The metric + used is `NUM_ENV_STEPS_TRAINED_LIFETIME` and the unit is `n` (see [1] + 4.1.1), where: `n = [circular_buffer_num_batches (N)] * + [circular_buffer_iterations_per_batch (K)] * [train batch size]` + For example, if you set `target_network_update_freq=2`, and N=4, K=2, + and `train_batch_size_per_learner=500`, then the target net is updated + every 2*4*2*500=8000 trained env steps (every 16 batch updates on each + learner). + The authors in [1] suggests that this setting is robust to a range of + choices (try values between 0.125 and 4). + target_network_update_freq: The frequency to update the target policy and + tune the kl loss coefficients that are used during training. After + setting this parameter, the algorithm waits for at least + `target_network_update_freq` number of environment samples to be trained + on before updating the target networks and tune the kl loss + coefficients. NOTE: This parameter is only applicable when using the + Learner API (enable_rl_module_and_learner=True). + tau: The factor by which to update the target policy network towards + the current policy network. Can range between 0 and 1. + e.g. updated_param = tau * current_param + (1 - tau) * target_param + target_worker_clipping: The maximum value for the target-worker-clipping + used for computing the IS ratio, described in [1] + IS = min(π(i) / π(target), ρ) * (π / π(i)) + circular_buffer_num_batches: The number of train batches that fit + into the circular buffer. Each such train batch can be sampled for + training max. `circular_buffer_iterations_per_batch` times. + circular_buffer_iterations_per_batch: The number of times any train + batch in the circular buffer can be sampled for training. A batch gets + evicted from the buffer either if it's the oldest batch in the buffer + and a new batch is added OR if the batch reaches this max. number of + being sampled. + + Returns: + This updated AlgorithmConfig object. + """ + if target_update_frequency != DEPRECATED_VALUE: + deprecation_warning( + old="target_update_frequency", + new="target_network_update_freq", + error=True, + ) + if use_critic != DEPRECATED_VALUE: + deprecation_warning( + old="use_critic", + help="`use_critic` no longer supported! APPO always uses a value " + "function (critic).", + error=True, + ) + + # Pass kwargs onto super's `training()` method. + super().training(**kwargs) + + if vtrace is not NotProvided: + self.vtrace = vtrace + if use_gae is not NotProvided: + self.use_gae = use_gae + if lambda_ is not NotProvided: + self.lambda_ = lambda_ + if clip_param is not NotProvided: + self.clip_param = clip_param + if use_kl_loss is not NotProvided: + self.use_kl_loss = use_kl_loss + if kl_coeff is not NotProvided: + self.kl_coeff = kl_coeff + if kl_target is not NotProvided: + self.kl_target = kl_target + if target_network_update_freq is not NotProvided: + self.target_network_update_freq = target_network_update_freq + if tau is not NotProvided: + self.tau = tau + if target_worker_clipping is not NotProvided: + self.target_worker_clipping = target_worker_clipping + if circular_buffer_num_batches is not NotProvided: + self.circular_buffer_num_batches = circular_buffer_num_batches + if circular_buffer_iterations_per_batch is not NotProvided: + self.circular_buffer_iterations_per_batch = ( + circular_buffer_iterations_per_batch + ) + + return self + + @override(IMPALAConfig) + def validate(self) -> None: + super().validate() + + # On new API stack, circular buffer should be used, not `minibatch_buffer_size`. + if self.enable_rl_module_and_learner: + if self.minibatch_buffer_size != 1 or self.replay_proportion != 0.0: + self._value_error( + "`minibatch_buffer_size/replay_proportion` not valid on new API " + "stack with APPO! " + "Use `circular_buffer_num_batches` for the number of train batches " + "in the circular buffer. To change the maximum number of times " + "any batch may be sampled, set " + "`circular_buffer_iterations_per_batch`." + ) + if self.num_multi_gpu_tower_stacks != 1: + self._value_error( + "`num_multi_gpu_tower_stacks` not supported on new API stack with " + "APPO! In order to train on multi-GPU, use " + "`config.learners(num_learners=[number of GPUs], " + "num_gpus_per_learner=1)`. To scale the throughput of batch-to-GPU-" + "pre-loading on each of your `Learners`, set " + "`num_gpu_loader_threads` to a higher number (recommended values: " + "1-8)." + ) + if self.learner_queue_size != 16: + self._value_error( + "`learner_queue_size` not supported on new API stack with " + "APPO! In order set the size of the circular buffer (which acts as " + "a 'learner queue'), use " + "`config.training(circular_buffer_num_batches=..)`. To change the " + "maximum number of times any batch may be sampled, set " + "`config.training(circular_buffer_iterations_per_batch=..)`." + ) + + @override(IMPALAConfig) + def get_default_learner_class(self): + if self.framework_str == "torch": + from ray.rllib.algorithms.appo.torch.appo_torch_learner import ( + APPOTorchLearner, + ) + + return APPOTorchLearner + elif self.framework_str in ["tf2", "tf"]: + raise ValueError( + "TensorFlow is no longer supported on the new API stack! " + "Use `framework='torch'`." + ) + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. " + "Use `framework='torch'`." + ) + + @override(IMPALAConfig) + def get_default_rl_module_spec(self) -> RLModuleSpec: + if self.framework_str == "torch": + from ray.rllib.algorithms.appo.torch.appo_torch_rl_module import ( + APPOTorchRLModule as RLModule, + ) + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. " + "Use either 'torch' or 'tf2'." + ) + + return RLModuleSpec(module_class=RLModule) + + @property + @override(AlgorithmConfig) + def _model_config_auto_includes(self): + return super()._model_config_auto_includes | {"vf_share_layers": False} + + +class APPO(IMPALA): + def __init__(self, config, *args, **kwargs): + """Initializes an APPO instance.""" + super().__init__(config, *args, **kwargs) + + # After init: Initialize target net. + + # TODO(avnishn): Does this need to happen in __init__? I think we can move it + # to setup() + if not self.config.enable_rl_module_and_learner: + self.env_runner.foreach_policy_to_train(lambda p, _: p.update_target()) + + @override(IMPALA) + def training_step(self) -> None: + if self.config.enable_rl_module_and_learner: + return super().training_step() + + train_results = super().training_step() + # Update the target network and the KL coefficient for the APPO-loss. + # The target network update frequency is calculated automatically by the product + # of `num_epochs` setting (usually 1 for APPO) and `minibatch_buffer_size`. + last_update = self._counters[LAST_TARGET_UPDATE_TS] + cur_ts = self._counters[ + ( + NUM_AGENT_STEPS_SAMPLED + if self.config.count_steps_by == "agent_steps" + else NUM_ENV_STEPS_SAMPLED + ) + ] + target_update_freq = self.config.num_epochs * self.config.minibatch_buffer_size + if cur_ts - last_update > target_update_freq: + self._counters[NUM_TARGET_UPDATES] += 1 + self._counters[LAST_TARGET_UPDATE_TS] = cur_ts + + # Update our target network. + self.env_runner.foreach_policy_to_train(lambda p, _: p.update_target()) + + # Also update the KL-coefficient for the APPO loss, if necessary. + if self.config.use_kl_loss: + + def update(pi, pi_id): + assert LEARNER_STATS_KEY not in train_results, ( + "{} should be nested under policy id key".format( + LEARNER_STATS_KEY + ), + train_results, + ) + if pi_id in train_results: + kl = train_results[pi_id][LEARNER_STATS_KEY].get("kl") + assert kl is not None, (train_results, pi_id) + # Make the actual `Policy.update_kl()` call. + pi.update_kl(kl) + else: + logger.warning("No data for {}, not updating kl".format(pi_id)) + + # Update KL on all trainable policies within the local (trainer) + # Worker. + self.env_runner.foreach_policy_to_train(update) + + return train_results + + @classmethod + @override(IMPALA) + def get_default_config(cls) -> APPOConfig: + return APPOConfig() + + @classmethod + @override(IMPALA) + def get_default_policy_class( + cls, config: AlgorithmConfig + ) -> Optional[Type[Policy]]: + if config["framework"] == "torch": + from ray.rllib.algorithms.appo.appo_torch_policy import APPOTorchPolicy + + return APPOTorchPolicy + elif config["framework"] == "tf": + if config.enable_rl_module_and_learner: + raise ValueError( + "RLlib's RLModule and Learner API is not supported for" + " tf1. Use " + "framework='tf2' instead." + ) + from ray.rllib.algorithms.appo.appo_tf_policy import APPOTF1Policy + + return APPOTF1Policy + else: + from ray.rllib.algorithms.appo.appo_tf_policy import APPOTF2Policy + + return APPOTF2Policy diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..235bc823209f8a47c9fe189c0c9cd5b6b4804170 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_learner.py @@ -0,0 +1,147 @@ +import abc +from typing import Any, Dict, Optional + +from ray.rllib.algorithms.appo.appo import APPOConfig +from ray.rllib.algorithms.appo.utils import CircularBuffer +from ray.rllib.algorithms.impala.impala_learner import IMPALALearner +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.learner.utils import update_target_network +from ray.rllib.core.rl_module.apis import TargetNetworkAPI, ValueFunctionAPI +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.utils.annotations import override +from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict +from ray.rllib.utils.metrics import ( + LAST_TARGET_UPDATE_TS, + NUM_ENV_STEPS_TRAINED_LIFETIME, + NUM_MODULE_STEPS_TRAINED, + NUM_TARGET_UPDATES, +) +from ray.rllib.utils.schedules.scheduler import Scheduler +from ray.rllib.utils.typing import ModuleID, ShouldModuleBeUpdatedFn + + +class APPOLearner(IMPALALearner): + """Adds KL coeff updates via `after_gradient_based_update()` to IMPALA logic. + + Framework-specific subclasses must override `_update_module_kl_coeff()`. + """ + + @override(IMPALALearner) + def build(self): + self._learner_thread_in_queue = CircularBuffer( + num_batches=self.config.circular_buffer_num_batches, + iterations_per_batch=self.config.circular_buffer_iterations_per_batch, + ) + + super().build() + + # Make target networks. + self.module.foreach_module( + lambda mid, mod: ( + mod.make_target_networks() + if isinstance(mod, TargetNetworkAPI) + else None + ) + ) + + # The current kl coefficients per module as (framework specific) tensor + # variables. + self.curr_kl_coeffs_per_module: LambdaDefaultDict[ + ModuleID, Scheduler + ] = LambdaDefaultDict( + lambda module_id: self._get_tensor_variable( + self.config.get_config_for_module(module_id).kl_coeff + ) + ) + + @override(Learner) + def add_module( + self, + *, + module_id: ModuleID, + module_spec: RLModuleSpec, + config_overrides: Optional[Dict] = None, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + ) -> MultiRLModuleSpec: + marl_spec = super().add_module( + module_id=module_id, + module_spec=module_spec, + config_overrides=config_overrides, + new_should_module_be_updated=new_should_module_be_updated, + ) + # Create target networks for added Module, if applicable. + if isinstance(self.module[module_id].unwrapped(), TargetNetworkAPI): + self.module[module_id].unwrapped().make_target_networks() + return marl_spec + + @override(IMPALALearner) + def remove_module(self, module_id: str) -> MultiRLModuleSpec: + marl_spec = super().remove_module(module_id) + self.curr_kl_coeffs_per_module.pop(module_id) + return marl_spec + + @override(Learner) + def after_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None: + """Updates the target Q Networks.""" + super().after_gradient_based_update(timesteps=timesteps) + + # TODO (sven): Maybe we should have a `after_gradient_based_update` + # method per module? + curr_timestep = timesteps.get(NUM_ENV_STEPS_TRAINED_LIFETIME, 0) + for module_id, module in self.module._rl_modules.items(): + config = self.config.get_config_for_module(module_id) + + last_update_ts_key = (module_id, LAST_TARGET_UPDATE_TS) + if isinstance(module.unwrapped(), TargetNetworkAPI) and ( + curr_timestep - self.metrics.peek(last_update_ts_key, default=0) + >= ( + config.target_network_update_freq + * config.circular_buffer_num_batches + * config.circular_buffer_iterations_per_batch + * config.train_batch_size_per_learner + ) + ): + for ( + main_net, + target_net, + ) in module.unwrapped().get_target_network_pairs(): + update_target_network( + main_net=main_net, + target_net=target_net, + tau=config.tau, + ) + # Increase lifetime target network update counter by one. + self.metrics.log_value((module_id, NUM_TARGET_UPDATES), 1, reduce="sum") + # Update the (single-value -> window=1) last updated timestep metric. + self.metrics.log_value(last_update_ts_key, curr_timestep, window=1) + + if ( + config.use_kl_loss + and self.metrics.peek((module_id, NUM_MODULE_STEPS_TRAINED), default=0) + > 0 + ): + self._update_module_kl_coeff(module_id=module_id, config=config) + + @classmethod + @override(Learner) + def rl_module_required_apis(cls) -> list[type]: + # In order for a PPOLearner to update an RLModule, it must implement the + # following APIs: + return [TargetNetworkAPI, ValueFunctionAPI] + + @abc.abstractmethod + def _update_module_kl_coeff(self, module_id: ModuleID, config: APPOConfig) -> None: + """Dynamically update the KL loss coefficients of each module. + + The update is completed using the mean KL divergence between the action + distributions current policy and old policy of each module. That action + distribution is computed during the most recent update/call to `compute_loss`. + + Args: + module_id: The module whose KL loss coefficient to update. + config: The AlgorithmConfig specific to the given `module_id`. + """ + + +AppoLearner = APPOLearner diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..de9b862a92abd122ccd95d9f698e47a4913807fc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_rl_module.py @@ -0,0 +1,11 @@ +# Backward compat import. +from ray.rllib.algorithms.appo.default_appo_rl_module import ( # noqa + DefaultAPPORLModule as APPORLModule, +) +from ray._common.deprecation import deprecation_warning + +deprecation_warning( + old="ray.rllib.algorithms.appo.appo_rl_module.APPORLModule", + new="ray.rllib.algorithms.appo.default_appo_rl_module.DefaultAPPORLModule", + error=False, +) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_tf_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_tf_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..eab4bfefeb2e2b0577550c89e3737380138bc1c7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_tf_policy.py @@ -0,0 +1,393 @@ +""" +TensorFlow policy class used for APPO. + +Adapted from VTraceTFPolicy to use the PPO surrogate loss. +Keep in sync with changes to VTraceTFPolicy. +""" + +import logging +from typing import Dict, List, Optional, Type, Union + +import gymnasium as gym +import numpy as np + +from ray.rllib.algorithms.appo.utils import make_appo_models +from ray.rllib.algorithms.impala import vtrace_tf as vtrace +from ray.rllib.algorithms.impala.impala_tf_policy import ( + VTraceClipGradients, + VTraceOptimizer, + _make_time_major, +) +from ray.rllib.evaluation.postprocessing import ( + Postprocessing, + compute_bootstrap_value, + compute_gae_for_sample_batch, +) +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import Categorical, TFActionDistribution +from ray.rllib.policy.dynamic_tf_policy_v2 import DynamicTFPolicyV2 +from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_mixins import ( + EntropyCoeffSchedule, + GradStatsMixin, + KLCoeffMixin, + LearningRateSchedule, + TargetNetworkMixin, + ValueNetworkMixin, +) +from ray.rllib.utils.annotations import ( + override, +) +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.tf_utils import explained_variance +from ray.rllib.utils.typing import TensorType + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +# TODO (sven): Deprecate once APPO and IMPALA fully on RLModules/Learner APIs. +def get_appo_tf_policy(name: str, base: type) -> type: + """Construct an APPOTFPolicy inheriting either dynamic or eager base policies. + + Args: + base: Base class for this policy. DynamicTFPolicyV2 or EagerTFPolicyV2. + + Returns: + A TF Policy to be used with Impala. + """ + + class APPOTFPolicy( + VTraceClipGradients, + VTraceOptimizer, + LearningRateSchedule, + KLCoeffMixin, + EntropyCoeffSchedule, + ValueNetworkMixin, + TargetNetworkMixin, + GradStatsMixin, + base, + ): + def __init__( + self, + observation_space, + action_space, + config, + existing_model=None, + existing_inputs=None, + ): + # First thing first, enable eager execution if necessary. + base.enable_eager_execution_if_necessary() + + # Although this is a no-op, we call __init__ here to make it clear + # that base.__init__ will use the make_model() call. + VTraceClipGradients.__init__(self) + VTraceOptimizer.__init__(self) + + # Initialize base class. + base.__init__( + self, + observation_space, + action_space, + config, + existing_inputs=existing_inputs, + existing_model=existing_model, + ) + + # TF LearningRateSchedule depends on self.framework, so initialize + # after base.__init__() is called. + LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"]) + EntropyCoeffSchedule.__init__( + self, config["entropy_coeff"], config["entropy_coeff_schedule"] + ) + ValueNetworkMixin.__init__(self, config) + KLCoeffMixin.__init__(self, config) + + GradStatsMixin.__init__(self) + + # Note: this is a bit ugly, but loss and optimizer initialization must + # happen after all the MixIns are initialized. + self.maybe_initialize_optimizer_and_loss() + + # Initiate TargetNetwork ops after loss initialization. + TargetNetworkMixin.__init__(self) + + @override(base) + def make_model(self) -> ModelV2: + return make_appo_models(self) + + @override(base) + def loss( + self, + model: Union[ModelV2, "tf.keras.Model"], + dist_class: Type[TFActionDistribution], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + model_out, _ = model(train_batch) + action_dist = dist_class(model_out, model) + + if isinstance(self.action_space, gym.spaces.Discrete): + is_multidiscrete = False + output_hidden_shape = [self.action_space.n] + elif isinstance(self.action_space, gym.spaces.multi_discrete.MultiDiscrete): + is_multidiscrete = True + output_hidden_shape = self.action_space.nvec.astype(np.int32) + else: + is_multidiscrete = False + output_hidden_shape = 1 + + def make_time_major(*args, **kw): + return _make_time_major( + self, train_batch.get(SampleBatch.SEQ_LENS), *args, **kw + ) + + actions = train_batch[SampleBatch.ACTIONS] + dones = train_batch[SampleBatch.TERMINATEDS] + rewards = train_batch[SampleBatch.REWARDS] + behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS] + + target_model_out, _ = self.target_model(train_batch) + prev_action_dist = dist_class(behaviour_logits, self.model) + values = self.model.value_function() + values_time_major = make_time_major(values) + bootstrap_values_time_major = make_time_major( + train_batch[SampleBatch.VALUES_BOOTSTRAPPED] + ) + bootstrap_value = bootstrap_values_time_major[-1] + + if self.is_recurrent(): + max_seq_len = tf.reduce_max(train_batch[SampleBatch.SEQ_LENS]) + mask = tf.sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len) + mask = tf.reshape(mask, [-1]) + mask = make_time_major(mask) + + def reduce_mean_valid(t): + return tf.reduce_mean(tf.boolean_mask(t, mask)) + + else: + reduce_mean_valid = tf.reduce_mean + + if self.config["vtrace"]: + logger.debug("Using V-Trace surrogate loss (vtrace=True)") + + # Prepare actions for loss. + loss_actions = ( + actions if is_multidiscrete else tf.expand_dims(actions, axis=1) + ) + + old_policy_behaviour_logits = tf.stop_gradient(target_model_out) + old_policy_action_dist = dist_class(old_policy_behaviour_logits, model) + + # Prepare KL for Loss + mean_kl = make_time_major(old_policy_action_dist.multi_kl(action_dist)) + + unpacked_behaviour_logits = tf.split( + behaviour_logits, output_hidden_shape, axis=1 + ) + unpacked_old_policy_behaviour_logits = tf.split( + old_policy_behaviour_logits, output_hidden_shape, axis=1 + ) + + # Compute vtrace on the CPU for better perf. + with tf.device("/cpu:0"): + vtrace_returns = vtrace.multi_from_logits( + behaviour_policy_logits=make_time_major( + unpacked_behaviour_logits + ), + target_policy_logits=make_time_major( + unpacked_old_policy_behaviour_logits + ), + actions=tf.unstack(make_time_major(loss_actions), axis=2), + discounts=tf.cast( + ~make_time_major(tf.cast(dones, tf.bool)), + tf.float32, + ) + * self.config["gamma"], + rewards=make_time_major(rewards), + values=values_time_major, + bootstrap_value=bootstrap_value, + dist_class=Categorical if is_multidiscrete else dist_class, + model=model, + clip_rho_threshold=tf.cast( + self.config["vtrace_clip_rho_threshold"], tf.float32 + ), + clip_pg_rho_threshold=tf.cast( + self.config["vtrace_clip_pg_rho_threshold"], tf.float32 + ), + ) + + actions_logp = make_time_major(action_dist.logp(actions)) + prev_actions_logp = make_time_major(prev_action_dist.logp(actions)) + old_policy_actions_logp = make_time_major( + old_policy_action_dist.logp(actions) + ) + + is_ratio = tf.clip_by_value( + tf.math.exp(prev_actions_logp - old_policy_actions_logp), 0.0, 2.0 + ) + logp_ratio = is_ratio * tf.exp(actions_logp - prev_actions_logp) + self._is_ratio = is_ratio + + advantages = vtrace_returns.pg_advantages + surrogate_loss = tf.minimum( + advantages * logp_ratio, + advantages + * tf.clip_by_value( + logp_ratio, + 1 - self.config["clip_param"], + 1 + self.config["clip_param"], + ), + ) + + action_kl = ( + tf.reduce_mean(mean_kl, axis=0) if is_multidiscrete else mean_kl + ) + mean_kl_loss = reduce_mean_valid(action_kl) + mean_policy_loss = -reduce_mean_valid(surrogate_loss) + + # The value function loss. + value_targets = vtrace_returns.vs + delta = values_time_major - value_targets + mean_vf_loss = 0.5 * reduce_mean_valid(tf.math.square(delta)) + + # The entropy loss. + actions_entropy = make_time_major(action_dist.multi_entropy()) + mean_entropy = reduce_mean_valid(actions_entropy) + + else: + logger.debug("Using PPO surrogate loss (vtrace=False)") + + # Prepare KL for Loss + mean_kl = make_time_major(prev_action_dist.multi_kl(action_dist)) + + logp_ratio = tf.math.exp( + make_time_major(action_dist.logp(actions)) + - make_time_major(prev_action_dist.logp(actions)) + ) + + advantages = make_time_major(train_batch[Postprocessing.ADVANTAGES]) + surrogate_loss = tf.minimum( + advantages * logp_ratio, + advantages + * tf.clip_by_value( + logp_ratio, + 1 - self.config["clip_param"], + 1 + self.config["clip_param"], + ), + ) + + action_kl = ( + tf.reduce_mean(mean_kl, axis=0) if is_multidiscrete else mean_kl + ) + mean_kl_loss = reduce_mean_valid(action_kl) + mean_policy_loss = -reduce_mean_valid(surrogate_loss) + + # The value function loss. + value_targets = make_time_major( + train_batch[Postprocessing.VALUE_TARGETS] + ) + delta = values_time_major - value_targets + mean_vf_loss = 0.5 * reduce_mean_valid(tf.math.square(delta)) + + # The entropy loss. + mean_entropy = reduce_mean_valid( + make_time_major(action_dist.multi_entropy()) + ) + + # The summed weighted loss. + total_loss = mean_policy_loss - mean_entropy * self.entropy_coeff + # Optional KL loss. + if self.config["use_kl_loss"]: + total_loss += self.kl_coeff * mean_kl_loss + # Optional vf loss (or in a separate term due to separate + # optimizers/networks). + loss_wo_vf = total_loss + if not self.config["_separate_vf_optimizer"]: + total_loss += mean_vf_loss * self.config["vf_loss_coeff"] + + # Store stats in policy for stats_fn. + self._total_loss = total_loss + self._loss_wo_vf = loss_wo_vf + self._mean_policy_loss = mean_policy_loss + # Backward compatibility: Deprecate policy._mean_kl. + self._mean_kl_loss = self._mean_kl = mean_kl_loss + self._mean_vf_loss = mean_vf_loss + self._mean_entropy = mean_entropy + self._value_targets = value_targets + + # Return one total loss or two losses: vf vs rest (policy + kl). + if self.config["_separate_vf_optimizer"]: + return loss_wo_vf, mean_vf_loss + else: + return total_loss + + @override(base) + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + values_batched = _make_time_major( + self, + train_batch.get(SampleBatch.SEQ_LENS), + self.model.value_function(), + ) + + stats_dict = { + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "total_loss": self._total_loss, + "policy_loss": self._mean_policy_loss, + "entropy": self._mean_entropy, + "var_gnorm": tf.linalg.global_norm(self.model.trainable_variables()), + "vf_loss": self._mean_vf_loss, + "vf_explained_var": explained_variance( + tf.reshape(self._value_targets, [-1]), + tf.reshape(values_batched, [-1]), + ), + "entropy_coeff": tf.cast(self.entropy_coeff, tf.float64), + } + + if self.config["vtrace"]: + is_stat_mean, is_stat_var = tf.nn.moments(self._is_ratio, [0, 1]) + stats_dict["mean_IS"] = is_stat_mean + stats_dict["var_IS"] = is_stat_var + + if self.config["use_kl_loss"]: + stats_dict["kl"] = self._mean_kl_loss + stats_dict["KL_Coeff"] = self.kl_coeff + + return stats_dict + + @override(base) + def postprocess_trajectory( + self, + sample_batch: SampleBatch, + other_agent_batches: Optional[SampleBatch] = None, + episode=None, + ): + # Call super's postprocess_trajectory first. + # sample_batch = super().postprocess_trajectory( + # sample_batch, other_agent_batches, episode + # ) + + if not self.config["vtrace"]: + sample_batch = compute_gae_for_sample_batch( + self, sample_batch, other_agent_batches, episode + ) + else: + # Add the Columns.VALUES_BOOTSTRAPPED column, which we'll need + # inside the loss for vtrace calculations. + sample_batch = compute_bootstrap_value(sample_batch, self) + + return sample_batch + + @override(base) + def get_batch_divisibility_req(self) -> int: + return self.config["rollout_fragment_length"] + + APPOTFPolicy.__name__ = name + APPOTFPolicy.__qualname__ = name + + return APPOTFPolicy + + +APPOTF1Policy = get_appo_tf_policy("APPOTF1Policy", DynamicTFPolicyV2) +APPOTF2Policy = get_appo_tf_policy("APPOTF2Policy", EagerTFPolicyV2) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_torch_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_torch_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..f150c6761cac2667be39503161ee4d7f2dd88420 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/appo_torch_policy.py @@ -0,0 +1,413 @@ +""" +PyTorch policy class used for APPO. + +Adapted from VTraceTFPolicy to use the PPO surrogate loss. +Keep in sync with changes to VTraceTFPolicy. +""" + +import logging +from typing import Any, Dict, List, Optional, Type, Union + +import gymnasium as gym +import numpy as np + +import ray +import ray.rllib.algorithms.impala.vtrace_torch as vtrace +from ray.rllib.algorithms.appo.utils import make_appo_models +from ray.rllib.algorithms.impala.impala_torch_policy import ( + VTraceOptimizer, + make_time_major, +) +from ray.rllib.evaluation.postprocessing import ( + Postprocessing, + compute_bootstrap_value, + compute_gae_for_sample_batch, +) +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.torch_action_dist import ( + TorchCategorical, + TorchDistributionWrapper, +) +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_mixins import ( + EntropyCoeffSchedule, + KLCoeffMixin, + LearningRateSchedule, + TargetNetworkMixin, + ValueNetworkMixin, +) +from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.torch_utils import ( + apply_grad_clipping, + explained_variance, + global_norm, + sequence_mask, +) +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +# TODO (sven): Deprecate once APPO and IMPALA fully on RLModules/Learner APIs. +class APPOTorchPolicy( + VTraceOptimizer, + LearningRateSchedule, + EntropyCoeffSchedule, + KLCoeffMixin, + ValueNetworkMixin, + TargetNetworkMixin, + TorchPolicyV2, +): + """PyTorch policy class used with APPO.""" + + def __init__(self, observation_space, action_space, config): + config = dict(ray.rllib.algorithms.appo.appo.APPOConfig().to_dict(), **config) + config["enable_rl_module_and_learner"] = False + config["enable_env_runner_and_connector_v2"] = False + + # Although this is a no-op, we call __init__ here to make it clear + # that base.__init__ will use the make_model() call. + VTraceOptimizer.__init__(self) + + lr_schedule_additional_args = [] + if config.get("_separate_vf_optimizer"): + lr_schedule_additional_args = ( + [config["_lr_vf"][0][1], config["_lr_vf"]] + if isinstance(config["_lr_vf"], (list, tuple)) + else [config["_lr_vf"], None] + ) + LearningRateSchedule.__init__( + self, config["lr"], config["lr_schedule"], *lr_schedule_additional_args + ) + + TorchPolicyV2.__init__( + self, + observation_space, + action_space, + config, + max_seq_len=config["model"]["max_seq_len"], + ) + + EntropyCoeffSchedule.__init__( + self, config["entropy_coeff"], config["entropy_coeff_schedule"] + ) + ValueNetworkMixin.__init__(self, config) + KLCoeffMixin.__init__(self, config) + + self._initialize_loss_from_dummy_batch() + + # Initiate TargetNetwork ops after loss initialization. + TargetNetworkMixin.__init__(self) + + @override(TorchPolicyV2) + def init_view_requirements(self): + self.view_requirements = self._get_default_view_requirements() + + @override(TorchPolicyV2) + def make_model(self) -> ModelV2: + return make_appo_models(self) + + @override(TorchPolicyV2) + def loss( + self, + model: ModelV2, + dist_class: Type[ActionDistribution], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + """Constructs the loss for APPO. + + With IS modifications and V-trace for Advantage Estimation. + + Args: + model (ModelV2): The Model to calculate the loss for. + dist_class (Type[ActionDistribution]): The action distr. class. + train_batch: The training data. + + Returns: + Union[TensorType, List[TensorType]]: A single loss tensor or a list + of loss tensors. + """ + target_model = self.target_models[model] + + model_out, _ = model(train_batch) + action_dist = dist_class(model_out, model) + + if isinstance(self.action_space, gym.spaces.Discrete): + is_multidiscrete = False + output_hidden_shape = [self.action_space.n] + elif isinstance(self.action_space, gym.spaces.multi_discrete.MultiDiscrete): + is_multidiscrete = True + output_hidden_shape = self.action_space.nvec.astype(np.int32) + else: + is_multidiscrete = False + output_hidden_shape = 1 + + def _make_time_major(*args, **kwargs): + return make_time_major( + self, train_batch.get(SampleBatch.SEQ_LENS), *args, **kwargs + ) + + actions = train_batch[SampleBatch.ACTIONS] + dones = train_batch[SampleBatch.TERMINATEDS] + rewards = train_batch[SampleBatch.REWARDS] + behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS] + + target_model_out, _ = target_model(train_batch) + + prev_action_dist = dist_class(behaviour_logits, model) + values = model.value_function() + values_time_major = _make_time_major(values) + bootstrap_values_time_major = _make_time_major( + train_batch[SampleBatch.VALUES_BOOTSTRAPPED] + ) + bootstrap_value = bootstrap_values_time_major[-1] + + if self.is_recurrent(): + max_seq_len = torch.max(train_batch[SampleBatch.SEQ_LENS]) + mask = sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len) + mask = torch.reshape(mask, [-1]) + mask = _make_time_major(mask) + num_valid = torch.sum(mask) + + def reduce_mean_valid(t): + return torch.sum(t[mask]) / num_valid + + else: + reduce_mean_valid = torch.mean + + if self.config["vtrace"]: + logger.debug("Using V-Trace surrogate loss (vtrace=True)") + + old_policy_behaviour_logits = target_model_out.detach() + old_policy_action_dist = dist_class(old_policy_behaviour_logits, model) + + if isinstance(output_hidden_shape, (list, tuple, np.ndarray)): + unpacked_behaviour_logits = torch.split( + behaviour_logits, list(output_hidden_shape), dim=1 + ) + unpacked_old_policy_behaviour_logits = torch.split( + old_policy_behaviour_logits, list(output_hidden_shape), dim=1 + ) + else: + unpacked_behaviour_logits = torch.chunk( + behaviour_logits, output_hidden_shape, dim=1 + ) + unpacked_old_policy_behaviour_logits = torch.chunk( + old_policy_behaviour_logits, output_hidden_shape, dim=1 + ) + + # Prepare actions for loss. + loss_actions = ( + actions if is_multidiscrete else torch.unsqueeze(actions, dim=1) + ) + + # Prepare KL for loss. + action_kl = _make_time_major(old_policy_action_dist.kl(action_dist)) + + # Compute vtrace on the CPU for better perf. + vtrace_returns = vtrace.multi_from_logits( + behaviour_policy_logits=_make_time_major(unpacked_behaviour_logits), + target_policy_logits=_make_time_major( + unpacked_old_policy_behaviour_logits + ), + actions=torch.unbind(_make_time_major(loss_actions), dim=2), + discounts=(1.0 - _make_time_major(dones).float()) + * self.config["gamma"], + rewards=_make_time_major(rewards), + values=values_time_major, + bootstrap_value=bootstrap_value, + dist_class=TorchCategorical if is_multidiscrete else dist_class, + model=model, + clip_rho_threshold=self.config["vtrace_clip_rho_threshold"], + clip_pg_rho_threshold=self.config["vtrace_clip_pg_rho_threshold"], + ) + + actions_logp = _make_time_major(action_dist.logp(actions)) + prev_actions_logp = _make_time_major(prev_action_dist.logp(actions)) + old_policy_actions_logp = _make_time_major( + old_policy_action_dist.logp(actions) + ) + is_ratio = torch.clamp( + torch.exp(prev_actions_logp - old_policy_actions_logp), 0.0, 2.0 + ) + logp_ratio = is_ratio * torch.exp(actions_logp - prev_actions_logp) + self._is_ratio = is_ratio + + advantages = vtrace_returns.pg_advantages.to(logp_ratio.device) + surrogate_loss = torch.min( + advantages * logp_ratio, + advantages + * torch.clamp( + logp_ratio, + 1 - self.config["clip_param"], + 1 + self.config["clip_param"], + ), + ) + + mean_kl_loss = reduce_mean_valid(action_kl) + mean_policy_loss = -reduce_mean_valid(surrogate_loss) + + # The value function loss. + value_targets = vtrace_returns.vs.to(values_time_major.device) + delta = values_time_major - value_targets + mean_vf_loss = 0.5 * reduce_mean_valid(torch.pow(delta, 2.0)) + + # The entropy loss. + mean_entropy = reduce_mean_valid(_make_time_major(action_dist.entropy())) + + else: + logger.debug("Using PPO surrogate loss (vtrace=False)") + + # Prepare KL for Loss + action_kl = _make_time_major(prev_action_dist.kl(action_dist)) + + actions_logp = _make_time_major(action_dist.logp(actions)) + prev_actions_logp = _make_time_major(prev_action_dist.logp(actions)) + logp_ratio = torch.exp(actions_logp - prev_actions_logp) + + advantages = _make_time_major(train_batch[Postprocessing.ADVANTAGES]) + surrogate_loss = torch.min( + advantages * logp_ratio, + advantages + * torch.clamp( + logp_ratio, + 1 - self.config["clip_param"], + 1 + self.config["clip_param"], + ), + ) + + mean_kl_loss = reduce_mean_valid(action_kl) + mean_policy_loss = -reduce_mean_valid(surrogate_loss) + + # The value function loss. + value_targets = _make_time_major(train_batch[Postprocessing.VALUE_TARGETS]) + delta = values_time_major - value_targets + mean_vf_loss = 0.5 * reduce_mean_valid(torch.pow(delta, 2.0)) + + # The entropy loss. + mean_entropy = reduce_mean_valid(_make_time_major(action_dist.entropy())) + + # The summed weighted loss. + total_loss = mean_policy_loss - mean_entropy * self.entropy_coeff + # Optional additional KL Loss + if self.config["use_kl_loss"]: + total_loss += self.kl_coeff * mean_kl_loss + + # Optional vf loss (or in a separate term due to separate + # optimizers/networks). + loss_wo_vf = total_loss + if not self.config["_separate_vf_optimizer"]: + total_loss += mean_vf_loss * self.config["vf_loss_coeff"] + + # Store values for stats function in model (tower), such that for + # multi-GPU, we do not override them during the parallel loss phase. + model.tower_stats["total_loss"] = total_loss + model.tower_stats["mean_policy_loss"] = mean_policy_loss + model.tower_stats["mean_kl_loss"] = mean_kl_loss + model.tower_stats["mean_vf_loss"] = mean_vf_loss + model.tower_stats["mean_entropy"] = mean_entropy + model.tower_stats["value_targets"] = value_targets + model.tower_stats["vf_explained_var"] = explained_variance( + torch.reshape(value_targets, [-1]), + torch.reshape(values_time_major, [-1]), + ) + + # Return one total loss or two losses: vf vs rest (policy + kl). + if self.config["_separate_vf_optimizer"]: + return loss_wo_vf, mean_vf_loss + else: + return total_loss + + @override(TorchPolicyV2) + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + """Stats function for APPO. Returns a dict with important loss stats. + + Args: + policy: The Policy to generate stats for. + train_batch: The SampleBatch (already) used for training. + + Returns: + Dict[str, TensorType]: The stats dict. + """ + stats_dict = { + "cur_lr": self.cur_lr, + "total_loss": torch.mean(torch.stack(self.get_tower_stats("total_loss"))), + "policy_loss": torch.mean( + torch.stack(self.get_tower_stats("mean_policy_loss")) + ), + "entropy": torch.mean(torch.stack(self.get_tower_stats("mean_entropy"))), + "entropy_coeff": self.entropy_coeff, + "var_gnorm": global_norm(self.model.trainable_variables()), + "vf_loss": torch.mean(torch.stack(self.get_tower_stats("mean_vf_loss"))), + "vf_explained_var": torch.mean( + torch.stack(self.get_tower_stats("vf_explained_var")) + ), + } + + if self.config["vtrace"]: + is_stat_mean = torch.mean(self._is_ratio, [0, 1]) + is_stat_var = torch.var(self._is_ratio, [0, 1]) + stats_dict["mean_IS"] = is_stat_mean + stats_dict["var_IS"] = is_stat_var + + if self.config["use_kl_loss"]: + stats_dict["kl"] = torch.mean( + torch.stack(self.get_tower_stats("mean_kl_loss")) + ) + stats_dict["KL_Coeff"] = self.kl_coeff + + return convert_to_numpy(stats_dict) + + @override(TorchPolicyV2) + def extra_action_out( + self, + input_dict: Dict[str, TensorType], + state_batches: List[TensorType], + model: TorchModelV2, + action_dist: TorchDistributionWrapper, + ) -> Dict[str, TensorType]: + return {SampleBatch.VF_PREDS: model.value_function()} + + @override(TorchPolicyV2) + def postprocess_trajectory( + self, + sample_batch: SampleBatch, + other_agent_batches: Optional[Dict[Any, SampleBatch]] = None, + episode=None, + ): + # Call super's postprocess_trajectory first. + # sample_batch = super().postprocess_trajectory( + # sample_batch, other_agent_batches, episode + # ) + + # Do all post-processing always with no_grad(). + # Not using this here will introduce a memory leak + # in torch (issue #6962). + with torch.no_grad(): + if not self.config["vtrace"]: + sample_batch = compute_gae_for_sample_batch( + self, sample_batch, other_agent_batches, episode + ) + else: + # Add the SampleBatch.VALUES_BOOTSTRAPPED column, which we'll need + # inside the loss for vtrace calculations. + sample_batch = compute_bootstrap_value(sample_batch, self) + + return sample_batch + + @override(TorchPolicyV2) + def extra_grad_process( + self, optimizer: "torch.optim.Optimizer", loss: TensorType + ) -> Dict[str, TensorType]: + return apply_grad_clipping(self, optimizer, loss) + + @override(TorchPolicyV2) + def get_batch_divisibility_req(self) -> int: + return self.config["rollout_fragment_length"] diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/default_appo_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/default_appo_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..e6eb13d23bf1a92672fbdc210cf1f38cecce91d7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/default_appo_rl_module.py @@ -0,0 +1,56 @@ +import abc +from typing import Any, Dict, List, Tuple + +from ray.rllib.algorithms.ppo.default_ppo_rl_module import DefaultPPORLModule +from ray.rllib.core.learner.utils import make_target_network +from ray.rllib.core.models.base import ACTOR, ENCODER_OUT +from ray.rllib.core.rl_module.apis import ( + TARGET_NETWORK_ACTION_DIST_INPUTS, + TargetNetworkAPI, +) +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.typing import NetworkType +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class DefaultAPPORLModule(DefaultPPORLModule, TargetNetworkAPI, abc.ABC): + """Default RLModule used by APPO, if user does not specify a custom RLModule. + + Users who want to train their RLModules with APPO may implement any RLModule + (or TorchRLModule) subclass as long as the custom class also implements the + `ValueFunctionAPI` (see ray.rllib.core.rl_module.apis.value_function_api.py) + and the `TargetNetworkAPI` (see + ray.rllib.core.rl_module.apis.target_network_api.py). + """ + + @override(TargetNetworkAPI) + def make_target_networks(self): + self._old_encoder = make_target_network(self.encoder) + self._old_pi = make_target_network(self.pi) + + @override(TargetNetworkAPI) + def get_target_network_pairs(self) -> List[Tuple[NetworkType, NetworkType]]: + return [ + (self.encoder, self._old_encoder), + (self.pi, self._old_pi), + ] + + @override(TargetNetworkAPI) + def forward_target(self, batch: Dict[str, Any]) -> Dict[str, Any]: + old_pi_inputs_encoded = self._old_encoder(batch)[ENCODER_OUT][ACTOR] + old_action_dist_logits = self._old_pi(old_pi_inputs_encoded) + return {TARGET_NETWORK_ACTION_DIST_INPUTS: old_action_dist_logits} + + @OverrideToImplementCustomLogic_CallToSuperRecommended + @override(DefaultPPORLModule) + def get_non_inference_attributes(self) -> List[str]: + # Get the NON inference-only attributes from the parent class. + ret = super().get_non_inference_attributes() + # Add the two (APPO) target networks to it (NOT needed in + # inference-only mode). + ret += ["_old_encoder", "_old_pi"] + return ret diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ef56a1f36fed7f1a4014f69e31a507e2053bfcd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/appo_torch_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/appo_torch_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b45ae05ae3eb681a6a82accbe0a7a1344b8822fe Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/appo_torch_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/appo_torch_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/appo_torch_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3c4cbb87cb7743166c661ab9d9287eaf89a3d33 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/appo_torch_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/default_appo_torch_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/default_appo_torch_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..28f0ff4b6fa95ea56db13dced176c33063cffb59 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/__pycache__/default_appo_torch_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/appo_torch_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/appo_torch_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..9e3bbfca3b925d80bfff3b93781cc3677d337cd9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/appo_torch_learner.py @@ -0,0 +1,234 @@ +"""Asynchronous Proximal Policy Optimization (APPO) + +The algorithm is described in [1] (under the name of "IMPACT"): + +Detailed documentation: +https://docs.ray.io/en/master/rllib-algorithms.html#appo + +[1] IMPACT: Importance Weighted Asynchronous Architectures with Clipped Target Networks. +Luo et al. 2020 +https://arxiv.org/pdf/1912.00167 +""" +from typing import Dict + +from ray.rllib.algorithms.appo.appo import ( + LEARNER_RESULTS_CURR_KL_COEFF_KEY, + LEARNER_RESULTS_KL_KEY, + APPOConfig, +) +from ray.rllib.algorithms.appo.appo_learner import APPOLearner +from ray.rllib.algorithms.impala.torch.impala_torch_learner import IMPALATorchLearner +from ray.rllib.algorithms.impala.torch.vtrace_torch_v2 import ( + make_time_major, + vtrace_torch, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.learner import ENTROPY_KEY, POLICY_LOSS_KEY, VF_LOSS_KEY +from ray.rllib.core.rl_module.apis import ( + TARGET_NETWORK_ACTION_DIST_INPUTS, + TargetNetworkAPI, + ValueFunctionAPI, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import ModuleID, TensorType + +torch, nn = try_import_torch() + + +class APPOTorchLearner(APPOLearner, IMPALATorchLearner): + """Implements APPO loss / update logic on top of IMPALATorchLearner.""" + + @override(IMPALATorchLearner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: APPOConfig, + batch: Dict, + fwd_out: Dict[str, TensorType], + ) -> TensorType: + module = self.module[module_id].unwrapped() + assert isinstance(module, TargetNetworkAPI) + assert isinstance(module, ValueFunctionAPI) + + # TODO (sven): Now that we do the +1ts trick to be less vulnerable about + # bootstrap values at the end of rollouts in the new stack, we might make + # this a more flexible, configurable parameter for users, e.g. + # `v_trace_seq_len` (independent of `rollout_fragment_length`). Separation + # of concerns (sampling vs learning). + rollout_frag_or_episode_len = config.get_rollout_fragment_length() + recurrent_seq_len = batch.get("seq_lens") + + loss_mask = batch[Columns.LOSS_MASK].float() + loss_mask_time_major = make_time_major( + loss_mask, + trajectory_len=rollout_frag_or_episode_len, + recurrent_seq_len=recurrent_seq_len, + ) + size_loss_mask = torch.sum(loss_mask) + + values = module.compute_values( + batch, embeddings=fwd_out.get(Columns.EMBEDDINGS) + ) + + action_dist_cls_train = module.get_train_action_dist_cls() + target_policy_dist = action_dist_cls_train.from_logits( + fwd_out[Columns.ACTION_DIST_INPUTS] + ) + + old_target_policy_dist = action_dist_cls_train.from_logits( + module.forward_target(batch)[TARGET_NETWORK_ACTION_DIST_INPUTS] + ) + old_target_policy_actions_logp = old_target_policy_dist.logp( + batch[Columns.ACTIONS] + ) + behaviour_actions_logp = batch[Columns.ACTION_LOGP] + target_actions_logp = target_policy_dist.logp(batch[Columns.ACTIONS]) + + behaviour_actions_logp_time_major = make_time_major( + behaviour_actions_logp, + trajectory_len=rollout_frag_or_episode_len, + recurrent_seq_len=recurrent_seq_len, + ) + target_actions_logp_time_major = make_time_major( + target_actions_logp, + trajectory_len=rollout_frag_or_episode_len, + recurrent_seq_len=recurrent_seq_len, + ) + old_actions_logp_time_major = make_time_major( + old_target_policy_actions_logp, + trajectory_len=rollout_frag_or_episode_len, + recurrent_seq_len=recurrent_seq_len, + ) + rewards_time_major = make_time_major( + batch[Columns.REWARDS], + trajectory_len=rollout_frag_or_episode_len, + recurrent_seq_len=recurrent_seq_len, + ) + values_time_major = make_time_major( + values, + trajectory_len=rollout_frag_or_episode_len, + recurrent_seq_len=recurrent_seq_len, + ) + assert Columns.VALUES_BOOTSTRAPPED not in batch + # Use as bootstrap values the vf-preds in the next "batch row", except + # for the very last row (which doesn't have a next row), for which the + # bootstrap value does not matter b/c it has a +1ts value at its end + # anyways. So we chose an arbitrary item (for simplicity of not having to + # move new data to the device). + bootstrap_values = torch.cat( + [ + values_time_major[0][1:], # 0th ts values from "next row" + values_time_major[0][0:1], # <- can use any arbitrary value here + ], + dim=0, + ) + + # The discount factor that is used should be gamma except for timesteps where + # the episode is terminated. In that case, the discount factor should be 0. + discounts_time_major = ( + 1.0 + - make_time_major( + batch[Columns.TERMINATEDS], + trajectory_len=rollout_frag_or_episode_len, + recurrent_seq_len=recurrent_seq_len, + ).float() + ) * config.gamma + + # Note that vtrace will compute the main loop on the CPU for better performance. + vtrace_adjusted_target_values, pg_advantages = vtrace_torch( + target_action_log_probs=old_actions_logp_time_major, + behaviour_action_log_probs=behaviour_actions_logp_time_major, + discounts=discounts_time_major, + rewards=rewards_time_major, + values=values_time_major, + bootstrap_values=bootstrap_values, + clip_pg_rho_threshold=config.vtrace_clip_pg_rho_threshold, + clip_rho_threshold=config.vtrace_clip_rho_threshold, + ) + pg_advantages = pg_advantages * loss_mask_time_major + + # The policy gradients loss. + is_ratio = torch.clip( + torch.exp(behaviour_actions_logp_time_major - old_actions_logp_time_major), + 0.0, + 2.0, + ) + logp_ratio = is_ratio * torch.exp( + target_actions_logp_time_major - behaviour_actions_logp_time_major + ) + + surrogate_loss = torch.minimum( + pg_advantages * logp_ratio, + pg_advantages + * torch.clip(logp_ratio, 1 - config.clip_param, 1 + config.clip_param), + ) + + if config.use_kl_loss: + action_kl = old_target_policy_dist.kl(target_policy_dist) * loss_mask + mean_kl_loss = torch.sum(action_kl) / size_loss_mask + else: + mean_kl_loss = 0.0 + mean_pi_loss = -(torch.sum(surrogate_loss) / size_loss_mask) + + # The baseline loss. + delta = values_time_major - vtrace_adjusted_target_values + vf_loss = 0.5 * torch.sum(torch.pow(delta, 2.0) * loss_mask_time_major) + mean_vf_loss = vf_loss / size_loss_mask + + # The entropy loss. + mean_entropy_loss = ( + -torch.sum(target_policy_dist.entropy() * loss_mask) / size_loss_mask + ) + + # The summed weighted loss. + total_loss = ( + mean_pi_loss + + (mean_vf_loss * config.vf_loss_coeff) + + ( + mean_entropy_loss + * self.entropy_coeff_schedulers_per_module[ + module_id + ].get_current_value() + ) + + (mean_kl_loss * self.curr_kl_coeffs_per_module[module_id]) + ) + + # Log important loss stats. + self.metrics.log_dict( + { + POLICY_LOSS_KEY: mean_pi_loss, + VF_LOSS_KEY: mean_vf_loss, + ENTROPY_KEY: -mean_entropy_loss, + LEARNER_RESULTS_KL_KEY: mean_kl_loss, + LEARNER_RESULTS_CURR_KL_COEFF_KEY: ( + self.curr_kl_coeffs_per_module[module_id] + ), + }, + key=module_id, + window=1, # <- single items (should not be mean/ema-reduced over time). + ) + # Return the total loss. + return total_loss + + @override(APPOLearner) + def _update_module_kl_coeff(self, module_id: ModuleID, config: APPOConfig) -> None: + # Update the current KL value based on the recently measured value. + # Increase. + kl = convert_to_numpy(self.metrics.peek((module_id, LEARNER_RESULTS_KL_KEY))) + kl_coeff_var = self.curr_kl_coeffs_per_module[module_id] + + if kl > 2.0 * config.kl_target: + # TODO (Kourosh) why not *2.0? + kl_coeff_var.data *= 1.5 + # Decrease. + elif kl < 0.5 * config.kl_target: + kl_coeff_var.data *= 0.5 + + self.metrics.log_value( + (module_id, LEARNER_RESULTS_CURR_KL_COEFF_KEY), + kl_coeff_var.item(), + window=1, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/appo_torch_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/appo_torch_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..3bb3f0ba7f405cd9e69e79c3f3e12800ba69a572 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/appo_torch_rl_module.py @@ -0,0 +1,13 @@ +# Backward compat import. +from ray.rllib.algorithms.appo.torch.default_appo_torch_rl_module import ( # noqa + DefaultAPPOTorchRLModule as APPOTorchRLModule, +) +from ray._common.deprecation import deprecation_warning + + +deprecation_warning( + old="ray.rllib.algorithms.appo.torch.appo_torch_rl_module.APPOTorchRLModule", + new="ray.rllib.algorithms.appo.torch.default_appo_torch_rl_module." + "DefaultAPPOTorchRLModule", + error=False, +) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/default_appo_torch_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/default_appo_torch_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..637c3de89d29ed98a3460a86c058ee20fac6b813 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/torch/default_appo_torch_rl_module.py @@ -0,0 +1,10 @@ +from ray.rllib.algorithms.appo.default_appo_rl_module import DefaultAPPORLModule +from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( + DefaultPPOTorchRLModule, +) +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class DefaultAPPOTorchRLModule(DefaultPPOTorchRLModule, DefaultAPPORLModule): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/appo/utils.py b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..81caccfcfa8e7110cb3bc7c7040210c31cf0faa4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/appo/utils.py @@ -0,0 +1,174 @@ +""" +[1] IMPACT: Importance Weighted Asynchronous Architectures with Clipped Target Networks. +Luo et al. 2020 +https://arxiv.org/pdf/1912.00167 +""" +import threading +import time +from collections import deque + +import numpy as np + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.metrics.ray_metrics import ( + DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + TimerAndPrometheusLogger, +) +from ray.util.metrics import Counter, Histogram + +POLICY_SCOPE = "func" +TARGET_POLICY_SCOPE = "target_func" + + +class CircularBuffer: + """A circular batch-wise buffer as described in [1] for APPO. + + The buffer holds at most N batches, which are sampled at random (uniformly). + If full and a new batch is added, the oldest batch is discarded. Also, each batch + currently in the buffer can be sampled at most K times (after which it is also + discarded). + """ + + def __init__(self, num_batches: int, iterations_per_batch: int): + # N from the paper (buffer size). + self.num_batches = num_batches + # K ("replay coefficient") from the paper. + self.iterations_per_batch = iterations_per_batch + + self._NxK = self.num_batches * self.iterations_per_batch + self._num_added = 0 + + self._buffer = deque([None for _ in range(self._NxK)], maxlen=self._NxK) + self._indices = set() + self._offset = self._NxK + self._lock = threading.Lock() + + self._rng = np.random.default_rng() + + # Ray metrics + self._metrics_circular_buffer_add_time = Histogram( + name="rllib_utils_circular_buffer_add_time", + description="Time spent in CircularBuffer.add()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_circular_buffer_add_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_circular_buffer_add_ts_dropped = Counter( + name="rllib_utils_circular_buffer_add_ts_dropped_counter", + description="Total number of env steps dropped by the CircularBuffer.", + tag_keys=("rllib",), + ) + self._metrics_circular_buffer_add_ts_dropped.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_circular_buffer_sample_time = Histogram( + name="rllib_utils_circular_buffer_sample_time", + description="Time spent in CircularBuffer.sample()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_circular_buffer_sample_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + def add(self, batch): + # Add buffer and k=0 information to the deque. + with TimerAndPrometheusLogger(self._metrics_circular_buffer_add_time): + with self._lock: + dropped_entry = self._buffer[0] + for _ in range(self.iterations_per_batch): + self._buffer.append(batch) + self._indices.add(self._offset) + self._indices.discard(self._offset - self._NxK) + self._offset += 1 + self._num_added += 1 + + # A valid entry (w/ a batch whose k has not been reach K yet) was dropped. + dropped_ts = 0 + if dropped_entry is not None: + dropped_ts = dropped_entry.env_steps() + if dropped_ts > 0: + self._metrics_circular_buffer_add_ts_dropped.inc(value=dropped_ts) + + return dropped_ts + + def sample(self): + # Only initially, the buffer may be empty -> Just wait for some time. + with TimerAndPrometheusLogger(self._metrics_circular_buffer_sample_time): + while len(self) == 0: + time.sleep(0.0001) + + # Sample a random buffer index. + with self._lock: + idx = self._rng.choice(list(self._indices)) + actual_buffer_idx = idx - self._offset + self._NxK + batch = self._buffer[actual_buffer_idx] + assert batch is not None, ( + idx, + actual_buffer_idx, + self._offset, + self._indices, + [b is None for b in self._buffer], + ) + self._buffer[actual_buffer_idx] = None + self._indices.discard(idx) + + # Return the sampled batch. + return batch + + @property + def filled(self): + """Whether the buffer has been filled once with at least `self.num_batches`.""" + with self._lock: + return self._num_added >= self.num_batches + + def __len__(self) -> int: + """Returns the number of actually valid (non-expired) batches in the buffer.""" + with self._lock: + return len(self._indices) + + +@OldAPIStack +def make_appo_models(policy) -> ModelV2: + """Builds model and target model for APPO. + + Returns: + ModelV2: The Model for the Policy to use. + Note: The target model will not be returned, just assigned to + `policy.target_model`. + """ + # Get the num_outputs for the following model construction calls. + _, logit_dim = ModelCatalog.get_action_dist( + policy.action_space, policy.config["model"] + ) + + # Construct the (main) model. + policy.model = ModelCatalog.get_model_v2( + policy.observation_space, + policy.action_space, + logit_dim, + policy.config["model"], + name=POLICY_SCOPE, + framework=policy.framework, + ) + policy.model_variables = policy.model.variables() + + # Construct the target model. + policy.target_model = ModelCatalog.get_model_v2( + policy.observation_space, + policy.action_space, + logit_dim, + policy.config["model"], + name=TARGET_POLICY_SCOPE, + framework=policy.framework, + ) + policy.target_model_variables = policy.target_model.variables() + + # Return only the model (not the target model). + return policy.model diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/callbacks.py b/lib/python3.12/site-packages/ray/rllib/algorithms/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..9330e66335d71b6932d5fffdb03e3841087237f2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/callbacks.py @@ -0,0 +1,7 @@ +# @OldAPIStack +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.callbacks.utils import _make_multi_callbacks + +# Backward compatibility +DefaultCallbacks = RLlibCallback +make_multi_callbacks = _make_multi_callbacks diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e132e3a8ae8baa2eb5d4787315ae3c30dff8d2e5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__init__.py @@ -0,0 +1,18 @@ +from ray.rllib.algorithms.marwil.marwil import ( + MARWIL, + MARWILConfig, +) +from ray.rllib.algorithms.marwil.marwil_tf_policy import ( + MARWILTF1Policy, + MARWILTF2Policy, +) +from ray.rllib.algorithms.marwil.marwil_torch_policy import MARWILTorchPolicy + +__all__ = [ + "MARWIL", + "MARWILConfig", + # @OldAPIStack + "MARWILTF1Policy", + "MARWILTF2Policy", + "MARWILTorchPolicy", +] diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c3f036238bad173ecd65d0c1f54d2e3a8308a9f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..03732a347957e320ccd5858f2bd11998e39f4ca7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bd5e592452d14952f53f3a2ba5f81c5f368ca19e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_tf_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_tf_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b4bece4e3c321433cfeb39fe14d2c894074cc76 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_tf_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_torch_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_torch_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9cc11b3cac1c4a984ecc35dea8e5404302b3d3e2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/__pycache__/marwil_torch_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil.py b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil.py new file mode 100644 index 0000000000000000000000000000000000000000..e54843213e64509ae98638b7606aa96a9fbee979 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil.py @@ -0,0 +1,536 @@ +from typing import Callable, Optional, Type, Union + +from typing_extensions import Self + +from ray._common.deprecation import deprecation_warning +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided +from ray.rllib.connectors.learner import ( + AddNextObservationsFromEpisodesToTrainBatch, + AddObservationsFromEpisodesToBatch, + AddOneTsToEpisodesAndTruncate, + GeneralAdvantageEstimation, +) +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.learner.training_data import TrainingData +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.execution.rollout_ops import ( + synchronous_parallel_sample, +) +from ray.rllib.execution.train_ops import ( + multi_gpu_train_one_step, + train_one_step, +) +from ray.rllib.policy.policy import Policy +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.metrics import ( + LEARNER_RESULTS, + LEARNER_UPDATE_TIMER, + NUM_AGENT_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED, + OFFLINE_SAMPLING_TIMER, + SAMPLE_TIMER, + SYNCH_WORKER_WEIGHTS_TIMER, + TIMERS, +) +from ray.rllib.utils.typing import ( + EnvType, + ResultDict, + RLModuleSpecType, +) +from ray.tune.logger import Logger + + +class MARWILConfig(AlgorithmConfig): + """Defines a configuration class from which a MARWIL Algorithm can be built. + + .. testcode:: + + import gymnasium as gym + import numpy as np + + from pathlib import Path + from ray.rllib.algorithms.marwil import MARWILConfig + + # Get the base path (to ray/rllib) + base_path = Path(__file__).parents[2] + # Get the path to the data in rllib folder. + data_path = base_path / "tests/data/cartpole/cartpole-v1_large" + + config = MARWILConfig() + # Enable the new API stack. + config.api_stack( + enable_rl_module_and_learner=True, + enable_env_runner_and_connector_v2=True, + ) + # Define the environment for which to learn a policy + # from offline data. + config.environment( + observation_space=gym.spaces.Box( + np.array([-4.8, -np.inf, -0.41887903, -np.inf]), + np.array([4.8, np.inf, 0.41887903, np.inf]), + shape=(4,), + dtype=np.float32, + ), + action_space=gym.spaces.Discrete(2), + ) + # Set the training parameters. + config.training( + beta=1.0, + lr=1e-5, + gamma=0.99, + # We must define a train batch size for each + # learner (here 1 local learner). + train_batch_size_per_learner=2000, + ) + # Define the data source for offline data. + config.offline_data( + input_=[data_path.as_posix()], + # Run exactly one update per training iteration. + dataset_num_iters_per_learner=1, + ) + + # Build an `Algorithm` object from the config and run 1 training + # iteration. + algo = config.build() + algo.train() + + .. testcode:: + + import gymnasium as gym + import numpy as np + + from pathlib import Path + from ray.rllib.algorithms.marwil import MARWILConfig + from ray import tune + + # Get the base path (to ray/rllib) + base_path = Path(__file__).parents[2] + # Get the path to the data in rllib folder. + data_path = base_path / "tests/data/cartpole/cartpole-v1_large" + + config = MARWILConfig() + # Enable the new API stack. + config.api_stack( + enable_rl_module_and_learner=True, + enable_env_runner_and_connector_v2=True, + ) + # Print out some default values + print(f"beta: {config.beta}") + # Update the config object. + config.training( + lr=tune.grid_search([1e-3, 1e-4]), + beta=0.75, + # We must define a train batch size for each + # learner (here 1 local learner). + train_batch_size_per_learner=2000, + ) + # Set the config's data path. + config.offline_data( + input_=[data_path.as_posix()], + # Set the number of updates to be run per learner + # per training step. + dataset_num_iters_per_learner=1, + ) + # Set the config's environment for evalaution. + config.environment( + observation_space=gym.spaces.Box( + np.array([-4.8, -np.inf, -0.41887903, -np.inf]), + np.array([4.8, np.inf, 0.41887903, np.inf]), + shape=(4,), + dtype=np.float32, + ), + action_space=gym.spaces.Discrete(2), + ) + # Set up a tuner to run the experiment. + tuner = tune.Tuner( + "MARWIL", + param_space=config, + run_config=tune.RunConfig( + stop={"training_iteration": 1}, + ), + ) + # Run the experiment. + tuner.fit() + """ + + def __init__(self, algo_class=None): + """Initializes a MARWILConfig instance.""" + self.exploration_config = { + # The Exploration class to use. In the simplest case, this is the name + # (str) of any class present in the `rllib.utils.exploration` package. + # You can also provide the python class directly or the full location + # of your class (e.g. "ray.rllib.utils.exploration.epsilon_greedy. + # EpsilonGreedy"). + "type": "StochasticSampling", + # Add constructor kwargs here (if any). + } + + super().__init__(algo_class=algo_class or MARWIL) + self._is_online = False + # fmt: off + # __sphinx_doc_begin__ + # MARWIL specific settings: + self.beta = 1.0 + self.bc_logstd_coeff = 0.0 + self.moving_average_sqd_adv_norm_update_rate = 1e-8 + self.moving_average_sqd_adv_norm_start = 100.0 + self.vf_coeff = 1.0 + self.model["vf_share_layers"] = False + self.grad_clip = None + + # Override some of AlgorithmConfig's default values with MARWIL-specific values. + + # You should override input_ to point to an offline dataset + # (see algorithm.py and algorithm_config.py). + # The dataset may have an arbitrary number of timesteps + # (and even episodes) per line. + # However, each line must only contain consecutive timesteps in + # order for MARWIL to be able to calculate accumulated + # discounted returns. It is ok, though, to have multiple episodes in + # the same line. + self.input_ = "sampler" + self.postprocess_inputs = True + self.lr = 1e-4 + self.lambda_ = 1.0 + self.train_batch_size = 2000 + + # Materialize only the data in raw format, but not the mapped data b/c + # MARWIL uses a connector to calculate values and therefore the module + # needs to be updated frequently. This updating would not work if we + # map the data once at the beginning. + # TODO (simon, sven): The module is only updated when the OfflinePreLearner + # gets reinitiated, i.e. when the iterator gets reinitiated. This happens + # frequently enough with a small dataset, but with a big one this does not + # update often enough. We might need to put model weigths every couple of + # iterations into the object storage (maybe also connector states). + self.materialize_data = True + self.materialize_mapped_data = False + # __sphinx_doc_end__ + # fmt: on + self._set_off_policy_estimation_methods = False + + @override(AlgorithmConfig) + def training( + self, + *, + beta: Optional[float] = NotProvided, + bc_logstd_coeff: Optional[float] = NotProvided, + moving_average_sqd_adv_norm_update_rate: Optional[float] = NotProvided, + moving_average_sqd_adv_norm_start: Optional[float] = NotProvided, + vf_coeff: Optional[float] = NotProvided, + grad_clip: Optional[float] = NotProvided, + **kwargs, + ) -> Self: + """Sets the training related configuration. + + Args: + beta: Scaling of advantages in exponential terms. When beta is 0.0, + MARWIL is reduced to behavior cloning (imitation learning); + see bc.py algorithm in this same directory. + bc_logstd_coeff: A coefficient to encourage higher action distribution + entropy for exploration. + moving_average_sqd_adv_norm_update_rate: The rate for updating the + squared moving average advantage norm (c^2). A higher rate leads + to faster updates of this moving avergage. + moving_average_sqd_adv_norm_start: Starting value for the + squared moving average advantage norm (c^2). + vf_coeff: Balancing value estimation loss and policy optimization loss. + grad_clip: If specified, clip the global norm of gradients by this amount. + + Returns: + This updated AlgorithmConfig object. + """ + # Pass kwargs onto super's `training()` method. + super().training(**kwargs) + if beta is not NotProvided: + self.beta = beta + if bc_logstd_coeff is not NotProvided: + self.bc_logstd_coeff = bc_logstd_coeff + if moving_average_sqd_adv_norm_update_rate is not NotProvided: + self.moving_average_sqd_adv_norm_update_rate = ( + moving_average_sqd_adv_norm_update_rate + ) + if moving_average_sqd_adv_norm_start is not NotProvided: + self.moving_average_sqd_adv_norm_start = moving_average_sqd_adv_norm_start + if vf_coeff is not NotProvided: + self.vf_coeff = vf_coeff + if grad_clip is not NotProvided: + self.grad_clip = grad_clip + return self + + @override(AlgorithmConfig) + def get_default_rl_module_spec(self) -> RLModuleSpecType: + if self.framework_str == "torch": + from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( + DefaultPPOTorchRLModule, + ) + + return RLModuleSpec(module_class=DefaultPPOTorchRLModule) + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. " + "Use 'torch' instead." + ) + + @override(AlgorithmConfig) + def get_default_learner_class(self) -> Union[Type["Learner"], str]: + if self.framework_str == "torch": + from ray.rllib.algorithms.marwil.torch.marwil_torch_learner import ( + MARWILTorchLearner, + ) + + return MARWILTorchLearner + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. " + "Use 'torch' instead." + ) + + @override(AlgorithmConfig) + def evaluation( + self, + **kwargs, + ) -> Self: + """Sets the evaluation related configuration. + Returns: + This updated AlgorithmConfig object. + """ + # Pass kwargs onto super's `evaluation()` method. + super().evaluation(**kwargs) + + if "off_policy_estimation_methods" in kwargs: + # User specified their OPE methods. + self._set_off_policy_estimation_methods = True + + return self + + @override(AlgorithmConfig) + def offline_data(self, **kwargs) -> Self: + + super().offline_data(**kwargs) + + # Check, if the passed in class incorporates the `OfflinePreLearner` + # interface. + if "prelearner_class" in kwargs: + from ray.rllib.offline.offline_data import OfflinePreLearner + + if not issubclass(kwargs.get("prelearner_class"), OfflinePreLearner): + raise ValueError( + f"`prelearner_class` {kwargs.get('prelearner_class')} is not a " + "subclass of `OfflinePreLearner`. Any class passed to " + "`prelearner_class` needs to implement the interface given by " + "`OfflinePreLearner`." + ) + + return self + + @override(AlgorithmConfig) + def build( + self, + env: Optional[Union[str, EnvType]] = None, + logger_creator: Optional[Callable[[], Logger]] = None, + ) -> "Algorithm": + if not self._set_off_policy_estimation_methods: + deprecation_warning( + old=r"MARWIL used to have off_policy_estimation_methods " + "is and wis by default. This has" + r"changed to off_policy_estimation_methods: \{\}." + "If you want to use an off-policy estimator, specify it in" + ".evaluation(off_policy_estimation_methods=...)", + error=False, + ) + return super().build(env, logger_creator) + + @override(AlgorithmConfig) + def build_learner_connector( + self, + input_observation_space, + input_action_space, + device=None, + ): + pipeline = super().build_learner_connector( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + device=device, + ) + + # Before anything, add one ts to each episode (and record this in the loss + # mask, so that the computations at this extra ts are not used to compute + # the loss). + pipeline.prepend(AddOneTsToEpisodesAndTruncate()) + + # Prepend the "add-NEXT_OBS-from-episodes-to-train-batch" connector piece (right + # after the corresponding "add-OBS-..." default piece). + pipeline.insert_after( + AddObservationsFromEpisodesToBatch, + AddNextObservationsFromEpisodesToTrainBatch(), + ) + + # At the end of the pipeline (when the batch is already completed), add the + # GAE connector, which performs a vf forward pass, then computes the GAE + # computations, and puts the results of this (advantages, value targets) + # directly back in the batch. This is then the batch used for + # `forward_train` and `compute_losses`. + pipeline.append( + GeneralAdvantageEstimation(gamma=self.gamma, lambda_=self.lambda_) + ) + + return pipeline + + @override(AlgorithmConfig) + def validate(self) -> None: + # Call super's validation method. + super().validate() + + if self.beta < 0.0 or self.beta > 1.0: + self._value_error("`beta` must be within 0.0 and 1.0!") + + if self.postprocess_inputs is False and self.beta > 0.0: + self._value_error( + "`postprocess_inputs` must be True for MARWIL (to " + "calculate accum., discounted returns)! Try setting " + "`config.offline_data(postprocess_inputs=True)`." + ) + + # Assert that for a local learner the number of iterations is 1. Note, + # this is needed because we have no iterators, but instead a single + # batch returned directly from the `OfflineData.sample` method. + if ( + self.num_learners == 0 + and not self.dataset_num_iters_per_learner + and self.enable_rl_module_and_learner + ): + self._value_error( + "When using a local Learner (`config.num_learners=0`), the number of " + "iterations per learner (`dataset_num_iters_per_learner`) has to be " + "defined! Set this hyperparameter through `config.offline_data(" + "dataset_num_iters_per_learner=...)`." + ) + + @property + def _model_auto_keys(self): + return super()._model_auto_keys | {"beta": self.beta, "vf_share_layers": False} + + +class MARWIL(Algorithm): + @classmethod + @override(Algorithm) + def get_default_config(cls) -> MARWILConfig: + return MARWILConfig() + + @classmethod + @override(Algorithm) + def get_default_policy_class( + cls, config: AlgorithmConfig + ) -> Optional[Type[Policy]]: + if config["framework"] == "torch": + from ray.rllib.algorithms.marwil.marwil_torch_policy import ( + MARWILTorchPolicy, + ) + + return MARWILTorchPolicy + elif config["framework"] == "tf": + from ray.rllib.algorithms.marwil.marwil_tf_policy import ( + MARWILTF1Policy, + ) + + return MARWILTF1Policy + else: + from ray.rllib.algorithms.marwil.marwil_tf_policy import MARWILTF2Policy + + return MARWILTF2Policy + + @override(Algorithm) + def training_step(self) -> None: + """Implements training logic for the new stack + + Note, this includes so far training with the `OfflineData` + class (multi-/single-learner setup) and evaluation on + `EnvRunner`s. Note further, evaluation on the dataset itself + using estimators is not implemented, yet. + """ + # Old API stack (Policy, RolloutWorker, Connector). + if not self.config.enable_env_runner_and_connector_v2: + return self._training_step_old_api_stack() + + # TODO (simon): Take care of sampler metrics: right + # now all rewards are `nan`, which possibly confuses + # the user that sth. is not right, although it is as + # we do not step the env. + with self.metrics.log_time((TIMERS, OFFLINE_SAMPLING_TIMER)): + # If we should use an iterator in the learner(s). Note, in case of + # multiple learners we must always return a list of iterators. + return_iterator = ( + self.config.num_learners > 0 + or self.config.dataset_num_iters_per_learner != 1 + ) + + # Sampling from offline data. + batch_or_iterator = self.offline_data.sample( + num_samples=self.config.train_batch_size_per_learner, + num_shards=self.config.num_learners, + # Return an iterator, if a `Learner` should update + # multiple times per RLlib iteration. + return_iterator=return_iterator, + ) + if return_iterator: + training_data = TrainingData(data_iterators=batch_or_iterator) + else: + training_data = TrainingData(batch=batch_or_iterator) + + with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)): + # Updating the policy. + learner_results = self.learner_group.update( + training_data=training_data, + minibatch_size=self.config.train_batch_size_per_learner, + num_iters=self.config.dataset_num_iters_per_learner, + **self.offline_data.iter_batches_kwargs, + ) + + # Log training results. + self.metrics.aggregate(learner_results, key=LEARNER_RESULTS) + + @OldAPIStack + def _training_step_old_api_stack(self) -> ResultDict: + """Implements training step for the old stack. + + Note, there is no hybrid stack anymore. If you need to use `RLModule`s, + use the new api stack. + """ + # Collect SampleBatches from sample workers. + with self._timers[SAMPLE_TIMER]: + train_batch = synchronous_parallel_sample(worker_set=self.env_runner_group) + train_batch = train_batch.as_multi_agent( + module_id=list(self.config.policies)[0] + ) + self._counters[NUM_AGENT_STEPS_SAMPLED] += train_batch.agent_steps() + self._counters[NUM_ENV_STEPS_SAMPLED] += train_batch.env_steps() + + # Train. + if self.config.simple_optimizer: + train_results = train_one_step(self, train_batch) + else: + train_results = multi_gpu_train_one_step(self, train_batch) + + # TODO: Move training steps counter update outside of `train_one_step()` method. + # # Update train step counters. + # self._counters[NUM_ENV_STEPS_TRAINED] += train_batch.env_steps() + # self._counters[NUM_AGENT_STEPS_TRAINED] += train_batch.agent_steps() + + global_vars = { + "timestep": self._counters[NUM_AGENT_STEPS_SAMPLED], + } + + # Update weights - after learning on the local worker - on all remote + # workers (only those policies that were actually trained). + if self.env_runner_group.num_remote_env_runners() > 0: + with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]: + self.env_runner_group.sync_weights( + policies=list(train_results.keys()), global_vars=global_vars + ) + + # Update global vars on local worker as well. + self.env_runner.set_global_vars(global_vars) + + return train_results diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..b98b0d090f66ba4729a3968b5cd0fa9b72d32952 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_learner.py @@ -0,0 +1,51 @@ +from typing import Dict, Optional + +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.rl_module.apis import ValueFunctionAPI +from ray.rllib.utils.annotations import override +from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict +from ray.rllib.utils.typing import ModuleID, ShouldModuleBeUpdatedFn, TensorType + +LEARNER_RESULTS_MOVING_AVG_SQD_ADV_NORM_KEY = "moving_avg_sqd_adv_norm" +LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY = "vf_explained_variance" + + +# TODO (simon): Check, if the norm update should be done inside +# the Learner. +class MARWILLearner(Learner): + @override(Learner) + def build(self) -> None: + super().build() + + # Dict mapping module IDs to the respective moving averages of squared + # advantages. + self.moving_avg_sqd_adv_norms_per_module: Dict[ + ModuleID, TensorType + ] = LambdaDefaultDict( + lambda module_id: self._get_tensor_variable( + self.config.get_config_for_module( + module_id + ).moving_average_sqd_adv_norm_start + ) + ) + + @override(Learner) + def remove_module( + self, + module_id: ModuleID, + *, + new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None, + ) -> None: + super().remove_module( + module_id, + new_should_module_be_updated=new_should_module_be_updated, + ) + # In case of BC (beta==0.0 and this property never being used), + self.moving_avg_sqd_adv_norms_per_module.pop(module_id, None) + + @classmethod + @override(Learner) + def rl_module_required_apis(cls) -> list[type]: + # In order for a PPOLearner to update an RLModule, it must implement the + # following APIs: + return [ValueFunctionAPI] diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_tf_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_tf_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..2dbdb6a0efd2ec1b23f3a7988660d50019fa483f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_tf_policy.py @@ -0,0 +1,251 @@ +import logging +from typing import Any, Dict, List, Optional, Type, Union + +from ray.rllib.evaluation.postprocessing import Postprocessing, compute_advantages +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import TFActionDistribution +from ray.rllib.policy.dynamic_tf_policy_v2 import DynamicTFPolicyV2 +from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2 +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_mixins import ( + ValueNetworkMixin, + compute_gradients, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import get_variable, try_import_tf +from ray.rllib.utils.tf_utils import explained_variance +from ray.rllib.utils.typing import ( + LocalOptimizer, + ModelGradients, + TensorType, +) + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +class PostprocessAdvantages: + """Marwil's custom trajectory post-processing mixin.""" + + def __init__(self): + pass + + def postprocess_trajectory( + self, + sample_batch: SampleBatch, + other_agent_batches: Optional[Dict[Any, SampleBatch]] = None, + episode=None, + ): + sample_batch = super().postprocess_trajectory( + sample_batch, other_agent_batches, episode + ) + + # Trajectory is actually complete -> last r=0.0. + if sample_batch[SampleBatch.TERMINATEDS][-1]: + last_r = 0.0 + # Trajectory has been truncated -> last r=VF estimate of last obs. + else: + # Input dict is provided to us automatically via the Model's + # requirements. It's a single-timestep (last one in trajectory) + # input_dict. + # Create an input dict according to the Model's requirements. + index = "last" if SampleBatch.NEXT_OBS in sample_batch else -1 + input_dict = sample_batch.get_single_step_input_dict( + self.view_requirements, index=index + ) + last_r = self._value(**input_dict) + + # Adds the "advantages" (which in the case of MARWIL are simply the + # discounted cumulative rewards) to the SampleBatch. + return compute_advantages( + sample_batch, + last_r, + self.config["gamma"], + # We just want the discounted cumulative rewards, so we won't need + # GAE nor critic (use_critic=True: Subtract vf-estimates from returns). + use_gae=False, + use_critic=False, + ) + + +class MARWILLoss: + def __init__( + self, + policy: Policy, + value_estimates: TensorType, + action_dist: ActionDistribution, + train_batch: SampleBatch, + vf_loss_coeff: float, + beta: float, + ): + # L = - A * log\pi_\theta(a|s) + logprobs = action_dist.logp(train_batch[SampleBatch.ACTIONS]) + if beta != 0.0: + cumulative_rewards = train_batch[Postprocessing.ADVANTAGES] + # Advantage Estimation. + adv = cumulative_rewards - value_estimates + adv_squared = tf.reduce_mean(tf.math.square(adv)) + # Value function's loss term (MSE). + self.v_loss = 0.5 * adv_squared + + # Perform moving averaging of advantage^2. + rate = policy.config["moving_average_sqd_adv_norm_update_rate"] + + # Update averaged advantage norm. + # Eager. + if policy.config["framework"] == "tf2": + update_term = adv_squared - policy._moving_average_sqd_adv_norm + policy._moving_average_sqd_adv_norm.assign_add(rate * update_term) + + # Exponentially weighted advantages. + c = tf.math.sqrt(policy._moving_average_sqd_adv_norm) + exp_advs = tf.math.exp(beta * (adv / (1e-8 + c))) + # Static graph. + else: + update_adv_norm = tf1.assign_add( + ref=policy._moving_average_sqd_adv_norm, + value=rate * (adv_squared - policy._moving_average_sqd_adv_norm), + ) + + # Exponentially weighted advantages. + with tf1.control_dependencies([update_adv_norm]): + exp_advs = tf.math.exp( + beta + * tf.math.divide( + adv, + 1e-8 + tf.math.sqrt(policy._moving_average_sqd_adv_norm), + ) + ) + exp_advs = tf.stop_gradient(exp_advs) + + self.explained_variance = tf.reduce_mean( + explained_variance(cumulative_rewards, value_estimates) + ) + + else: + # Value function's loss term (MSE). + self.v_loss = tf.constant(0.0) + exp_advs = 1.0 + + # logprob loss alone tends to push action distributions to + # have very low entropy, resulting in worse performance for + # unfamiliar situations. + # A scaled logstd loss term encourages stochasticity, thus + # alleviate the problem to some extent. + logstd_coeff = policy.config["bc_logstd_coeff"] + if logstd_coeff > 0.0: + logstds = tf.reduce_sum(action_dist.log_std, axis=1) + else: + logstds = 0.0 + + self.p_loss = -1.0 * tf.reduce_mean( + exp_advs * (logprobs + logstd_coeff * logstds) + ) + + self.total_loss = self.p_loss + vf_loss_coeff * self.v_loss + + +# We need this builder function because we want to share the same +# custom logics between TF1 dynamic and TF2 eager policies. +def get_marwil_tf_policy(name: str, base: type) -> type: + """Construct a MARWILTFPolicy inheriting either dynamic or eager base policies. + + Args: + base: Base class for this policy. DynamicTFPolicyV2 or EagerTFPolicyV2. + + Returns: + A TF Policy to be used with MAML. + """ + + class MARWILTFPolicy(ValueNetworkMixin, PostprocessAdvantages, base): + def __init__( + self, + observation_space, + action_space, + config, + existing_model=None, + existing_inputs=None, + ): + # First thing first, enable eager execution if necessary. + base.enable_eager_execution_if_necessary() + + # Initialize base class. + base.__init__( + self, + observation_space, + action_space, + config, + existing_inputs=existing_inputs, + existing_model=existing_model, + ) + + ValueNetworkMixin.__init__(self, config) + PostprocessAdvantages.__init__(self) + + # Not needed for pure BC. + if config["beta"] != 0.0: + # Set up a tf-var for the moving avg (do this here to make it work + # with eager mode); "c^2" in the paper. + self._moving_average_sqd_adv_norm = get_variable( + config["moving_average_sqd_adv_norm_start"], + framework="tf", + tf_name="moving_average_of_advantage_norm", + trainable=False, + ) + + # Note: this is a bit ugly, but loss and optimizer initialization must + # happen after all the MixIns are initialized. + self.maybe_initialize_optimizer_and_loss() + + @override(base) + def loss( + self, + model: Union[ModelV2, "tf.keras.Model"], + dist_class: Type[TFActionDistribution], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + model_out, _ = model(train_batch) + action_dist = dist_class(model_out, model) + value_estimates = model.value_function() + + self._marwil_loss = MARWILLoss( + self, + value_estimates, + action_dist, + train_batch, + self.config["vf_coeff"], + self.config["beta"], + ) + + return self._marwil_loss.total_loss + + @override(base) + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + stats = { + "policy_loss": self._marwil_loss.p_loss, + "total_loss": self._marwil_loss.total_loss, + } + if self.config["beta"] != 0.0: + stats["moving_average_sqd_adv_norm"] = self._moving_average_sqd_adv_norm + stats["vf_explained_var"] = self._marwil_loss.explained_variance + stats["vf_loss"] = self._marwil_loss.v_loss + + return stats + + @override(base) + def compute_gradients_fn( + self, optimizer: LocalOptimizer, loss: TensorType + ) -> ModelGradients: + return compute_gradients(self, optimizer, loss) + + MARWILTFPolicy.__name__ = name + MARWILTFPolicy.__qualname__ = name + + return MARWILTFPolicy + + +MARWILTF1Policy = get_marwil_tf_policy("MARWILTF1Policy", DynamicTFPolicyV2) +MARWILTF2Policy = get_marwil_tf_policy("MARWILTF2Policy", EagerTFPolicyV2) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_torch_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_torch_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..219a0b176d918fd34a123ef1ea66d921660a19e4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/marwil_torch_policy.py @@ -0,0 +1,132 @@ +from typing import Dict, List, Type, Union + +from ray.rllib.algorithms.marwil.marwil_tf_policy import PostprocessAdvantages +from ray.rllib.evaluation.postprocessing import Postprocessing +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_mixins import ValueNetworkMixin +from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.torch_utils import apply_grad_clipping, explained_variance +from ray.rllib.utils.typing import TensorType + +torch, _ = try_import_torch() + + +class MARWILTorchPolicy(ValueNetworkMixin, PostprocessAdvantages, TorchPolicyV2): + """PyTorch policy class used with Marwil.""" + + def __init__(self, observation_space, action_space, config): + TorchPolicyV2.__init__( + self, + observation_space, + action_space, + config, + max_seq_len=config["model"]["max_seq_len"], + ) + + ValueNetworkMixin.__init__(self, config) + PostprocessAdvantages.__init__(self) + + # Not needed for pure BC. + if config["beta"] != 0.0: + # Set up a torch-var for the squared moving avg. advantage norm. + self._moving_average_sqd_adv_norm = torch.tensor( + [config["moving_average_sqd_adv_norm_start"]], + dtype=torch.float32, + requires_grad=False, + ).to(self.device) + + # TODO: Don't require users to call this manually. + self._initialize_loss_from_dummy_batch() + + @override(TorchPolicyV2) + def loss( + self, + model: ModelV2, + dist_class: Type[TorchDistributionWrapper], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + model_out, _ = model(train_batch) + action_dist = dist_class(model_out, model) + actions = train_batch[SampleBatch.ACTIONS] + # log\pi_\theta(a|s) + logprobs = action_dist.logp(actions) + + # Advantage estimation. + if self.config["beta"] != 0.0: + cumulative_rewards = train_batch[Postprocessing.ADVANTAGES] + state_values = model.value_function() + adv = cumulative_rewards - state_values + adv_squared_mean = torch.mean(torch.pow(adv, 2.0)) + + explained_var = explained_variance(cumulative_rewards, state_values) + ev = torch.mean(explained_var) + model.tower_stats["explained_variance"] = ev + + # Policy loss. + # Update averaged advantage norm. + rate = self.config["moving_average_sqd_adv_norm_update_rate"] + self._moving_average_sqd_adv_norm = ( + rate * (adv_squared_mean.detach() - self._moving_average_sqd_adv_norm) + + self._moving_average_sqd_adv_norm + ) + model.tower_stats[ + "_moving_average_sqd_adv_norm" + ] = self._moving_average_sqd_adv_norm + # Exponentially weighted advantages. + exp_advs = torch.exp( + self.config["beta"] + * (adv / (1e-8 + torch.pow(self._moving_average_sqd_adv_norm, 0.5))) + ).detach() + # Value loss. + v_loss = 0.5 * adv_squared_mean + else: + # Policy loss (simple BC loss term). + exp_advs = 1.0 + # Value loss. + v_loss = 0.0 + model.tower_stats["v_loss"] = v_loss + # logprob loss alone tends to push action distributions to + # have very low entropy, resulting in worse performance for + # unfamiliar situations. + # A scaled logstd loss term encourages stochasticity, thus + # alleviate the problem to some extent. + logstd_coeff = self.config["bc_logstd_coeff"] + if logstd_coeff > 0.0: + logstds = torch.mean(action_dist.log_std, dim=1) + else: + logstds = 0.0 + + p_loss = -torch.mean(exp_advs * (logprobs + logstd_coeff * logstds)) + model.tower_stats["p_loss"] = p_loss + # Combine both losses. + self.v_loss = v_loss + self.p_loss = p_loss + total_loss = p_loss + self.config["vf_coeff"] * v_loss + model.tower_stats["total_loss"] = total_loss + return total_loss + + @override(TorchPolicyV2) + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + stats = { + "policy_loss": self.get_tower_stats("p_loss")[0].item(), + "total_loss": self.get_tower_stats("total_loss")[0].item(), + } + if self.config["beta"] != 0.0: + stats["moving_average_sqd_adv_norm"] = self.get_tower_stats( + "_moving_average_sqd_adv_norm" + )[0].item() + stats["vf_explained_var"] = self.get_tower_stats("explained_variance")[ + 0 + ].item() + stats["vf_loss"] = self.get_tower_stats("v_loss")[0].item() + return convert_to_numpy(stats) + + def extra_grad_process( + self, optimizer: "torch.optim.Optimizer", loss: TensorType + ) -> Dict[str, TensorType]: + return apply_grad_clipping(self, optimizer, loss) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5414db8f6d240aaf7ae4be2c76843fc7ba926208 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__pycache__/marwil_torch_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__pycache__/marwil_torch_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45e72e23ffb83a7f0920c798deabaafaefe0967e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/__pycache__/marwil_torch_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/marwil_torch_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/marwil_torch_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..a61eb3bde13596b515497aac396383b4cbb08068 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/marwil/torch/marwil_torch_learner.py @@ -0,0 +1,138 @@ +from typing import Any, Dict, Optional + +from ray.rllib.algorithms.marwil.marwil import MARWILConfig +from ray.rllib.algorithms.marwil.marwil_learner import ( + LEARNER_RESULTS_MOVING_AVG_SQD_ADV_NORM_KEY, + LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY, + MARWILLearner, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.learner import POLICY_LOSS_KEY, VF_LOSS_KEY +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.evaluation.postprocessing import Postprocessing +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.torch_utils import explained_variance +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +class MARWILTorchLearner(MARWILLearner, TorchLearner): + """Implements torch-specific MARWIL loss on top of MARWILLearner. + + This class implements the MARWIL loss under `self.compute_loss_for_module()`. + """ + + def compute_loss_for_module( + self, + *, + module_id: str, + config: Optional[MARWILConfig] = None, + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType] + ) -> TensorType: + module = self.module[module_id].unwrapped() + + # Possibly apply masking to some sub loss terms and to the total loss term + # at the end. Masking could be used for RNN-based model (zero padded `batch`) + # and for PPO's batched value function (and bootstrap value) computations, + # for which we add an additional (artificial) timestep to each episode to + # simplify the actual computation. + if Columns.LOSS_MASK in batch: + num_valid = torch.sum(batch[Columns.LOSS_MASK]) + + def possibly_masked_mean(data_): + return torch.sum(data_[batch[Columns.LOSS_MASK]]) / num_valid + + else: + possibly_masked_mean = torch.mean + + action_dist_class_train = module.get_train_action_dist_cls() + curr_action_dist = action_dist_class_train.from_logits( + fwd_out[Columns.ACTION_DIST_INPUTS] + ) + + log_probs = curr_action_dist.logp(batch[Columns.ACTIONS]) + + # If beta is zero, we fall back to BC. + if config.beta == 0.0: + # Value function's loss term. + mean_vf_loss = 0.0 + # Policy's loss term. + exp_weighted_advantages = 1.0 + # Otherwise, compute advantages. + else: + # cumulative_rewards = batch[Columns.ADVANTAGES] + value_fn_out = module.compute_values( + batch, embeddings=fwd_out.get(Columns.EMBEDDINGS) + ) + advantages = batch[Columns.VALUE_TARGETS] - value_fn_out + advantages_squared_mean = possibly_masked_mean(torch.pow(advantages, 2.0)) + + # Compute the value loss. + mean_vf_loss = 0.5 * advantages_squared_mean + + # Compute the policy loss. + self.moving_avg_sqd_adv_norms_per_module[module_id] = ( + config.moving_average_sqd_adv_norm_update_rate + * ( + advantages_squared_mean.detach() + - self.moving_avg_sqd_adv_norms_per_module[module_id] + ) + + self.moving_avg_sqd_adv_norms_per_module[module_id] + ) + # Exponentially weighted advantages. + # TODO (simon): Check, if we need the mask here. + exp_weighted_advantages = torch.exp( + config.beta + * ( + advantages + / ( + 1e-8 + + torch.pow( + self.moving_avg_sqd_adv_norms_per_module[module_id], 0.5 + ) + ) + ) + ).detach() + + # Note, using solely a log-probability loss term tends to push the action + # distributions to have very low entropy, which results in worse performance + # specifically in unknown situations. + # Scaling the loss term with the logarithm of the action distribution's + # standard deviation encourages stochasticity in the policy. + if config.bc_logstd_coeff > 0.0: + log_stds = possibly_masked_mean(curr_action_dist.log_std, dim=1) + else: + log_stds = 0.0 + + # Compute the policy loss. + policy_loss = -possibly_masked_mean( + exp_weighted_advantages * (log_probs + config.bc_logstd_coeff * log_stds) + ) + + # Compute the total loss. + total_loss = policy_loss + config.vf_coeff * mean_vf_loss + + # Log import loss stats. In case of the BC loss this is simply + # the policy loss. + if config.beta == 0.0: + self.metrics.log_value((module_id, POLICY_LOSS_KEY), policy_loss, window=1) + # Log more stats, if using the MARWIL loss. + else: + ma_sqd_adv_norms = self.moving_avg_sqd_adv_norms_per_module[module_id] + self.metrics.log_dict( + { + POLICY_LOSS_KEY: policy_loss, + VF_LOSS_KEY: mean_vf_loss, + LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY: explained_variance( + batch[Postprocessing.VALUE_TARGETS], value_fn_out + ), + LEARNER_RESULTS_MOVING_AVG_SQD_ADV_NORM_KEY: ma_sqd_adv_norms, + }, + key=module_id, + window=1, # <- single items (should not be mean/ema-reduced over time). + ) + + # Return the total loss. + return total_loss diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/mock.py b/lib/python3.12/site-packages/ray/rllib/algorithms/mock.py new file mode 100644 index 0000000000000000000000000000000000000000..ba2ac262af21ab62898c661614cc8d4a58954dbf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/mock.py @@ -0,0 +1,154 @@ +import os +import pickle +import time + +import numpy as np + +from ray.rllib.algorithms.algorithm import Algorithm, AlgorithmConfig +from ray.rllib.utils.annotations import override +from ray.tune import result as tune_result + + +class _MockTrainer(Algorithm): + """Mock Algorithm for use in tests.""" + + @classmethod + @override(Algorithm) + def get_default_config(cls) -> AlgorithmConfig: + return ( + AlgorithmConfig() + .framework("tf") + .update_from_dict( + { + "mock_error": False, + "persistent_error": False, + "test_variable": 1, + "user_checkpoint_freq": 0, + "sleep": 0, + } + ) + ) + + @classmethod + def default_resource_request(cls, config: AlgorithmConfig): + return None + + @override(Algorithm) + def setup(self, config): + self.callbacks = self.config.callbacks_class() + + # Add needed properties. + self.info = None + self.restored = False + + @override(Algorithm) + def step(self): + if ( + self.config.mock_error + and self.iteration == 1 + and (self.config.persistent_error or not self.restored) + ): + raise Exception("mock error") + if self.config.sleep: + time.sleep(self.config.sleep) + result = dict( + episode_reward_mean=10, episode_len_mean=10, timesteps_this_iter=10, info={} + ) + if self.config.user_checkpoint_freq > 0 and self.iteration > 0: + if self.iteration % self.config.user_checkpoint_freq == 0: + result.update({tune_result.SHOULD_CHECKPOINT: True}) + return result + + @override(Algorithm) + def save_checkpoint(self, checkpoint_dir): + path = os.path.join(checkpoint_dir, "mock_agent.pkl") + with open(path, "wb") as f: + pickle.dump(self.info, f) + + @override(Algorithm) + def load_checkpoint(self, checkpoint_dir): + path = os.path.join(checkpoint_dir, "mock_agent.pkl") + with open(path, "rb") as f: + info = pickle.load(f) + self.info = info + self.restored = True + + @staticmethod + @override(Algorithm) + def _get_env_id_and_creator(env_specifier, config): + # No env to register. + return None, None + + def set_info(self, info): + self.info = info + return info + + def get_info(self, sess=None): + return self.info + + +class _SigmoidFakeData(_MockTrainer): + """Algorithm that returns sigmoid learning curves. + + This can be helpful for evaluating early stopping algorithms.""" + + @classmethod + @override(Algorithm) + def get_default_config(cls) -> AlgorithmConfig: + return AlgorithmConfig().update_from_dict( + { + "width": 100, + "height": 100, + "offset": 0, + "iter_time": 10, + "iter_timesteps": 1, + } + ) + + def step(self): + i = max(0, self.iteration - self.config.offset) + v = np.tanh(float(i) / self.config.width) + v *= self.config.height + return dict( + episode_reward_mean=v, + episode_len_mean=v, + timesteps_this_iter=self.config.iter_timesteps, + time_this_iter_s=self.config.iter_time, + info={}, + ) + + +class _ParameterTuningTrainer(_MockTrainer): + @classmethod + @override(Algorithm) + def get_default_config(cls) -> AlgorithmConfig: + return AlgorithmConfig().update_from_dict( + { + "reward_amt": 10, + "dummy_param": 10, + "dummy_param2": 15, + "iter_time": 10, + "iter_timesteps": 1, + } + ) + + def step(self): + return dict( + episode_reward_mean=self.config.reward_amt * self.iteration, + episode_len_mean=self.config.reward_amt, + timesteps_this_iter=self.config.iter_timesteps, + time_this_iter_s=self.config.iter_time, + info={}, + ) + + +def _algorithm_import_failed(trace): + """Returns dummy Algorithm class for if PyTorch etc. is not installed.""" + + class _AlgorithmImportFailed(Algorithm): + _name = "AlgorithmImportFailed" + + def setup(self, config): + raise ImportError(trace) + + return _AlgorithmImportFailed diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9ed907f5dd1e40cdd09786b1a823e4cdee1d8ca1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__init__.py @@ -0,0 +1,12 @@ +from ray.rllib.algorithms.ppo.ppo import PPO, PPOConfig +from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy, PPOTF2Policy +from ray.rllib.algorithms.ppo.ppo_torch_policy import PPOTorchPolicy + +__all__ = [ + "PPO", + "PPOConfig", + # @OldAPIStack + "PPOTF1Policy", + "PPOTF2Policy", + "PPOTorchPolicy", +] diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08d936ce8ddb4f109aaa6cfbbf4fef4157f256b1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/default_ppo_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/default_ppo_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c83c22ae9953395d8a021fd42135d498f757332 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/default_ppo_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd933136185d32edb6182e824d30d185a826cf4a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_catalog.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_catalog.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1da65e3c3199e9b77d051c64e1e28e32ffc23698 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_catalog.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8dec709fff2cc7d8b5f04b0523e833c922d1ad23 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63da56e596cde2d881928a5c90d3793ab8c2d8c8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_tf_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_tf_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fbe8471b97d1c51cb4192f1e5175ccd420fe0d47 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_tf_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_torch_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_torch_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ef1dd1c7bec1115de7a7d4019949388c16aa784 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/__pycache__/ppo_torch_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/default_ppo_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/default_ppo_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..5ac176452f368cca2b45308116eaf4244bab648e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/default_ppo_rl_module.py @@ -0,0 +1,62 @@ +import abc +from typing import List + +from ray.rllib.core.models.configs import RecurrentEncoderConfig +from ray.rllib.core.rl_module.apis import InferenceOnlyAPI, ValueFunctionAPI +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class DefaultPPORLModule(RLModule, InferenceOnlyAPI, ValueFunctionAPI, abc.ABC): + """Default RLModule used by PPO, if user does not specify a custom RLModule. + + Users who want to train their RLModules with PPO may implement any RLModule + (or TorchRLModule) subclass as long as the custom class also implements the + `ValueFunctionAPI` (see ray.rllib.core.rl_module.apis.value_function_api.py) + """ + + @override(RLModule) + def setup(self): + # __sphinx_doc_begin__ + # If we have a stateful model, states for the critic need to be collected + # during sampling and `inference-only` needs to be `False`. Note, at this + # point the encoder is not built, yet and therefore `is_stateful()` does + # not work. + is_stateful = isinstance( + self.catalog.actor_critic_encoder_config.base_encoder_config, + RecurrentEncoderConfig, + ) + if is_stateful: + self.inference_only = False + # If this is an `inference_only` Module, we'll have to pass this information + # to the encoder config as well. + if self.inference_only and self.framework == "torch": + self.catalog.actor_critic_encoder_config.inference_only = True + + # Build models from catalog. + self.encoder = self.catalog.build_actor_critic_encoder(framework=self.framework) + self.pi = self.catalog.build_pi_head(framework=self.framework) + self.vf = self.catalog.build_vf_head(framework=self.framework) + # __sphinx_doc_end__ + + @override(RLModule) + def get_initial_state(self) -> dict: + if hasattr(self.encoder, "get_initial_state"): + return self.encoder.get_initial_state() + else: + return {} + + @OverrideToImplementCustomLogic_CallToSuperRecommended + @override(InferenceOnlyAPI) + def get_non_inference_attributes(self) -> List[str]: + """Return attributes, which are NOT inference-only (only used for training).""" + return ["vf"] + ( + [] + if self.model_config.get("vf_share_layers") + else ["encoder.critic_encoder"] + ) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo.py new file mode 100644 index 0000000000000000000000000000000000000000..d94bd0a5c6f376ab2400e6ef5d77299e2ac0cfa5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo.py @@ -0,0 +1,559 @@ +""" +Proximal Policy Optimization (PPO) +================================== + +This file defines the distributed Algorithm class for proximal policy +optimization. +See `ppo_[tf|torch]_policy.py` for the definition of the policy loss. + +Detailed documentation: https://docs.ray.io/en/master/rllib-algorithms.html#ppo +""" + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union + +from typing_extensions import Self + +from ray._common.deprecation import DEPRECATED_VALUE +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.execution.rollout_ops import ( + standardize_fields, + synchronous_parallel_sample, +) +from ray.rllib.execution.train_ops import ( + multi_gpu_train_one_step, + train_one_step, +) +from ray.rllib.policy.policy import Policy +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.metrics import ( + ALL_MODULES, + ENV_RUNNER_RESULTS, + ENV_RUNNER_SAMPLING_TIMER, + LEARNER_RESULTS, + LEARNER_UPDATE_TIMER, + NUM_AGENT_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + SAMPLE_TIMER, + SYNCH_WORKER_WEIGHTS_TIMER, + TIMERS, +) +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.schedules.scheduler import Scheduler +from ray.rllib.utils.typing import ResultDict +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.rllib.core.learner.learner import Learner + + +logger = logging.getLogger(__name__) + +LEARNER_RESULTS_VF_LOSS_UNCLIPPED_KEY = "vf_loss_unclipped" +LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY = "vf_explained_var" +LEARNER_RESULTS_KL_KEY = "mean_kl_loss" +LEARNER_RESULTS_CURR_KL_COEFF_KEY = "curr_kl_coeff" +LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY = "curr_entropy_coeff" + + +class PPOConfig(AlgorithmConfig): + """Defines a configuration class from which a PPO Algorithm can be built. + + .. testcode:: + + from ray.rllib.algorithms.ppo import PPOConfig + + config = PPOConfig() + config.environment("CartPole-v1") + config.env_runners(num_env_runners=1) + config.training( + gamma=0.9, lr=0.01, kl_coeff=0.3, train_batch_size_per_learner=256 + ) + + # Build a Algorithm object from the config and run 1 training iteration. + algo = config.build() + algo.train() + + .. testcode:: + + from ray.rllib.algorithms.ppo import PPOConfig + from ray import tune + + config = ( + PPOConfig() + # Set the config object's env. + .environment(env="CartPole-v1") + # Update the config object's training parameters. + .training( + lr=0.001, clip_param=0.2 + ) + ) + + tune.Tuner( + "PPO", + run_config=tune.RunConfig(stop={"training_iteration": 1}), + param_space=config, + ).fit() + + .. testoutput:: + :hide: + + ... + """ + + def __init__(self, algo_class=None): + """Initializes a PPOConfig instance.""" + self.exploration_config = { + # The Exploration class to use. In the simplest case, this is the name + # (str) of any class present in the `rllib.utils.exploration` package. + # You can also provide the python class directly or the full location + # of your class (e.g. "ray.rllib.utils.exploration.epsilon_greedy. + # EpsilonGreedy"). + "type": "StochasticSampling", + # Add constructor kwargs here (if any). + } + + super().__init__(algo_class=algo_class or PPO) + + # fmt: off + # __sphinx_doc_begin__ + self.lr = 5e-5 + self.rollout_fragment_length = "auto" + self.train_batch_size = 4000 + + # PPO specific settings: + self.use_critic = True + self.use_gae = True + self.num_epochs = 30 + self.minibatch_size = 128 + self.shuffle_batch_per_epoch = True + self.lambda_ = 1.0 + self.use_kl_loss = True + self.kl_coeff = 0.2 + self.kl_target = 0.01 + self.vf_loss_coeff = 1.0 + self.entropy_coeff = 0.0 + self.clip_param = 0.3 + self.vf_clip_param = 10.0 + self.grad_clip = None + + # Override some of AlgorithmConfig's default values with PPO-specific values. + self.num_env_runners = 2 + # __sphinx_doc_end__ + # fmt: on + + self.model["vf_share_layers"] = False # @OldAPIStack + self.entropy_coeff_schedule = None # @OldAPIStack + self.lr_schedule = None # @OldAPIStack + + # Deprecated keys. + self.sgd_minibatch_size = DEPRECATED_VALUE + self.vf_share_layers = DEPRECATED_VALUE + + @override(AlgorithmConfig) + def get_default_rl_module_spec(self) -> RLModuleSpec: + if self.framework_str == "torch": + from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( + DefaultPPOTorchRLModule, + ) + + return RLModuleSpec(module_class=DefaultPPOTorchRLModule) + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. " + "Use either 'torch' or 'tf2'." + ) + + @override(AlgorithmConfig) + def get_default_learner_class(self) -> Union[Type["Learner"], str]: + if self.framework_str == "torch": + from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import ( + PPOTorchLearner, + ) + + return PPOTorchLearner + elif self.framework_str in ["tf2", "tf"]: + raise ValueError( + "TensorFlow is no longer supported on the new API stack! " + "Use `framework='torch'`." + ) + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. " + "Use `framework='torch'`." + ) + + @override(AlgorithmConfig) + def training( + self, + *, + use_critic: Optional[bool] = NotProvided, + use_gae: Optional[bool] = NotProvided, + lambda_: Optional[float] = NotProvided, + use_kl_loss: Optional[bool] = NotProvided, + kl_coeff: Optional[float] = NotProvided, + kl_target: Optional[float] = NotProvided, + vf_loss_coeff: Optional[float] = NotProvided, + entropy_coeff: Optional[float] = NotProvided, + entropy_coeff_schedule: Optional[List[List[Union[int, float]]]] = NotProvided, + clip_param: Optional[float] = NotProvided, + vf_clip_param: Optional[float] = NotProvided, + grad_clip: Optional[float] = NotProvided, + # @OldAPIStack + lr_schedule: Optional[List[List[Union[int, float]]]] = NotProvided, + # Deprecated. + vf_share_layers=DEPRECATED_VALUE, + **kwargs, + ) -> Self: + """Sets the training related configuration. + + Args: + use_critic: Should use a critic as a baseline (otherwise don't use value + baseline; required for using GAE). + use_gae: If true, use the Generalized Advantage Estimator (GAE) + with a value function, see https://arxiv.org/pdf/1506.02438.pdf. + lambda_: The lambda parameter for General Advantage Estimation (GAE). + Defines the exponential weight used between actually measured rewards + vs value function estimates over multiple time steps. Specifically, + `lambda_` balances short-term, low-variance estimates against long-term, + high-variance returns. A `lambda_` of 0.0 makes the GAE rely only on + immediate rewards (and vf predictions from there on, reducing variance, + but increasing bias), while a `lambda_` of 1.0 only incorporates vf + predictions at the truncation points of the given episodes or episode + chunks (reducing bias but increasing variance). + use_kl_loss: Whether to use the KL-term in the loss function. + kl_coeff: Initial coefficient for KL divergence. + kl_target: Target value for KL divergence. + vf_loss_coeff: Coefficient of the value function loss. IMPORTANT: you must + tune this if you set vf_share_layers=True inside your model's config. + entropy_coeff: The entropy coefficient (float) or entropy coefficient + schedule in the format of + [[timestep, coeff-value], [timestep, coeff-value], ...] + In case of a schedule, intermediary timesteps will be assigned to + linearly interpolated coefficient values. A schedule config's first + entry must start with timestep 0, i.e.: [[0, initial_value], [...]]. + clip_param: The PPO clip parameter. + vf_clip_param: Clip param for the value function. Note that this is + sensitive to the scale of the rewards. If your expected V is large, + increase this. + grad_clip: If specified, clip the global norm of gradients by this amount. + + Returns: + This updated AlgorithmConfig object. + """ + # Pass kwargs onto super's `training()` method. + super().training(**kwargs) + + if use_critic is not NotProvided: + self.use_critic = use_critic + # TODO (Kourosh) This is experimental. + # Don't forget to remove .use_critic from algorithm config. + if use_gae is not NotProvided: + self.use_gae = use_gae + if lambda_ is not NotProvided: + self.lambda_ = lambda_ + if use_kl_loss is not NotProvided: + self.use_kl_loss = use_kl_loss + if kl_coeff is not NotProvided: + self.kl_coeff = kl_coeff + if kl_target is not NotProvided: + self.kl_target = kl_target + if vf_loss_coeff is not NotProvided: + self.vf_loss_coeff = vf_loss_coeff + if entropy_coeff is not NotProvided: + self.entropy_coeff = entropy_coeff + if clip_param is not NotProvided: + self.clip_param = clip_param + if vf_clip_param is not NotProvided: + self.vf_clip_param = vf_clip_param + if grad_clip is not NotProvided: + self.grad_clip = grad_clip + + # TODO (sven): Remove these once new API stack is only option for PPO. + if lr_schedule is not NotProvided: + self.lr_schedule = lr_schedule + if entropy_coeff_schedule is not NotProvided: + self.entropy_coeff_schedule = entropy_coeff_schedule + + return self + + @override(AlgorithmConfig) + def validate(self) -> None: + # Call super's validation method. + super().validate() + + # Synchronous sampling, on-policy/PPO algos -> Check mismatches between + # `rollout_fragment_length` and `train_batch_size_per_learner` to avoid user + # confusion. + # TODO (sven): Make rollout_fragment_length a property and create a private + # attribute to store (possibly) user provided value (or "auto") in. Deprecate + # `self.get_rollout_fragment_length()`. + self.validate_train_batch_size_vs_rollout_fragment_length() + + # SGD minibatch size must be smaller than train_batch_size (b/c + # we subsample a batch of `minibatch_size` from the train-batch for + # each `num_epochs`). + if ( + not self.enable_rl_module_and_learner + and self.minibatch_size > self.train_batch_size + ): + self._value_error( + f"`minibatch_size` ({self.minibatch_size}) must be <= " + f"`train_batch_size` ({self.train_batch_size}). In PPO, the train batch" + f" will be split into {self.minibatch_size} chunks, each of which " + f"is iterated over (used for updating the policy) {self.num_epochs} " + "times." + ) + elif self.enable_rl_module_and_learner: + mbs = self.minibatch_size + tbs = self.train_batch_size_per_learner or self.train_batch_size + if isinstance(mbs, int) and isinstance(tbs, int) and mbs > tbs: + self._value_error( + f"`minibatch_size` ({mbs}) must be <= " + f"`train_batch_size_per_learner` ({tbs}). In PPO, the train batch" + f" will be split into {mbs} chunks, each of which is iterated over " + f"(used for updating the policy) {self.num_epochs} times." + ) + + # Episodes may only be truncated (and passed into PPO's + # `postprocessing_fn`), iff generalized advantage estimation is used + # (value function estimate at end of truncated episode to estimate + # remaining value). + if ( + not self.in_evaluation + and self.batch_mode == "truncate_episodes" + and not self.use_gae + ): + self._value_error( + "Episode truncation is not supported without a value " + "function (to estimate the return at the end of the truncated" + " trajectory). Consider setting " + "batch_mode=complete_episodes." + ) + + # New API stack checks. + if self.enable_rl_module_and_learner: + # `lr_schedule` checking. + if self.lr_schedule is not None: + self._value_error( + "`lr_schedule` is deprecated and must be None! Use the " + "`lr` setting to setup a schedule." + ) + if self.entropy_coeff_schedule is not None: + self._value_error( + "`entropy_coeff_schedule` is deprecated and must be None! Use the " + "`entropy_coeff` setting to setup a schedule." + ) + Scheduler.validate( + fixed_value_or_schedule=self.entropy_coeff, + setting_name="entropy_coeff", + description="entropy coefficient", + ) + if isinstance(self.entropy_coeff, float) and self.entropy_coeff < 0.0: + self._value_error("`entropy_coeff` must be >= 0.0") + + @property + @override(AlgorithmConfig) + def _model_config_auto_includes(self) -> Dict[str, Any]: + return super()._model_config_auto_includes | {"vf_share_layers": False} + + +class PPO(Algorithm): + @classmethod + @override(Algorithm) + def get_default_config(cls) -> PPOConfig: + return PPOConfig() + + @classmethod + @override(Algorithm) + def get_default_policy_class( + cls, config: AlgorithmConfig + ) -> Optional[Type[Policy]]: + if config["framework"] == "torch": + + from ray.rllib.algorithms.ppo.ppo_torch_policy import PPOTorchPolicy + + return PPOTorchPolicy + elif config["framework"] == "tf": + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy + + return PPOTF1Policy + else: + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF2Policy + + return PPOTF2Policy + + @override(Algorithm) + def training_step(self) -> None: + # Old API stack (Policy, RolloutWorker, Connector). + if not self.config.enable_env_runner_and_connector_v2: + return self._training_step_old_api_stack() + + # Collect batches from sample workers until we have a full batch. + with self.metrics.log_time((TIMERS, ENV_RUNNER_SAMPLING_TIMER)): + # Sample in parallel from the workers. + if self.config.count_steps_by == "agent_steps": + episodes, env_runner_results = synchronous_parallel_sample( + worker_set=self.env_runner_group, + max_agent_steps=self.config.total_train_batch_size, + sample_timeout_s=self.config.sample_timeout_s, + _uses_new_env_runners=( + self.config.enable_env_runner_and_connector_v2 + ), + _return_metrics=True, + ) + else: + episodes, env_runner_results = synchronous_parallel_sample( + worker_set=self.env_runner_group, + max_env_steps=self.config.total_train_batch_size, + sample_timeout_s=self.config.sample_timeout_s, + _uses_new_env_runners=( + self.config.enable_env_runner_and_connector_v2 + ), + _return_metrics=True, + ) + # Return early if all our workers failed. + if not episodes: + return + + # Reduce EnvRunner metrics over the n EnvRunners. + self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS) + + # Perform a learner update step on the collected episodes. + with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)): + learner_results = self.learner_group.update( + episodes=episodes, + timesteps={ + NUM_ENV_STEPS_SAMPLED_LIFETIME: ( + self.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME) + ) + ), + }, + num_epochs=self.config.num_epochs, + minibatch_size=self.config.minibatch_size, + shuffle_batch_per_epoch=self.config.shuffle_batch_per_epoch, + ) + self.metrics.aggregate(learner_results, key=LEARNER_RESULTS) + + # Update weights - after learning on the local worker - on all remote + # workers. + with self.metrics.log_time((TIMERS, SYNCH_WORKER_WEIGHTS_TIMER)): + # The train results's loss keys are ModuleIDs to their loss values. + # But we also return a total_loss key at the same level as the ModuleID + # keys. So we need to subtract that to get the correct set of ModuleIDs to + # update. + # TODO (sven): We should not be using `learner_results` as a messenger + # to infer which modules to update. `policies_to_train` might also NOT work + # as it might be a very large set (100s of Modules) vs a smaller Modules + # set that's present in the current train batch. + modules_to_update = set(learner_results[0].keys()) - {ALL_MODULES} + self.env_runner_group.sync_weights( + # Sync weights from learner_group to all EnvRunners. + from_worker_or_learner_group=self.learner_group, + policies=modules_to_update, + inference_only=True, + ) + + @OldAPIStack + def _training_step_old_api_stack(self) -> ResultDict: + # Collect batches from sample workers until we have a full batch. + with self._timers[SAMPLE_TIMER]: + if self.config.count_steps_by == "agent_steps": + train_batch = synchronous_parallel_sample( + worker_set=self.env_runner_group, + max_agent_steps=self.config.total_train_batch_size, + sample_timeout_s=self.config.sample_timeout_s, + ) + else: + train_batch = synchronous_parallel_sample( + worker_set=self.env_runner_group, + max_env_steps=self.config.total_train_batch_size, + sample_timeout_s=self.config.sample_timeout_s, + ) + # Return early if all our workers failed. + if not train_batch: + return {} + train_batch = train_batch.as_multi_agent() + self._counters[NUM_AGENT_STEPS_SAMPLED] += train_batch.agent_steps() + self._counters[NUM_ENV_STEPS_SAMPLED] += train_batch.env_steps() + # Standardize advantages. + train_batch = standardize_fields(train_batch, ["advantages"]) + + if self.config.simple_optimizer: + train_results = train_one_step(self, train_batch) + else: + train_results = multi_gpu_train_one_step(self, train_batch) + + policies_to_update = list(train_results.keys()) + + global_vars = { + "timestep": self._counters[NUM_AGENT_STEPS_SAMPLED], + # TODO (sven): num_grad_updates per each policy should be + # accessible via `train_results` (and get rid of global_vars). + "num_grad_updates_per_policy": { + pid: self.env_runner.policy_map[pid].num_grad_updates + for pid in policies_to_update + }, + } + + # Update weights - after learning on the local worker - on all remote + # workers. + with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]: + if self.env_runner_group.num_remote_workers() > 0: + from_worker_or_learner_group = None + self.env_runner_group.sync_weights( + from_worker_or_learner_group=from_worker_or_learner_group, + policies=policies_to_update, + global_vars=global_vars, + ) + + # For each policy: Update KL scale and warn about possible issues + for policy_id, policy_info in train_results.items(): + # Update KL loss with dynamic scaling + # for each (possibly multiagent) policy we are training + kl_divergence = policy_info[LEARNER_STATS_KEY].get("kl") + self.get_policy(policy_id).update_kl(kl_divergence) + + # Warn about excessively high value function loss + scaled_vf_loss = ( + self.config.vf_loss_coeff * policy_info[LEARNER_STATS_KEY]["vf_loss"] + ) + policy_loss = policy_info[LEARNER_STATS_KEY]["policy_loss"] + if ( + log_once("ppo_warned_lr_ratio") + and self.config.get("model", {}).get("vf_share_layers") + and scaled_vf_loss > 100 + ): + logger.warning( + "The magnitude of your value function loss for policy: {} is " + "extremely large ({}) compared to the policy loss ({}). This " + "can prevent the policy from learning. Consider scaling down " + "the VF loss by reducing vf_loss_coeff, or disabling " + "vf_share_layers.".format(policy_id, scaled_vf_loss, policy_loss) + ) + # Warn about bad clipping configs. + train_batch.policy_batches[policy_id].set_get_interceptor(None) + mean_reward = train_batch.policy_batches[policy_id]["rewards"].mean() + if ( + log_once("ppo_warned_vf_clip") + and mean_reward > self.config.vf_clip_param + ): + self.warned_vf_clip = True + logger.warning( + f"The mean reward returned from the environment is {mean_reward}" + f" but the vf_clip_param is set to {self.config['vf_clip_param']}." + f" Consider increasing it for policy: {policy_id} to improve" + " value function convergence." + ) + + # Update global vars on local worker as well. + # TODO (simon): At least in RolloutWorker obsolete I guess as called in + # `sync_weights()` called above if remote workers. Can we call this + # where `set_weights()` is called on the local_worker? + self.env_runner.set_global_vars(global_vars) + + return train_results diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_catalog.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..e88b761427a2b4fbb4ed6b3365a2f92a909c5395 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_catalog.py @@ -0,0 +1,195 @@ +# __sphinx_doc_begin__ +import gymnasium as gym + +from ray.rllib.core.models.base import ActorCriticEncoder, Encoder, Model +from ray.rllib.core.models.catalog import Catalog +from ray.rllib.core.models.configs import ( + ActorCriticEncoderConfig, + FreeLogStdMLPHeadConfig, + MLPHeadConfig, +) +from ray.rllib.utils import override +from ray.rllib.utils.annotations import OverrideToImplementCustomLogic + + +def _check_if_diag_gaussian(action_distribution_cls, framework, no_error=False): + if framework == "torch": + from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchDiagGaussian, + ) + + is_diag_gaussian = issubclass(action_distribution_cls, TorchDiagGaussian) + if no_error: + return is_diag_gaussian + else: + assert is_diag_gaussian, ( + f"free_log_std is only supported for DiagGaussian action " + f"distributions. Found action distribution: {action_distribution_cls}." + ) + else: + raise ValueError(f"Framework {framework} not supported for free_log_std.") + + +class PPOCatalog(Catalog): + """The Catalog class used to build models for PPO. + + PPOCatalog provides the following models: + - ActorCriticEncoder: The encoder used to encode the observations. + - Pi Head: The head used to compute the policy logits. + - Value Function Head: The head used to compute the value function. + + The ActorCriticEncoder is a wrapper around Encoders to produce separate outputs + for the policy and value function. See implementations of DefaultPPORLModule for + more details. + + Any custom ActorCriticEncoder can be built by overriding the + build_actor_critic_encoder() method. Alternatively, the ActorCriticEncoderConfig + at PPOCatalog.actor_critic_encoder_config can be overridden to build a custom + ActorCriticEncoder during RLModule runtime. + + Any custom head can be built by overriding the build_pi_head() and build_vf_head() + methods. Alternatively, the PiHeadConfig and VfHeadConfig can be overridden to + build custom heads during RLModule runtime. + + Any module built for exploration or inference is built with the flag + `ìnference_only=True` and does not contain a value network. This flag can be set + in the `SingleAgentModuleSpec` through the `inference_only` boolean flag. + In case that the actor-critic-encoder is not shared between the policy and value + function, the inference-only module will contain only the actor encoder network. + """ + + def __init__( + self, + observation_space: gym.Space, + action_space: gym.Space, + model_config_dict: dict, + ): + """Initializes the PPOCatalog. + + Args: + observation_space: The observation space of the Encoder. + action_space: The action space for the Pi Head. + model_config_dict: The model config to use. + """ + super().__init__( + observation_space=observation_space, + action_space=action_space, + model_config_dict=model_config_dict, + ) + # Replace EncoderConfig by ActorCriticEncoderConfig + self.actor_critic_encoder_config = ActorCriticEncoderConfig( + base_encoder_config=self._encoder_config, + shared=self._model_config_dict["vf_share_layers"], + ) + + self.pi_and_vf_head_hiddens = self._model_config_dict["head_fcnet_hiddens"] + self.pi_and_vf_head_activation = self._model_config_dict[ + "head_fcnet_activation" + ] + + # We don't have the exact (framework specific) action dist class yet and thus + # cannot determine the exact number of output nodes (action space) required. + # -> Build pi config only in the `self.build_pi_head` method. + self.pi_head_config = None + + self.vf_head_config = MLPHeadConfig( + input_dims=self.latent_dims, + hidden_layer_dims=self.pi_and_vf_head_hiddens, + hidden_layer_activation=self.pi_and_vf_head_activation, + output_layer_activation="linear", + output_layer_dim=1, + ) + + @OverrideToImplementCustomLogic + def build_actor_critic_encoder(self, framework: str) -> ActorCriticEncoder: + """Builds the ActorCriticEncoder. + + The default behavior is to build the encoder from the encoder_config. + This can be overridden to build a custom ActorCriticEncoder as a means of + configuring the behavior of a PPORLModule implementation. + + Args: + framework: The framework to use. Either "torch" or "tf2". + + Returns: + The ActorCriticEncoder. + """ + return self.actor_critic_encoder_config.build(framework=framework) + + @override(Catalog) + def build_encoder(self, framework: str) -> Encoder: + """Builds the encoder. + + Since PPO uses an ActorCriticEncoder, this method should not be implemented. + """ + raise NotImplementedError( + "Use PPOCatalog.build_actor_critic_encoder() instead for PPO." + ) + + @OverrideToImplementCustomLogic + def build_pi_head(self, framework: str) -> Model: + """Builds the policy head. + + The default behavior is to build the head from the pi_head_config. + This can be overridden to build a custom policy head as a means of configuring + the behavior of a PPORLModule implementation. + + Args: + framework: The framework to use. Either "torch" or "tf2". + + Returns: + The policy head. + """ + # Get action_distribution_cls to find out about the output dimension for pi_head + action_distribution_cls = self.get_action_dist_cls(framework=framework) + if self._model_config_dict["free_log_std"]: + _check_if_diag_gaussian( + action_distribution_cls=action_distribution_cls, framework=framework + ) + is_diag_gaussian = True + else: + is_diag_gaussian = _check_if_diag_gaussian( + action_distribution_cls=action_distribution_cls, + framework=framework, + no_error=True, + ) + required_output_dim = action_distribution_cls.required_input_dim( + space=self.action_space, model_config=self._model_config_dict + ) + # Now that we have the action dist class and number of outputs, we can define + # our pi-config and build the pi head. + pi_head_config_class = ( + FreeLogStdMLPHeadConfig + if self._model_config_dict["free_log_std"] + else MLPHeadConfig + ) + self.pi_head_config = pi_head_config_class( + input_dims=self.latent_dims, + hidden_layer_dims=self.pi_and_vf_head_hiddens, + hidden_layer_activation=self.pi_and_vf_head_activation, + output_layer_dim=required_output_dim, + output_layer_activation="linear", + clip_log_std=is_diag_gaussian, + log_std_clip_param=self._model_config_dict.get("log_std_clip_param", 20), + ) + + return self.pi_head_config.build(framework=framework) + + @OverrideToImplementCustomLogic + def build_vf_head(self, framework: str) -> Model: + """Builds the value function head. + + The default behavior is to build the head from the vf_head_config. + This can be overridden to build a custom value function head as a means of + configuring the behavior of a PPORLModule implementation. + + Args: + framework: The framework to use. Either "torch" or "tf2". + + Returns: + The value function head. + """ + return self.vf_head_config.build(framework=framework) + + +# __sphinx_doc_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..ef16f71c98bb73cb41252200b785879664127932 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_learner.py @@ -0,0 +1,149 @@ +import abc +from typing import Any, Dict + +from ray.rllib.algorithms.ppo.ppo import ( + LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY, + LEARNER_RESULTS_KL_KEY, + PPOConfig, +) +from ray.rllib.connectors.learner import ( + AddOneTsToEpisodesAndTruncate, + GeneralAdvantageEstimation, +) +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict +from ray.rllib.utils.metrics import ( + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_MODULE_STEPS_TRAINED, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.schedules.scheduler import Scheduler +from ray.rllib.utils.typing import ModuleID, TensorType + + +class PPOLearner(Learner): + @override(Learner) + def build(self) -> None: + super().build() + + # Dict mapping module IDs to the respective entropy Scheduler instance. + self.entropy_coeff_schedulers_per_module: Dict[ + ModuleID, Scheduler + ] = LambdaDefaultDict( + lambda module_id: Scheduler( + fixed_value_or_schedule=( + self.config.get_config_for_module(module_id).entropy_coeff + ), + framework=self.framework, + device=self._device, + ) + ) + + # Set up KL coefficient variables (per module). + # Note that the KL coeff is not controlled by a Scheduler, but seeks + # to stay close to a given kl_target value. + self.curr_kl_coeffs_per_module: Dict[ModuleID, TensorType] = LambdaDefaultDict( + lambda module_id: self._get_tensor_variable( + self.config.get_config_for_module(module_id).kl_coeff + ) + ) + + # Extend all episodes by one artificial timestep to allow the value function net + # to compute the bootstrap values (and add a mask to the batch to know, which + # slots to mask out). + if ( + self._learner_connector is not None + and self.config.add_default_connectors_to_learner_pipeline + ): + # Before anything, add one ts to each episode (and record this in the loss + # mask, so that the computations at this extra ts are not used to compute + # the loss). + self._learner_connector.prepend(AddOneTsToEpisodesAndTruncate()) + # At the end of the pipeline (when the batch is already completed), add the + # GAE connector, which performs a vf forward pass, then computes the GAE + # computations, and puts the results of this (advantages, value targets) + # directly back in the batch. This is then the batch used for + # `forward_train` and `compute_losses`. + self._learner_connector.append( + GeneralAdvantageEstimation( + gamma=self.config.gamma, lambda_=self.config.lambda_ + ) + ) + + @override(Learner) + def remove_module(self, module_id: ModuleID, **kwargs): + marl_spec = super().remove_module(module_id, **kwargs) + + self.entropy_coeff_schedulers_per_module.pop(module_id, None) + self.curr_kl_coeffs_per_module.pop(module_id, None) + + return marl_spec + + @OverrideToImplementCustomLogic_CallToSuperRecommended + @override(Learner) + def after_gradient_based_update( + self, + *, + timesteps: Dict[str, Any], + ) -> None: + super().after_gradient_based_update(timesteps=timesteps) + + for module_id, module in self.module._rl_modules.items(): + config = self.config.get_config_for_module(module_id) + + # Update entropy coefficient via our Scheduler. + new_entropy_coeff = self.entropy_coeff_schedulers_per_module[ + module_id + ].update(timestep=timesteps.get(NUM_ENV_STEPS_SAMPLED_LIFETIME, 0)) + self.metrics.log_value( + (module_id, LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY), + new_entropy_coeff, + window=1, + ) + if ( + config.use_kl_loss + and self.metrics.peek((module_id, NUM_MODULE_STEPS_TRAINED), default=0) + > 0 + and (module_id, LEARNER_RESULTS_KL_KEY) in self.metrics + ): + kl_loss = convert_to_numpy( + self.metrics.peek((module_id, LEARNER_RESULTS_KL_KEY)) + ) + self._update_module_kl_coeff( + module_id=module_id, + config=config, + kl_loss=kl_loss, + ) + + @classmethod + @override(Learner) + def rl_module_required_apis(cls) -> list[type]: + # In order for a PPOLearner to update an RLModule, it must implement the + # following APIs: + return [ValueFunctionAPI] + + @abc.abstractmethod + def _update_module_kl_coeff( + self, + *, + module_id: ModuleID, + config: PPOConfig, + kl_loss: float, + ) -> None: + """Dynamically update the KL loss coefficients of each module. + + The update is completed using the mean KL divergence between the action + distributions current policy and old policy of each module. That action + distribution is computed during the most recent update/call to `compute_loss`. + + Args: + module_id: The module whose KL loss coefficient to update. + config: The AlgorithmConfig specific to the given `module_id`. + kl_loss: The mean KL loss of the module, computed inside + `compute_loss_for_module()`. + """ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..631bf29fdd6295beda8b4a4d9e2432f2e02fda04 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_rl_module.py @@ -0,0 +1,11 @@ +# Backward compat import. +from ray.rllib.algorithms.ppo.default_ppo_rl_module import ( # noqa + DefaultPPORLModule as PPORLModule, +) +from ray._common.deprecation import deprecation_warning + +deprecation_warning( + old="ray.rllib.algorithms.ppo.ppo_rl_module.PPORLModule", + new="ray.rllib.algorithms.ppo.default_ppo_rl_module.DefaultPPORLModule", + error=False, +) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_tf_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_tf_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..76e8d0161689a4e79fb9673cfce65dff8743354c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_tf_policy.py @@ -0,0 +1,235 @@ +""" +TensorFlow policy class used for PPO. +""" + +import logging +from typing import Dict, List, Type, Union + +from ray.rllib.evaluation.postprocessing import ( + Postprocessing, + compute_gae_for_sample_batch, +) +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import TFActionDistribution +from ray.rllib.policy.dynamic_tf_policy_v2 import DynamicTFPolicyV2 +from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_mixins import ( + EntropyCoeffSchedule, + KLCoeffMixin, + LearningRateSchedule, + ValueNetworkMixin, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.tf_utils import explained_variance, warn_if_infinite_kl_divergence +from ray.rllib.utils.typing import AlgorithmConfigDict, TensorType, TFPolicyV2Type + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +def validate_config(config: AlgorithmConfigDict) -> None: + """Executed before Policy is "initialized" (at beginning of constructor). + Args: + config: The Policy's config. + """ + # If vf_share_layers is True, inform about the need to tune vf_loss_coeff. + if config.get("model", {}).get("vf_share_layers") is True: + logger.info( + "`vf_share_layers=True` in your model. " + "Therefore, remember to tune the value of `vf_loss_coeff`!" + ) + + +# We need this builder function because we want to share the same +# custom logics between TF1 dynamic and TF2 eager policies. +def get_ppo_tf_policy(name: str, base: TFPolicyV2Type) -> TFPolicyV2Type: + """Construct a PPOTFPolicy inheriting either dynamic or eager base policies. + + Args: + base: Base class for this policy. DynamicTFPolicyV2 or EagerTFPolicyV2. + + Returns: + A TF Policy to be used with PPO. + """ + + class PPOTFPolicy( + EntropyCoeffSchedule, + LearningRateSchedule, + KLCoeffMixin, + ValueNetworkMixin, + base, + ): + def __init__( + self, + observation_space, + action_space, + config, + existing_model=None, + existing_inputs=None, + ): + # First thing first, enable eager execution if necessary. + base.enable_eager_execution_if_necessary() + + # TODO: Move into Policy API, if needed at all here. Why not move this into + # `PPOConfig`?. + validate_config(config) + + # Initialize base class. + base.__init__( + self, + observation_space, + action_space, + config, + existing_inputs=existing_inputs, + existing_model=existing_model, + ) + + # Initialize MixIns. + ValueNetworkMixin.__init__(self, config) + EntropyCoeffSchedule.__init__( + self, config["entropy_coeff"], config["entropy_coeff_schedule"] + ) + LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"]) + KLCoeffMixin.__init__(self, config) + + # Note: this is a bit ugly, but loss and optimizer initialization must + # happen after all the MixIns are initialized. + self.maybe_initialize_optimizer_and_loss() + + @override(base) + def loss( + self, + model: Union[ModelV2, "tf.keras.Model"], + dist_class: Type[TFActionDistribution], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + if isinstance(model, tf.keras.Model): + logits, state, extra_outs = model(train_batch) + value_fn_out = extra_outs[SampleBatch.VF_PREDS] + else: + logits, state = model(train_batch) + value_fn_out = model.value_function() + + curr_action_dist = dist_class(logits, model) + + # RNN case: Mask away 0-padded chunks at end of time axis. + if state: + # Derive max_seq_len from the data itself, not from the seq_lens + # tensor. This is in case e.g. seq_lens=[2, 3], but the data is still + # 0-padded up to T=5 (as it's the case for attention nets). + B = tf.shape(train_batch[SampleBatch.SEQ_LENS])[0] + max_seq_len = tf.shape(logits)[0] // B + + mask = tf.sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len) + mask = tf.reshape(mask, [-1]) + + def reduce_mean_valid(t): + return tf.reduce_mean(tf.boolean_mask(t, mask)) + + # non-RNN case: No masking. + else: + mask = None + reduce_mean_valid = tf.reduce_mean + + prev_action_dist = dist_class( + train_batch[SampleBatch.ACTION_DIST_INPUTS], model + ) + + logp_ratio = tf.exp( + curr_action_dist.logp(train_batch[SampleBatch.ACTIONS]) + - train_batch[SampleBatch.ACTION_LOGP] + ) + + # Only calculate kl loss if necessary (kl-coeff > 0.0). + if self.config["kl_coeff"] > 0.0: + action_kl = prev_action_dist.kl(curr_action_dist) + mean_kl_loss = reduce_mean_valid(action_kl) + warn_if_infinite_kl_divergence(self, mean_kl_loss) + else: + mean_kl_loss = tf.constant(0.0) + + curr_entropy = curr_action_dist.entropy() + mean_entropy = reduce_mean_valid(curr_entropy) + + surrogate_loss = tf.minimum( + train_batch[Postprocessing.ADVANTAGES] * logp_ratio, + train_batch[Postprocessing.ADVANTAGES] + * tf.clip_by_value( + logp_ratio, + 1 - self.config["clip_param"], + 1 + self.config["clip_param"], + ), + ) + + # Compute a value function loss. + if self.config["use_critic"]: + vf_loss = tf.math.square( + value_fn_out - train_batch[Postprocessing.VALUE_TARGETS] + ) + vf_loss_clipped = tf.clip_by_value( + vf_loss, + 0, + self.config["vf_clip_param"], + ) + mean_vf_loss = reduce_mean_valid(vf_loss_clipped) + # Ignore the value function. + else: + vf_loss_clipped = mean_vf_loss = tf.constant(0.0) + + total_loss = reduce_mean_valid( + -surrogate_loss + + self.config["vf_loss_coeff"] * vf_loss_clipped + - self.entropy_coeff * curr_entropy + ) + # Add mean_kl_loss (already processed through `reduce_mean_valid`), + # if necessary. + if self.config["kl_coeff"] > 0.0: + total_loss += self.kl_coeff * mean_kl_loss + + # Store stats in policy for stats_fn. + self._total_loss = total_loss + self._mean_policy_loss = reduce_mean_valid(-surrogate_loss) + self._mean_vf_loss = mean_vf_loss + self._mean_entropy = mean_entropy + # Backward compatibility: Deprecate self._mean_kl. + self._mean_kl_loss = self._mean_kl = mean_kl_loss + self._value_fn_out = value_fn_out + + return total_loss + + @override(base) + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + return { + "cur_kl_coeff": tf.cast(self.kl_coeff, tf.float64), + "cur_lr": tf.cast(self.cur_lr, tf.float64), + "total_loss": self._total_loss, + "policy_loss": self._mean_policy_loss, + "vf_loss": self._mean_vf_loss, + "vf_explained_var": explained_variance( + train_batch[Postprocessing.VALUE_TARGETS], self._value_fn_out + ), + "kl": self._mean_kl_loss, + "entropy": self._mean_entropy, + "entropy_coeff": tf.cast(self.entropy_coeff, tf.float64), + } + + @override(base) + def postprocess_trajectory( + self, sample_batch, other_agent_batches=None, episode=None + ): + sample_batch = super().postprocess_trajectory(sample_batch) + return compute_gae_for_sample_batch( + self, sample_batch, other_agent_batches, episode + ) + + PPOTFPolicy.__name__ = name + PPOTFPolicy.__qualname__ = name + + return PPOTFPolicy + + +PPOTF1Policy = get_ppo_tf_policy("PPOTF1Policy", DynamicTFPolicyV2) +PPOTF2Policy = get_ppo_tf_policy("PPOTF2Policy", EagerTFPolicyV2) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_torch_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_torch_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..26a52dbe4d2b1bcd6f66b501bff9d3c410cc5923 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_torch_policy.py @@ -0,0 +1,217 @@ +import logging +from typing import Dict, List, Type, Union + +import ray +from ray.rllib.algorithms.ppo.ppo_tf_policy import validate_config +from ray.rllib.evaluation.postprocessing import ( + Postprocessing, + compute_gae_for_sample_batch, +) +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_mixins import ( + EntropyCoeffSchedule, + KLCoeffMixin, + LearningRateSchedule, + ValueNetworkMixin, +) +from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.torch_utils import ( + apply_grad_clipping, + explained_variance, + sequence_mask, + warn_if_infinite_kl_divergence, +) +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +class PPOTorchPolicy( + ValueNetworkMixin, + LearningRateSchedule, + EntropyCoeffSchedule, + KLCoeffMixin, + TorchPolicyV2, +): + """PyTorch policy class used with PPO.""" + + def __init__(self, observation_space, action_space, config): + config = dict(ray.rllib.algorithms.ppo.ppo.PPOConfig().to_dict(), **config) + validate_config(config) + + TorchPolicyV2.__init__( + self, + observation_space, + action_space, + config, + max_seq_len=config["model"]["max_seq_len"], + ) + + ValueNetworkMixin.__init__(self, config) + LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"]) + EntropyCoeffSchedule.__init__( + self, config["entropy_coeff"], config["entropy_coeff_schedule"] + ) + KLCoeffMixin.__init__(self, config) + + self._initialize_loss_from_dummy_batch() + + @override(TorchPolicyV2) + def loss( + self, + model: ModelV2, + dist_class: Type[ActionDistribution], + train_batch: SampleBatch, + ) -> Union[TensorType, List[TensorType]]: + """Compute loss for Proximal Policy Objective. + + Args: + model: The Model to calculate the loss for. + dist_class: The action distr. class. + train_batch: The training data. + + Returns: + The PPO loss tensor given the input batch. + """ + + logits, state = model(train_batch) + curr_action_dist = dist_class(logits, model) + + # RNN case: Mask away 0-padded chunks at end of time axis. + if state: + B = len(train_batch[SampleBatch.SEQ_LENS]) + max_seq_len = logits.shape[0] // B + mask = sequence_mask( + train_batch[SampleBatch.SEQ_LENS], + max_seq_len, + time_major=model.is_time_major(), + ) + mask = torch.reshape(mask, [-1]) + num_valid = torch.sum(mask) + + def reduce_mean_valid(t): + return torch.sum(t[mask]) / num_valid + + # non-RNN case: No masking. + else: + mask = None + reduce_mean_valid = torch.mean + + prev_action_dist = dist_class( + train_batch[SampleBatch.ACTION_DIST_INPUTS], model + ) + + logp_ratio = torch.exp( + curr_action_dist.logp(train_batch[SampleBatch.ACTIONS]) + - train_batch[SampleBatch.ACTION_LOGP] + ) + + # Only calculate kl loss if necessary (kl-coeff > 0.0). + if self.config["kl_coeff"] > 0.0: + action_kl = prev_action_dist.kl(curr_action_dist) + mean_kl_loss = reduce_mean_valid(action_kl) + # TODO smorad: should we do anything besides warn? Could discard KL term + # for this update + warn_if_infinite_kl_divergence(self, mean_kl_loss) + else: + mean_kl_loss = torch.tensor(0.0, device=logp_ratio.device) + + curr_entropy = curr_action_dist.entropy() + mean_entropy = reduce_mean_valid(curr_entropy) + + surrogate_loss = torch.min( + train_batch[Postprocessing.ADVANTAGES] * logp_ratio, + train_batch[Postprocessing.ADVANTAGES] + * torch.clamp( + logp_ratio, 1 - self.config["clip_param"], 1 + self.config["clip_param"] + ), + ) + + # Compute a value function loss. + if self.config["use_critic"]: + value_fn_out = model.value_function() + vf_loss = torch.pow( + value_fn_out - train_batch[Postprocessing.VALUE_TARGETS], 2.0 + ) + vf_loss_clipped = torch.clamp(vf_loss, 0, self.config["vf_clip_param"]) + mean_vf_loss = reduce_mean_valid(vf_loss_clipped) + # Ignore the value function. + else: + value_fn_out = torch.tensor(0.0).to(surrogate_loss.device) + vf_loss_clipped = mean_vf_loss = torch.tensor(0.0).to(surrogate_loss.device) + + total_loss = reduce_mean_valid( + -surrogate_loss + + self.config["vf_loss_coeff"] * vf_loss_clipped + - self.entropy_coeff * curr_entropy + ) + + # Add mean_kl_loss (already processed through `reduce_mean_valid`), + # if necessary. + if self.config["kl_coeff"] > 0.0: + total_loss += self.kl_coeff * mean_kl_loss + + # Store values for stats function in model (tower), such that for + # multi-GPU, we do not override them during the parallel loss phase. + model.tower_stats["total_loss"] = total_loss + model.tower_stats["mean_policy_loss"] = reduce_mean_valid(-surrogate_loss) + model.tower_stats["mean_vf_loss"] = mean_vf_loss + model.tower_stats["vf_explained_var"] = explained_variance( + train_batch[Postprocessing.VALUE_TARGETS], value_fn_out + ) + model.tower_stats["mean_entropy"] = mean_entropy + model.tower_stats["mean_kl_loss"] = mean_kl_loss + + return total_loss + + # TODO: Make this an event-style subscription (e.g.: + # "after_gradients_computed"). + @override(TorchPolicyV2) + def extra_grad_process(self, local_optimizer, loss): + return apply_grad_clipping(self, local_optimizer, loss) + + @override(TorchPolicyV2) + def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]: + return convert_to_numpy( + { + "cur_kl_coeff": self.kl_coeff, + "cur_lr": self.cur_lr, + "total_loss": torch.mean( + torch.stack(self.get_tower_stats("total_loss")) + ), + "policy_loss": torch.mean( + torch.stack(self.get_tower_stats("mean_policy_loss")) + ), + "vf_loss": torch.mean( + torch.stack(self.get_tower_stats("mean_vf_loss")) + ), + "vf_explained_var": torch.mean( + torch.stack(self.get_tower_stats("vf_explained_var")) + ), + "kl": torch.mean(torch.stack(self.get_tower_stats("mean_kl_loss"))), + "entropy": torch.mean( + torch.stack(self.get_tower_stats("mean_entropy")) + ), + "entropy_coeff": self.entropy_coeff, + } + ) + + @override(TorchPolicyV2) + def postprocess_trajectory( + self, sample_batch, other_agent_batches=None, episode=None + ): + # Do all post-processing always with no_grad(). + # Not using this here will introduce a memory leak + # in torch (issue #6962). + # TODO: no_grad still necessary? + with torch.no_grad(): + return compute_gae_for_sample_batch( + self, sample_batch, other_agent_batches, episode + ) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cedf2cd9bfd97225468f7da1ef38057492e18590 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/default_ppo_torch_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/default_ppo_torch_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cfa34a67ca77268e176556b8bb7d5771b9cd4a7f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/default_ppo_torch_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/ppo_torch_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/ppo_torch_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5094973405bdbaa36ef3f8e3ba346983848c09cf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/ppo_torch_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/ppo_torch_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/ppo_torch_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0dccfcaefe6e74871260df5bc95b482afb4fe4dd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/__pycache__/ppo_torch_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/default_ppo_torch_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/default_ppo_torch_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..0ca3093887e62562f7a9697db8b77f3ba1d6d44b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/default_ppo_torch_rl_module.py @@ -0,0 +1,73 @@ +from typing import Any, Dict, Optional + +from ray.rllib.algorithms.ppo.default_ppo_rl_module import DefaultPPORLModule +from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog +from ray.rllib.core.columns import Columns +from ray.rllib.core.models.base import ACTOR, CRITIC, ENCODER_OUT +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.core.rl_module.torch import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType +from ray.util.annotations import DeveloperAPI + +torch, nn = try_import_torch() + + +@DeveloperAPI +class DefaultPPOTorchRLModule(TorchRLModule, DefaultPPORLModule): + def __init__(self, *args, **kwargs): + catalog_class = kwargs.pop("catalog_class", None) + if catalog_class is None: + catalog_class = PPOCatalog + super().__init__(*args, **kwargs, catalog_class=catalog_class) + + @override(RLModule) + def _forward(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """Default forward pass (used for inference and exploration).""" + output = {} + # Encoder forward pass. + encoder_outs = self.encoder(batch) + # Stateful encoder? + if Columns.STATE_OUT in encoder_outs: + output[Columns.STATE_OUT] = encoder_outs[Columns.STATE_OUT] + # Pi head. + output[Columns.ACTION_DIST_INPUTS] = self.pi(encoder_outs[ENCODER_OUT][ACTOR]) + return output + + @override(RLModule) + def _forward_train(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + """Train forward pass (keep embeddings for possible shared value func. call).""" + output = {} + encoder_outs = self.encoder(batch) + output[Columns.EMBEDDINGS] = encoder_outs[ENCODER_OUT][CRITIC] + if Columns.STATE_OUT in encoder_outs: + output[Columns.STATE_OUT] = encoder_outs[Columns.STATE_OUT] + output[Columns.ACTION_DIST_INPUTS] = self.pi(encoder_outs[ENCODER_OUT][ACTOR]) + return output + + @override(ValueFunctionAPI) + def compute_values( + self, + batch: Dict[str, Any], + embeddings: Optional[Any] = None, + ) -> TensorType: + if embeddings is None: + # Separate vf-encoder. + if hasattr(self.encoder, "critic_encoder"): + batch_ = batch + if self.is_stateful(): + # The recurrent encoders expect a `(state_in, h)` key in the + # input dict while the key returned is `(state_in, critic, h)`. + batch_ = batch.copy() + batch_[Columns.STATE_IN] = batch[Columns.STATE_IN][CRITIC] + embeddings = self.encoder.critic_encoder(batch_)[ENCODER_OUT] + # Shared encoder. + else: + embeddings = self.encoder(batch)[ENCODER_OUT][CRITIC] + + # Value head. + vf_out = self.vf(embeddings) + # Squeeze out last dimension (single node value head). + return vf_out.squeeze(-1) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/ppo_torch_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/ppo_torch_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..4e7a806f98ab79d9fea47b8f1904ffc8f74a1575 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/ppo_torch_learner.py @@ -0,0 +1,173 @@ +import logging +from typing import Any, Dict + +import numpy as np + +from ray.rllib.algorithms.ppo.ppo import ( + LEARNER_RESULTS_CURR_KL_COEFF_KEY, + LEARNER_RESULTS_KL_KEY, + LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY, + LEARNER_RESULTS_VF_LOSS_UNCLIPPED_KEY, + PPOConfig, +) +from ray.rllib.algorithms.ppo.ppo_learner import PPOLearner +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.learner import ENTROPY_KEY, POLICY_LOSS_KEY, VF_LOSS_KEY +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.evaluation.postprocessing import Postprocessing +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.torch_utils import explained_variance +from ray.rllib.utils.typing import ModuleID, TensorType + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +class PPOTorchLearner(PPOLearner, TorchLearner): + """Implements torch-specific PPO loss logic on top of PPOLearner. + + This class implements the ppo loss under `self.compute_loss_for_module()`. + """ + + @override(TorchLearner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: PPOConfig, + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + module = self.module[module_id].unwrapped() + + # Possibly apply masking to some sub loss terms and to the total loss term + # at the end. Masking could be used for RNN-based model (zero padded `batch`) + # and for PPO's batched value function (and bootstrap value) computations, + # for which we add an (artificial) timestep to each episode to + # simplify the actual computation. + if Columns.LOSS_MASK in batch: + mask = batch[Columns.LOSS_MASK] + num_valid = torch.sum(mask) + + def possibly_masked_mean(data_): + return torch.sum(data_[mask]) / num_valid + + else: + possibly_masked_mean = torch.mean + + action_dist_class_train = module.get_train_action_dist_cls() + action_dist_class_exploration = module.get_exploration_action_dist_cls() + + curr_action_dist = action_dist_class_train.from_logits( + fwd_out[Columns.ACTION_DIST_INPUTS] + ) + # TODO (sven): We should ideally do this in the LearnerConnector (separation of + # concerns: Only do things on the EnvRunners that are required for computing + # actions, do NOT do anything on the EnvRunners that's only required for a + # training update). + prev_action_dist = action_dist_class_exploration.from_logits( + batch[Columns.ACTION_DIST_INPUTS] + ) + + logp_ratio = torch.exp( + curr_action_dist.logp(batch[Columns.ACTIONS]) - batch[Columns.ACTION_LOGP] + ) + + # Only calculate kl loss if necessary (kl-coeff > 0.0). + if config.use_kl_loss: + action_kl = prev_action_dist.kl(curr_action_dist) + mean_kl_loss = possibly_masked_mean(action_kl) + else: + mean_kl_loss = torch.tensor(0.0, device=logp_ratio.device) + + curr_entropy = curr_action_dist.entropy() + mean_entropy = possibly_masked_mean(curr_entropy) + + surrogate_loss = torch.min( + batch[Postprocessing.ADVANTAGES] * logp_ratio, + batch[Postprocessing.ADVANTAGES] + * torch.clamp(logp_ratio, 1 - config.clip_param, 1 + config.clip_param), + ) + + # Compute a value function loss. + if config.use_critic: + value_fn_out = module.compute_values( + batch, embeddings=fwd_out.get(Columns.EMBEDDINGS) + ) + vf_loss = torch.pow(value_fn_out - batch[Postprocessing.VALUE_TARGETS], 2.0) + vf_loss_clipped = torch.clamp(vf_loss, 0, config.vf_clip_param) + mean_vf_loss = possibly_masked_mean(vf_loss_clipped) + mean_vf_unclipped_loss = possibly_masked_mean(vf_loss) + # Ignore the value function -> Set all to 0.0. + else: + z = torch.tensor(0.0, device=surrogate_loss.device) + value_fn_out = mean_vf_unclipped_loss = vf_loss_clipped = mean_vf_loss = z + + total_loss = possibly_masked_mean( + -surrogate_loss + + config.vf_loss_coeff * vf_loss_clipped + - ( + self.entropy_coeff_schedulers_per_module[module_id].get_current_value() + * curr_entropy + ) + ) + + # Add mean_kl_loss (already processed through `possibly_masked_mean`), + # if necessary. + if config.use_kl_loss: + total_loss += self.curr_kl_coeffs_per_module[module_id] * mean_kl_loss + + # Log important loss stats. + self.metrics.log_dict( + { + POLICY_LOSS_KEY: -possibly_masked_mean(surrogate_loss), + VF_LOSS_KEY: mean_vf_loss, + LEARNER_RESULTS_VF_LOSS_UNCLIPPED_KEY: mean_vf_unclipped_loss, + LEARNER_RESULTS_VF_EXPLAINED_VAR_KEY: explained_variance( + batch[Postprocessing.VALUE_TARGETS], value_fn_out + ), + ENTROPY_KEY: mean_entropy, + LEARNER_RESULTS_KL_KEY: mean_kl_loss, + }, + key=module_id, + window=1, # <- single items (should not be mean/ema-reduced over time). + ) + # Return the total loss. + return total_loss + + @override(PPOLearner) + def _update_module_kl_coeff( + self, + *, + module_id: ModuleID, + config: PPOConfig, + kl_loss: float, + ) -> None: + if np.isnan(kl_loss): + logger.warning( + f"KL divergence for Module {module_id} is non-finite, this " + "will likely destabilize your model and the training " + "process. Action(s) in a specific state have near-zero " + "probability. This can happen naturally in deterministic " + "environments where the optimal policy has zero mass for a " + "specific action. To fix this issue, consider setting " + "`kl_coeff` to 0.0 or increasing `entropy_coeff` in your " + "config." + ) + + # Update the KL coefficient. + curr_var = self.curr_kl_coeffs_per_module[module_id] + if kl_loss > 2.0 * config.kl_target: + # TODO (Kourosh) why not 2? + curr_var.data *= 1.5 + elif kl_loss < 0.5 * config.kl_target: + curr_var.data *= 0.5 + + # Log the updated KL-coeff value. + self.metrics.log_value( + (module_id, LEARNER_RESULTS_CURR_KL_COEFF_KEY), + curr_var.item(), + window=1, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/ppo_torch_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/ppo_torch_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..66acb9e5fb5ad722ca54e3ae9534d800ccf9271e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/torch/ppo_torch_rl_module.py @@ -0,0 +1,13 @@ +# Backward compat import. +from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( # noqa + DefaultPPOTorchRLModule as PPOTorchRLModule, +) +from ray._common.deprecation import deprecation_warning + + +deprecation_warning( + old="ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module.PPOTorchRLModule", + new="ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module." + "DefaultPPOTorchRLModule", + error=False, +) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/registry.py b/lib/python3.12/site-packages/ray/rllib/algorithms/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..c349489d165c27e920928bf8d6b5d4f70362300a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/registry.py @@ -0,0 +1,185 @@ +"""Registry of algorithm names for tune.Tuner(trainable=[..]).""" + +import importlib +import re + + +def _import_appo(): + import ray.rllib.algorithms.appo as appo + + return appo.APPO, appo.APPO.get_default_config() + + +def _import_bc(): + import ray.rllib.algorithms.bc as bc + + return bc.BC, bc.BC.get_default_config() + + +def _import_cql(): + import ray.rllib.algorithms.cql as cql + + return cql.CQL, cql.CQL.get_default_config() + + +def _import_dqn(): + import ray.rllib.algorithms.dqn as dqn + + return dqn.DQN, dqn.DQN.get_default_config() + + +def _import_dreamerv3(): + import ray.rllib.algorithms.dreamerv3 as dreamerv3 + + return dreamerv3.DreamerV3, dreamerv3.DreamerV3.get_default_config() + + +def _import_impala(): + import ray.rllib.algorithms.impala as impala + + return impala.IMPALA, impala.IMPALA.get_default_config() + + +def _import_iql(): + import ray.rllib.algorithms.iql as iql + + return iql.IQL, iql.IQL.get_default_config() + + +def _import_marwil(): + import ray.rllib.algorithms.marwil as marwil + + return marwil.MARWIL, marwil.MARWIL.get_default_config() + + +def _import_ppo(): + import ray.rllib.algorithms.ppo as ppo + + return ppo.PPO, ppo.PPO.get_default_config() + + +def _import_sac(): + import ray.rllib.algorithms.sac as sac + + return sac.SAC, sac.SAC.get_default_config() + + +ALGORITHMS = { + "APPO": _import_appo, + "BC": _import_bc, + "CQL": _import_cql, + "DQN": _import_dqn, + "DreamerV3": _import_dreamerv3, + "IMPALA": _import_impala, + "IQL": _import_iql, + "MARWIL": _import_marwil, + "PPO": _import_ppo, + "SAC": _import_sac, +} + + +ALGORITHMS_CLASS_TO_NAME = { + "APPO": "APPO", + "BC": "BC", + "CQL": "CQL", + "DQN": "DQN", + "DreamerV3": "DreamerV3", + "Impala": "IMPALA", + "IQL": "IQL", + "IMPALA": "IMPALA", + "MARWIL": "MARWIL", + "PPO": "PPO", + "SAC": "SAC", +} + + +def _get_algorithm_class(alg: str) -> type: + # This helps us get around a circular import (tune calls rllib._register_all when + # checking if a rllib Trainable is registered) + if alg in ALGORITHMS: + return ALGORITHMS[alg]()[0] + elif alg == "script": + from ray.tune import script_runner + + return script_runner.ScriptRunner + elif alg == "__fake": + from ray.rllib.algorithms.mock import _MockTrainer + + return _MockTrainer + elif alg == "__sigmoid_fake_data": + from ray.rllib.algorithms.mock import _SigmoidFakeData + + return _SigmoidFakeData + elif alg == "__parameter_tuning": + from ray.rllib.algorithms.mock import _ParameterTuningTrainer + + return _ParameterTuningTrainer + else: + raise Exception("Unknown algorithm {}.".format(alg)) + + +# Dict mapping policy names to where the class is located, relative to rllib.algorithms. +# TODO(jungong) : Finish migrating all the policies to PolicyV2, so we can list +# all the TF eager policies here. +POLICIES = { + "APPOTF1Policy": "appo.appo_tf_policy", + "APPOTF2Policy": "appo.appo_tf_policy", + "APPOTorchPolicy": "appo.appo_torch_policy", + "CQLTFPolicy": "cql.cql_tf_policy", + "CQLTorchPolicy": "cql.cql_torch_policy", + "DQNTFPolicy": "dqn.dqn_tf_policy", + "DQNTorchPolicy": "dqn.dqn_torch_policy", + "ImpalaTF1Policy": "impala.impala_tf_policy", + "ImpalaTF2Policy": "impala.impala_tf_policy", + "ImpalaTorchPolicy": "impala.impala_torch_policy", + "MARWILTF1Policy": "marwil.marwil_tf_policy", + "MARWILTF2Policy": "marwil.marwil_tf_policy", + "MARWILTorchPolicy": "marwil.marwil_torch_policy", + "SACTFPolicy": "sac.sac_tf_policy", + "SACTorchPolicy": "sac.sac_torch_policy", + "PPOTF1Policy": "ppo.ppo_tf_policy", + "PPOTF2Policy": "ppo.ppo_tf_policy", + "PPOTorchPolicy": "ppo.ppo_torch_policy", +} + + +def get_policy_class_name(policy_class: type): + """Returns a string name for the provided policy class. + + Args: + policy_class: RLlib policy class, e.g. A3CTorchPolicy, DQNTFPolicy, etc. + + Returns: + A string name uniquely mapped to the given policy class. + """ + # TF2 policy classes may get automatically converted into new class types + # that have eager tracing capability. + # These policy classes have the "_traced" postfix in their names. + # When checkpointing these policy classes, we should save the name of the + # original policy class instead. So that users have the choice of turning + # on eager tracing during inference time. + name = re.sub("_traced$", "", policy_class.__name__) + if name in POLICIES: + return name + return None + + +def get_policy_class(name: str): + """Return an actual policy class given the string name. + + Args: + name: string name of the policy class. + + Returns: + Actual policy class for the given name. + """ + if name not in POLICIES: + return None + + path = POLICIES[name] + module = importlib.import_module("ray.rllib.algorithms." + path) + + if not hasattr(module, name): + return None + + return getattr(module, name) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8170f353edd12a02ae10ea205b21520f78d97b27 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__init__.py @@ -0,0 +1,10 @@ +from ray.rllib.algorithms.sac.sac import SAC, SACConfig +from ray.rllib.algorithms.sac.sac_tf_policy import SACTFPolicy +from ray.rllib.algorithms.sac.sac_torch_policy import SACTorchPolicy + +__all__ = [ + "SAC", + "SACTFPolicy", + "SACTorchPolicy", + "SACConfig", +] diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0f9407e273d602f31dbb50da150a74f997d1ca4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/default_sac_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/default_sac_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c95d5519ed1eeca3e998d2a1f50b2090fc18baa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/default_sac_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c44c2104d1df26869d84c460a53a9c868b335329 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_catalog.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_catalog.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd842aab5118f03832adbaf9ef141643ff286e40 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_catalog.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..773cc57afa2a426f0c6b04f0395f630f4acecd68 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_tf_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_tf_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c817d4006d9823518a320da29848adb6f3fb3d0f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_tf_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_tf_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_tf_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a27e9fa9e103d5311d49370958bbef032906b003 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_tf_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_torch_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_torch_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9812d39e8f1f028bc6209a3f227ff33c4e9aeffa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_torch_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_torch_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_torch_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77a2617da4d921acf879d67c622b0f01b0cc0637 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/__pycache__/sac_torch_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/default_sac_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/default_sac_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..76f02a1e4c7f1fa662ce9834dae40c00990cd1f7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/default_sac_rl_module.py @@ -0,0 +1,141 @@ +from abc import abstractmethod +from typing import Any, Dict, List, Tuple + +from ray.rllib.core.learner.utils import make_target_network +from ray.rllib.core.models.base import Encoder, Model +from ray.rllib.core.rl_module.apis import InferenceOnlyAPI, QNetAPI, TargetNetworkAPI +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + override, +) +from ray.rllib.utils.typing import NetworkType +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class DefaultSACRLModule(RLModule, InferenceOnlyAPI, TargetNetworkAPI, QNetAPI): + """`RLModule` for the Soft-Actor-Critic (SAC) algorithm. + + It consists of several architectures, each in turn composed of + two networks: an encoder and a head. + + The policy (actor) contains a state encoder (`pi_encoder`) and + a head (`pi_head`) that feeds into an action distribution (a + squashed Gaussian, i.e. outputs define the location and the log + scale parameters). + + In addition, two (or four in case `twin_q=True`) Q networks are + defined, the second one (and fourth, if `twin_q=True`) of them the + Q target network(s). All of these in turn are - similar to the + policy network - composed of an encoder and a head network. Each of + the encoders forms a state-action encoding that feeds into the + corresponding value heads to result in an estimation of the soft + action-value of SAC. + + The following graphics show the forward passes through this module: + [obs] -> [pi_encoder] -> [pi_head] -> [action_dist_inputs] + [obs, action] -> [qf_encoder] -> [qf_head] -> [q-value] + [obs, action] -> [qf_target_encoder] -> [qf_target_head] + -> [q-target-value] + --- + If `twin_q=True`: + [obs, action] -> [qf_twin_encoder] -> [qf_twin_head] -> [q-twin-value] + [obs, action] -> [qf_target_twin_encoder] -> [qf_target_twin_head] + -> [q-target-twin-value] + """ + + @override(RLModule) + def setup(self): + # If a twin Q architecture should be used. + self.twin_q = self.model_config["twin_q"] + + # Build the encoder for the policy. + self.pi_encoder = self.catalog.build_encoder(framework=self.framework) + + if not self.inference_only or self.framework != "torch": + # SAC needs a separate Q network encoder (besides the pi network). + # This is because the Q network also takes the action as input + # (concatenated with the observations). + self.qf_encoder = self.catalog.build_qf_encoder(framework=self.framework) + + # If necessary, build also a twin Q encoders. + if self.twin_q: + self.qf_twin_encoder = self.catalog.build_qf_encoder( + framework=self.framework + ) + + # Build heads. + self.pi = self.catalog.build_pi_head(framework=self.framework) + + if not self.inference_only or self.framework != "torch": + self.qf = self.catalog.build_qf_head(framework=self.framework) + # If necessary build also a twin Q heads. + if self.twin_q: + self.qf_twin = self.catalog.build_qf_head(framework=self.framework) + + @override(TargetNetworkAPI) + def make_target_networks(self): + self.target_qf_encoder = make_target_network(self.qf_encoder) + self.target_qf = make_target_network(self.qf) + if self.twin_q: + self.target_qf_twin_encoder = make_target_network(self.qf_twin_encoder) + self.target_qf_twin = make_target_network(self.qf_twin) + + @override(InferenceOnlyAPI) + def get_non_inference_attributes(self) -> List[str]: + ret = ["qf", "target_qf", "qf_encoder", "target_qf_encoder"] + if self.twin_q: + ret += [ + "qf_twin", + "target_qf_twin", + "qf_twin_encoder", + "target_qf_twin_encoder", + ] + return ret + + @override(TargetNetworkAPI) + def get_target_network_pairs(self) -> List[Tuple[NetworkType, NetworkType]]: + """Returns target Q and Q network(s) to update the target network(s).""" + return [ + (self.qf_encoder, self.target_qf_encoder), + (self.qf, self.target_qf), + ] + ( + # If we have twin networks we need to update them, too. + [ + (self.qf_twin_encoder, self.target_qf_twin_encoder), + (self.qf_twin, self.target_qf_twin), + ] + if self.twin_q + else [] + ) + + # TODO (simon): SAC does not support RNNs, yet. + @override(RLModule) + def get_initial_state(self) -> dict: + # if hasattr(self.pi_encoder, "get_initial_state"): + # return { + # ACTOR: self.pi_encoder.get_initial_state(), + # CRITIC: self.qf_encoder.get_initial_state(), + # CRITIC_TARGET: self.qf_target_encoder.get_initial_state(), + # } + # else: + # return {} + return {} + + @abstractmethod + @OverrideToImplementCustomLogic + def _qf_forward_train_helper( + self, batch: Dict[str, Any], encoder: Encoder, head: Model + ) -> Dict[str, Any]: + """Executes the forward pass for Q networks. + + Args: + batch: Dict containing a concatenated tensor with observations + and actions under the key `SampleBatch.OBS`. + encoder: An `Encoder` model for the Q state-action encoder. + head: A `Model` for the Q head. + + Returns: + The estimated Q-value using the `encoder` and `head` networks. + """ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac.py new file mode 100644 index 0000000000000000000000000000000000000000..d464e95889db3621b1c538c1b9126563f29c60e1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac.py @@ -0,0 +1,591 @@ +import logging +from typing import Any, Dict, Optional, Tuple, Type, Union + +from typing_extensions import Self + +from ray._common.deprecation import DEPRECATED_VALUE, deprecation_warning +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided +from ray.rllib.algorithms.dqn.dqn import DQN +from ray.rllib.algorithms.sac.sac_tf_policy import SACTFPolicy +from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import ( + AddObservationsFromEpisodesToBatch, +) +from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import ( # noqa + AddNextObservationsFromEpisodesToTrainBatch, +) +from ray.rllib.core.learner import Learner +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.policy.policy import Policy +from ray.rllib.utils import deep_update +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_tfp +from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer +from ray.rllib.utils.typing import LearningRateOrSchedule, RLModuleSpecType + +tf1, tf, tfv = try_import_tf() +tfp = try_import_tfp() + +logger = logging.getLogger(__name__) + + +class SACConfig(AlgorithmConfig): + """Defines a configuration class from which an SAC Algorithm can be built. + + .. testcode:: + + config = ( + SACConfig() + .environment("Pendulum-v1") + .env_runners(num_env_runners=1) + .training( + gamma=0.9, + actor_lr=0.001, + critic_lr=0.002, + train_batch_size_per_learner=32, + ) + ) + # Build the SAC algo object from the config and run 1 training iteration. + algo = config.build() + algo.train() + """ + + def __init__(self, algo_class=None): + self.exploration_config = { + # The Exploration class to use. In the simplest case, this is the name + # (str) of any class present in the `rllib.utils.exploration` package. + # You can also provide the python class directly or the full location + # of your class (e.g. "ray.rllib.utils.exploration.epsilon_greedy. + # EpsilonGreedy"). + "type": "StochasticSampling", + # Add constructor kwargs here (if any). + } + + super().__init__(algo_class=algo_class or SAC) + + # fmt: off + # __sphinx_doc_begin__ + # SAC-specific config settings. + # `.training()` + self.twin_q = True + self.q_model_config = { + "fcnet_hiddens": [256, 256], + "fcnet_activation": "relu", + "post_fcnet_hiddens": [], + "post_fcnet_activation": None, + "custom_model": None, # Use this to define custom Q-model(s). + "custom_model_config": {}, + } + self.policy_model_config = { + "fcnet_hiddens": [256, 256], + "fcnet_activation": "relu", + "post_fcnet_hiddens": [], + "post_fcnet_activation": None, + "custom_model": None, # Use this to define a custom policy model. + "custom_model_config": {}, + } + self.clip_actions = False + self.tau = 5e-3 + self.initial_alpha = 1.0 + self.target_entropy = "auto" + self.n_step = 1 + + # Replay buffer configuration. + self.replay_buffer_config = { + "type": "PrioritizedEpisodeReplayBuffer", + # Size of the replay buffer. Note that if async_updates is set, + # then each worker will have a replay buffer of this size. + "capacity": int(1e6), + "alpha": 0.6, + # Beta parameter for sampling from prioritized replay buffer. + "beta": 0.4, + } + + self.store_buffer_in_checkpoints = False + self.training_intensity = None + self.optimization = { + "actor_learning_rate": 3e-4, + "critic_learning_rate": 3e-4, + "entropy_learning_rate": 3e-4, + } + self.actor_lr = 3e-5 + self.critic_lr = 3e-4 + self.alpha_lr = 3e-4 + # Set `lr` parameter to `None` and ensure it is not used. + self.lr = None + self.grad_clip = None + self.target_network_update_freq = 0 + + # .env_runners() + # Set to `self.n_step`, if 'auto'. + self.rollout_fragment_length = "auto" + + # .training() + self.train_batch_size_per_learner = 256 + self.train_batch_size = 256 # @OldAPIstack + self.num_steps_sampled_before_learning_starts = 1500 + + # .reporting() + self.min_time_s_per_iteration = 1 + self.min_sample_timesteps_per_iteration = 100 + # __sphinx_doc_end__ + # fmt: on + + self._deterministic_loss = False + self._use_beta_distribution = False + + self.use_state_preprocessor = DEPRECATED_VALUE + self.worker_side_prioritization = DEPRECATED_VALUE + + @override(AlgorithmConfig) + def training( + self, + *, + twin_q: Optional[bool] = NotProvided, + q_model_config: Optional[Dict[str, Any]] = NotProvided, + policy_model_config: Optional[Dict[str, Any]] = NotProvided, + tau: Optional[float] = NotProvided, + initial_alpha: Optional[float] = NotProvided, + target_entropy: Optional[Union[str, float]] = NotProvided, + n_step: Optional[Union[int, Tuple[int, int]]] = NotProvided, + store_buffer_in_checkpoints: Optional[bool] = NotProvided, + replay_buffer_config: Optional[Dict[str, Any]] = NotProvided, + training_intensity: Optional[float] = NotProvided, + clip_actions: Optional[bool] = NotProvided, + grad_clip: Optional[float] = NotProvided, + optimization_config: Optional[Dict[str, Any]] = NotProvided, + actor_lr: Optional[LearningRateOrSchedule] = NotProvided, + critic_lr: Optional[LearningRateOrSchedule] = NotProvided, + alpha_lr: Optional[LearningRateOrSchedule] = NotProvided, + target_network_update_freq: Optional[int] = NotProvided, + _deterministic_loss: Optional[bool] = NotProvided, + _use_beta_distribution: Optional[bool] = NotProvided, + num_steps_sampled_before_learning_starts: Optional[int] = NotProvided, + **kwargs, + ) -> Self: + """Sets the training related configuration. + + Args: + twin_q: Use two Q-networks (instead of one) for action-value estimation. + Note: Each Q-network will have its own target network. + q_model_config: Model configs for the Q network(s). These will override + MODEL_DEFAULTS. This is treated just as the top-level `model` dict in + setting up the Q-network(s) (2 if twin_q=True). + That means, you can do for different observation spaces: + `obs=Box(1D)` -> `Tuple(Box(1D) + Action)` -> `concat` -> `post_fcnet` + obs=Box(3D) -> Tuple(Box(3D) + Action) -> vision-net -> concat w/ action + -> post_fcnet + obs=Tuple(Box(1D), Box(3D)) -> Tuple(Box(1D), Box(3D), Action) + -> vision-net -> concat w/ Box(1D) and action -> post_fcnet + You can also have SAC use your custom_model as Q-model(s), by simply + specifying the `custom_model` sub-key in below dict (just like you would + do in the top-level `model` dict. + policy_model_config: Model options for the policy function (see + `q_model_config` above for details). The difference to `q_model_config` + above is that no action concat'ing is performed before the post_fcnet + stack. + tau: Update the target by \tau * policy + (1-\tau) * target_policy. + initial_alpha: Initial value to use for the entropy weight alpha. + target_entropy: Target entropy lower bound. If "auto", will be set + to `-|A|` (e.g. -2.0 for Discrete(2), -3.0 for Box(shape=(3,))). + This is the inverse of reward scale, and will be optimized + automatically. + n_step: N-step target updates. If >1, sars' tuples in trajectories will be + postprocessed to become sa[discounted sum of R][s t+n] tuples. An + integer will be interpreted as a fixed n-step value. If a tuple of 2 + ints is provided here, the n-step value will be drawn for each sample(!) + in the train batch from a uniform distribution over the closed interval + defined by `[n_step[0], n_step[1]]`. + store_buffer_in_checkpoints: Set this to True, if you want the contents of + your buffer(s) to be stored in any saved checkpoints as well. + Warnings will be created if: + - This is True AND restoring from a checkpoint that contains no buffer + data. + - This is False AND restoring from a checkpoint that does contain + buffer data. + replay_buffer_config: Replay buffer config. + Examples: + { + "_enable_replay_buffer_api": True, + "type": "MultiAgentReplayBuffer", + "capacity": 50000, + "replay_batch_size": 32, + "replay_sequence_length": 1, + } + - OR - + { + "_enable_replay_buffer_api": True, + "type": "MultiAgentPrioritizedReplayBuffer", + "capacity": 50000, + "prioritized_replay_alpha": 0.6, + "prioritized_replay_beta": 0.4, + "prioritized_replay_eps": 1e-6, + "replay_sequence_length": 1, + } + - Where - + prioritized_replay_alpha: Alpha parameter controls the degree of + prioritization in the buffer. In other words, when a buffer sample has + a higher temporal-difference error, with how much more probability + should it drawn to use to update the parametrized Q-network. 0.0 + corresponds to uniform probability. Setting much above 1.0 may quickly + result as the sampling distribution could become heavily “pointy” with + low entropy. + prioritized_replay_beta: Beta parameter controls the degree of + importance sampling which suppresses the influence of gradient updates + from samples that have higher probability of being sampled via alpha + parameter and the temporal-difference error. + prioritized_replay_eps: Epsilon parameter sets the baseline probability + for sampling so that when the temporal-difference error of a sample is + zero, there is still a chance of drawing the sample. + training_intensity: The intensity with which to update the model (vs + collecting samples from the env). + If None, uses "natural" values of: + `train_batch_size` / (`rollout_fragment_length` x `num_env_runners` x + `num_envs_per_env_runner`). + If not None, will make sure that the ratio between timesteps inserted + into and sampled from th buffer matches the given values. + Example: + training_intensity=1000.0 + train_batch_size=250 + rollout_fragment_length=1 + num_env_runners=1 (or 0) + num_envs_per_env_runner=1 + -> natural value = 250 / 1 = 250.0 + -> will make sure that replay+train op will be executed 4x asoften as + rollout+insert op (4 * 250 = 1000). + See: rllib/algorithms/dqn/dqn.py::calculate_rr_weights for further + details. + clip_actions: Whether to clip actions. If actions are already normalized, + this should be set to False. + grad_clip: If not None, clip gradients during optimization at this value. + optimization_config: Config dict for optimization. Set the supported keys + `actor_learning_rate`, `critic_learning_rate`, and + `entropy_learning_rate` in here. + actor_lr: The learning rate (float) or learning rate schedule for the + policy in the format of + [[timestep, lr-value], [timestep, lr-value], ...] In case of a + schedule, intermediary timesteps will be assigned to linearly + interpolated learning rate values. A schedule config's first entry + must start with timestep 0, i.e.: [[0, initial_value], [...]]. + Note: It is common practice (two-timescale approach) to use a smaller + learning rate for the policy than for the critic to ensure that the + critic gives adequate values for improving the policy. + Note: If you require a) more than one optimizer (per RLModule), + b) optimizer types that are not Adam, c) a learning rate schedule that + is not a linearly interpolated, piecewise schedule as described above, + or d) specifying c'tor arguments of the optimizer that are not the + learning rate (e.g. Adam's epsilon), then you must override your + Learner's `configure_optimizer_for_module()` method and handle + lr-scheduling yourself. + The default value is 3e-5, one decimal less than the respective + learning rate of the critic (see `critic_lr`). + critic_lr: The learning rate (float) or learning rate schedule for the + critic in the format of + [[timestep, lr-value], [timestep, lr-value], ...] In case of a + schedule, intermediary timesteps will be assigned to linearly + interpolated learning rate values. A schedule config's first entry + must start with timestep 0, i.e.: [[0, initial_value], [...]]. + Note: It is common practice (two-timescale approach) to use a smaller + learning rate for the policy than for the critic to ensure that the + critic gives adequate values for improving the policy. + Note: If you require a) more than one optimizer (per RLModule), + b) optimizer types that are not Adam, c) a learning rate schedule that + is not a linearly interpolated, piecewise schedule as described above, + or d) specifying c'tor arguments of the optimizer that are not the + learning rate (e.g. Adam's epsilon), then you must override your + Learner's `configure_optimizer_for_module()` method and handle + lr-scheduling yourself. + The default value is 3e-4, one decimal higher than the respective + learning rate of the actor (policy) (see `actor_lr`). + alpha_lr: The learning rate (float) or learning rate schedule for the + hyperparameter alpha in the format of + [[timestep, lr-value], [timestep, lr-value], ...] In case of a + schedule, intermediary timesteps will be assigned to linearly + interpolated learning rate values. A schedule config's first entry + must start with timestep 0, i.e.: [[0, initial_value], [...]]. + Note: If you require a) more than one optimizer (per RLModule), + b) optimizer types that are not Adam, c) a learning rate schedule that + is not a linearly interpolated, piecewise schedule as described above, + or d) specifying c'tor arguments of the optimizer that are not the + learning rate (e.g. Adam's epsilon), then you must override your + Learner's `configure_optimizer_for_module()` method and handle + lr-scheduling yourself. + The default value is 3e-4, identical to the critic learning rate (`lr`). + target_network_update_freq: Update the target network every + `target_network_update_freq` steps. + num_steps_sampled_before_learning_starts: Number of timesteps (int) + that we collect from the runners before we start sampling the + replay buffers for learning. Whether we count this in agent steps + or environment steps depends on the value of + `config.multi_agent(count_steps_by=...)`. + _deterministic_loss: Whether the loss should be calculated deterministically + (w/o the stochastic action sampling step). True only useful for + continuous actions and for debugging. + _use_beta_distribution: Use a Beta-distribution instead of a + `SquashedGaussian` for bounded, continuous action spaces (not + recommended; for debugging only). + + Returns: + This updated AlgorithmConfig object. + """ + # Pass kwargs onto super's `training()` method. + super().training(**kwargs) + + if twin_q is not NotProvided: + self.twin_q = twin_q + if q_model_config is not NotProvided: + self.q_model_config.update(q_model_config) + if policy_model_config is not NotProvided: + self.policy_model_config.update(policy_model_config) + if tau is not NotProvided: + self.tau = tau + if initial_alpha is not NotProvided: + self.initial_alpha = initial_alpha + if target_entropy is not NotProvided: + self.target_entropy = target_entropy + if n_step is not NotProvided: + self.n_step = n_step + if store_buffer_in_checkpoints is not NotProvided: + self.store_buffer_in_checkpoints = store_buffer_in_checkpoints + 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"] + if training_intensity is not NotProvided: + self.training_intensity = training_intensity + if clip_actions is not NotProvided: + self.clip_actions = clip_actions + if grad_clip is not NotProvided: + self.grad_clip = grad_clip + if optimization_config is not NotProvided: + self.optimization = optimization_config + if actor_lr is not NotProvided: + self.actor_lr = actor_lr + if critic_lr is not NotProvided: + self.critic_lr = critic_lr + if alpha_lr is not NotProvided: + self.alpha_lr = alpha_lr + if target_network_update_freq is not NotProvided: + self.target_network_update_freq = target_network_update_freq + if _deterministic_loss is not NotProvided: + self._deterministic_loss = _deterministic_loss + if _use_beta_distribution is not NotProvided: + self._use_beta_distribution = _use_beta_distribution + if num_steps_sampled_before_learning_starts is not NotProvided: + self.num_steps_sampled_before_learning_starts = ( + num_steps_sampled_before_learning_starts + ) + + return self + + @override(AlgorithmConfig) + def validate(self) -> None: + # Call super's validation method. + super().validate() + + # Check rollout_fragment_length to be compatible with n_step. + if isinstance(self.n_step, tuple): + min_rollout_fragment_length = self.n_step[1] + else: + min_rollout_fragment_length = self.n_step + + if ( + not self.in_evaluation + and self.rollout_fragment_length != "auto" + and self.rollout_fragment_length + < min_rollout_fragment_length # (self.n_step or 1) + ): + raise ValueError( + f"Your `rollout_fragment_length` ({self.rollout_fragment_length}) is " + f"smaller than needed for `n_step` ({self.n_step})! If `n_step` is " + f"an integer try setting `rollout_fragment_length={self.n_step}`. If " + "`n_step` is a tuple, try setting " + f"`rollout_fragment_length={self.n_step[1]}`." + ) + + if self.use_state_preprocessor != DEPRECATED_VALUE: + deprecation_warning( + old="config['use_state_preprocessor']", + error=False, + ) + self.use_state_preprocessor = DEPRECATED_VALUE + + if self.grad_clip is not None and self.grad_clip <= 0.0: + raise ValueError("`grad_clip` value must be > 0.0!") + + if self.framework in ["tf", "tf2"] and tfp is None: + logger.warning( + "You need `tensorflow_probability` in order to run SAC! " + "Install it via `pip install tensorflow_probability`. Your " + f"tf.__version__={tf.__version__ if tf else None}." + "Trying to import tfp results in the following error:" + ) + try_import_tfp(error=True) + + # Validate that we use the corresponding `EpisodeReplayBuffer` when using + # episodes. + if ( + self.enable_env_runner_and_connector_v2 + and self.replay_buffer_config["type"] + not in [ + "EpisodeReplayBuffer", + "PrioritizedEpisodeReplayBuffer", + "MultiAgentEpisodeReplayBuffer", + "MultiAgentPrioritizedEpisodeReplayBuffer", + ] + and not ( + # TODO (simon): Set up an indicator `is_offline_new_stack` that + # includes all these variable checks. + self.input_ + and ( + isinstance(self.input_, str) + or ( + isinstance(self.input_, list) + and isinstance(self.input_[0], str) + ) + ) + and self.input_ != "sampler" + and self.enable_rl_module_and_learner + ) + ): + raise ValueError( + "When using the new `EnvRunner API` the replay buffer must be of type " + "`EpisodeReplayBuffer`." + ) + elif not self.enable_env_runner_and_connector_v2 and ( + ( + isinstance(self.replay_buffer_config["type"], str) + and "Episode" in self.replay_buffer_config["type"] + ) + or ( + isinstance(self.replay_buffer_config["type"], type) + and issubclass(self.replay_buffer_config["type"], EpisodeReplayBuffer) + ) + ): + raise ValueError( + "When using the old API stack the replay buffer must not be of type " + "`EpisodeReplayBuffer`! We suggest you use the following config to run " + "SAC on the old API stack: `config.training(replay_buffer_config={" + "'type': 'MultiAgentPrioritizedReplayBuffer', " + "'prioritized_replay_alpha': [alpha], " + "'prioritized_replay_beta': [beta], " + "'prioritized_replay_eps': [eps], " + "})`." + ) + + if self.enable_rl_module_and_learner: + if self.lr is not None: + raise ValueError( + "Basic learning rate parameter `lr` is not `None`. For SAC " + "use the specific learning rate parameters `actor_lr`, `critic_lr` " + "and `alpha_lr`, for the actor, critic, and the hyperparameter " + "`alpha`, respectively and set `config.lr` to None." + ) + # Warn about new API stack on by default. + logger.warning( + "You are running SAC on the new API stack! This is the new default " + "behavior for this algorithm. If you don't want to use the new API " + "stack, set `config.api_stack(enable_rl_module_and_learner=False, " + "enable_env_runner_and_connector_v2=False)`. For a detailed " + "migration guide, see here: https://docs.ray.io/en/master/rllib/new-api-stack-migration-guide.html" # noqa + ) + + @override(AlgorithmConfig) + def get_rollout_fragment_length(self, worker_index: int = 0) -> int: + if self.rollout_fragment_length == "auto": + return ( + self.n_step[1] + if isinstance(self.n_step, (tuple, list)) + else self.n_step + ) + else: + return self.rollout_fragment_length + + @override(AlgorithmConfig) + def get_default_rl_module_spec(self) -> RLModuleSpecType: + if self.framework_str == "torch": + from ray.rllib.algorithms.sac.torch.default_sac_torch_rl_module import ( + DefaultSACTorchRLModule, + ) + + return RLModuleSpec(module_class=DefaultSACTorchRLModule) + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. Use `torch`." + ) + + @override(AlgorithmConfig) + def get_default_learner_class(self) -> Union[Type["Learner"], str]: + if self.framework_str == "torch": + from ray.rllib.algorithms.sac.torch.sac_torch_learner import SACTorchLearner + + return SACTorchLearner + else: + raise ValueError( + f"The framework {self.framework_str} is not supported. Use `torch`." + ) + + @override(AlgorithmConfig) + def build_learner_connector( + self, + input_observation_space, + input_action_space, + device=None, + ): + pipeline = super().build_learner_connector( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + device=device, + ) + + # Prepend the "add-NEXT_OBS-from-episodes-to-train-batch" connector piece (right + # after the corresponding "add-OBS-..." default piece). + pipeline.insert_after( + AddObservationsFromEpisodesToBatch, + AddNextObservationsFromEpisodesToTrainBatch(), + ) + + return pipeline + + @property + def _model_config_auto_includes(self): + return super()._model_config_auto_includes | {"twin_q": self.twin_q} + + +class SAC(DQN): + """Soft Actor Critic (SAC) Algorithm class. + + This file defines the distributed Algorithm class for the soft actor critic + algorithm. + See `sac_[tf|torch]_policy.py` for the definition of the policy loss. + + Detailed documentation: + https://docs.ray.io/en/master/rllib-algorithms.html#sac + """ + + def __init__(self, *args, **kwargs): + self._allow_unknown_subkeys += ["policy_model_config", "q_model_config"] + super().__init__(*args, **kwargs) + + @classmethod + @override(DQN) + def get_default_config(cls) -> SACConfig: + return SACConfig() + + @classmethod + @override(DQN) + def get_default_policy_class( + cls, config: AlgorithmConfig + ) -> Optional[Type[Policy]]: + if config["framework"] == "torch": + from ray.rllib.algorithms.sac.sac_torch_policy import SACTorchPolicy + + return SACTorchPolicy + else: + return SACTFPolicy diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_catalog.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..2ebe4470af18740c5838b23199dfbaec159c2628 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_catalog.py @@ -0,0 +1,288 @@ +from typing import Callable + +import gymnasium as gym +import numpy as np + +# TODO (simon): Store this function somewhere more central as many +# algorithms will use it. +from ray.rllib.algorithms.ppo.ppo_catalog import _check_if_diag_gaussian +from ray.rllib.core.distribution.distribution import Distribution +from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchCategorical, + TorchSquashedGaussian, +) +from ray.rllib.core.models.base import Encoder, Model +from ray.rllib.core.models.catalog import Catalog +from ray.rllib.core.models.configs import ( + FreeLogStdMLPHeadConfig, + MLPEncoderConfig, + MLPHeadConfig, +) +from ray.rllib.utils.annotations import OverrideToImplementCustomLogic, override + + +# TODO (simon): Check, if we can directly derive from DQNCatalog. +# This should work as we need a qf and qf_target. +# TODO (simon): Add CNNEnocders for Image observations. +class SACCatalog(Catalog): + """The catalog class used to build models for SAC. + + SACCatalog provides the following models: + - Encoder: The encoder used to encode the observations for the actor + network (`pi`). For this we use the default encoder from the Catalog. + - Q-Function Encoder: The encoder used to encode the observations and + actions for the soft Q-function network. + - Target Q-Function Encoder: The encoder used to encode the observations + and actions for the target soft Q-function network. + - Pi Head: The head used to compute the policy logits. This network outputs + the mean and log-std for the action distribution (a Squashed Gaussian). + - Q-Function Head: The head used to compute the soft Q-values. + - Target Q-Function Head: The head used to compute the target soft Q-values. + + Any custom Encoder to be used for the policy network can be built by overriding + the build_encoder() method. Alternatively the `encoder_config` can be overridden + by using the `model_config_dict`. + + Any custom Q-Function Encoder can be built by overriding the build_qf_encoder(). + Important: The Q-Function Encoder must encode both the state and the action. The + same holds true for the target Q-Function Encoder. + + Any custom head can be built by overriding the build_pi_head() and build_qf_head(). + + Any module built for exploration or inference is built with the flag + `ìnference_only=True` and does not contain any Q-function. This flag can be set + in the `model_config_dict` with the key `ray.rllib.core.rl_module.INFERENCE_ONLY`. + """ + + def __init__( + self, + observation_space: gym.Space, + action_space: gym.Space, + model_config_dict: dict, + view_requirements: dict = None, + ): + """Initializes the SACCatalog. + + Args: + observation_space: The observation space of the Encoder. + action_space: The action space for the Pi Head. + model_config_dict: The model config to use. + """ + assert view_requirements is None, ( + "Instead, use the new ConnectorV2 API to pick whatever information " + "you need from the running episodes" + ) + + super().__init__( + observation_space=observation_space, + action_space=action_space, + model_config_dict=model_config_dict, + ) + if not isinstance(self.action_space, (gym.spaces.Box, gym.spaces.Discrete)): + self._raise_unsupported_action_space_error() + + # Define the heads. + self.pi_and_qf_head_hiddens = self._model_config_dict["head_fcnet_hiddens"] + self.pi_and_qf_head_activation = self._model_config_dict[ + "head_fcnet_activation" + ] + + # We don't have the exact (framework specific) action dist class yet and thus + # cannot determine the exact number of output nodes (action space) required. + # -> Build pi config only in the `self.build_pi_head` method. + self.pi_head_config = None + + # SAC-Discrete: The Q-function outputs q-values for each action + # SAC-Continuous: The Q-function outputs a single value (the Q-value for the + # action taken). + required_qf_output_dim = ( + self.action_space.n + if isinstance(self.action_space, gym.spaces.Discrete) + else 1 + ) + + # TODO (simon): Implement in a later step a q network with + # different `head_fcnet_hiddens` than pi. + # TODO (simon): These latent_dims could be different for the + # q function, value function, and pi head. + # Here we consider the simple case of identical encoders. + self.qf_head_config = MLPHeadConfig( + input_dims=self.latent_dims, + hidden_layer_dims=self.pi_and_qf_head_hiddens, + hidden_layer_activation=self.pi_and_qf_head_activation, + output_layer_activation="linear", + output_layer_dim=required_qf_output_dim, + ) + + @OverrideToImplementCustomLogic + def build_qf_encoder(self, framework: str) -> Encoder: + """Builds the Q-function encoder. + + In contrast to PPO, SAC needs a different encoder for Pi and + Q-function as the Q-function in the continuous case has to + encode actions, too. Therefore the Q-function uses its own + encoder config. + Note, the Pi network uses the base encoder from the `Catalog`. + + Args: + framework: The framework to use. Either `torch` or `tf2`. + + Returns: + The encoder for the Q-network. + """ + + # Compute the required dimension for the action space. + if isinstance(self.action_space, gym.spaces.Box): + required_action_dim = self.action_space.shape[0] + elif isinstance(self.action_space, gym.spaces.Discrete): + # for discrete action spaces, we don't need to encode the action + # because the Q-function will output a value for each action + required_action_dim = 0 + else: + self._raise_unsupported_action_space_error() + + # Encoder input for the Q-network contains state and action. We + # need to infer the shape for the input from the state and action + # spaces + if not ( + isinstance(self.observation_space, gym.spaces.Box) + and len(self.observation_space.shape) == 1 + ): + raise ValueError("The observation space is not supported by RLlib's SAC.") + + input_space = gym.spaces.Box( + -np.inf, + np.inf, + (self.observation_space.shape[0] + required_action_dim,), + dtype=np.float32, + ) + + self.qf_encoder_hiddens = self._model_config_dict["fcnet_hiddens"][:-1] + self.qf_encoder_activation = self._model_config_dict["fcnet_activation"] + + self.qf_encoder_config = MLPEncoderConfig( + input_dims=input_space.shape, + hidden_layer_dims=self.qf_encoder_hiddens, + hidden_layer_activation=self.qf_encoder_activation, + output_layer_dim=self.latent_dims[0], + output_layer_activation=self.qf_encoder_activation, + ) + + return self.qf_encoder_config.build(framework=framework) + + @OverrideToImplementCustomLogic + def build_pi_head(self, framework: str) -> Model: + """Builds the policy head. + + The default behavior is to build the head from the pi_head_config. + This can be overridden to build a custom policy head as a means of configuring + the behavior of the DefaultSACRLModule implementation. + + Args: + framework: The framework to use. Either "torch" or "tf2". + + Returns: + The policy head. + """ + # Get action_distribution_cls to find out about the output dimension for pi_head + action_distribution_cls = self.get_action_dist_cls(framework=framework) + BUILD_MAP: dict[ + type[gym.spaces.Space], Callable[[str, Distribution], Model] + ] = { + gym.spaces.Discrete: self._build_pi_head_discrete, + gym.spaces.Box: self._build_pi_head_continuous, + } + try: + # Try to get the build function for the action space type. + return BUILD_MAP[type(self.action_space)]( + framework, action_distribution_cls + ) + except KeyError: + # If the action space type is not supported, raise an error. + self._raise_unsupported_action_space_error() + + def _build_pi_head_continuous( + self, framework: str, action_distribution_cls: Distribution + ) -> Model: + """Builds the policy head for continuous action spaces.""" + # Get action_distribution_cls to find out about the output dimension for pi_head + # TODO (simon): CHeck, if this holds also for Squashed Gaussian. + if self._model_config_dict["free_log_std"]: + _check_if_diag_gaussian( + action_distribution_cls=action_distribution_cls, framework=framework + ) + is_diag_gaussian = True + else: + is_diag_gaussian = _check_if_diag_gaussian( + action_distribution_cls=action_distribution_cls, + framework=framework, + no_error=True, + ) + required_output_dim = action_distribution_cls.required_input_dim( + space=self.action_space, model_config=self._model_config_dict + ) + # Now that we have the action dist class and number of outputs, we can define + # our pi-config and build the pi head. + pi_head_config_class = ( + FreeLogStdMLPHeadConfig + if self._model_config_dict["free_log_std"] + else MLPHeadConfig + ) + self.pi_head_config = pi_head_config_class( + input_dims=self.latent_dims, + hidden_layer_dims=self.pi_and_qf_head_hiddens, + hidden_layer_activation=self.pi_and_qf_head_activation, + output_layer_dim=required_output_dim, + output_layer_activation="linear", + clip_log_std=is_diag_gaussian, + log_std_clip_param=self._model_config_dict.get("log_std_clip_param", 20), + ) + + return self.pi_head_config.build(framework=framework) + + def _build_pi_head_discrete( + self, framework: str, action_distribution_cls: Distribution + ) -> Model: + """Builds the policy head for discrete action spaces. The module outputs logits for Categorical + distribution. + """ + required_output_dim = action_distribution_cls.required_input_dim( + space=self.action_space, model_config=self._model_config_dict + ) + self.pi_head_config = MLPHeadConfig( + input_dims=self.latent_dims, + hidden_layer_dims=self.pi_and_qf_head_hiddens, + hidden_layer_activation=self.pi_and_qf_head_activation, + output_layer_dim=required_output_dim, + output_layer_activation="linear", + ) + return self.pi_head_config.build(framework=framework) + + @OverrideToImplementCustomLogic + def build_qf_head(self, framework: str) -> Model: + """Build the Q function head.""" + + return self.qf_head_config.build(framework=framework) + + @override(Catalog) + def get_action_dist_cls(self, framework: str) -> Distribution: + """Returns the action distribution class to use for the given framework. TorchSquashedGaussian + for continuous action spaces and TorchCategorical for discrete action spaces.""" + # TODO (KIY): Catalog.get_action_dist_cls should return a type[Distribution] instead of a Distribution instance. + assert framework == "torch" + + if isinstance(self.action_space, gym.spaces.Box): + # For continuous action spaces, we use a Squashed Gaussian. + return TorchSquashedGaussian + elif isinstance(self.action_space, gym.spaces.Discrete): + # For discrete action spaces, we use a Categorical distribution. + return TorchCategorical + else: + self._raise_unsupported_action_space_error() + + def _raise_unsupported_action_space_error(self): + """Raises an error if the action space is not supported.""" + raise ValueError( + f"SAC only supports Box and Discrete action spaces. " + f"Got: {type(self.action_space)}" + ) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..abbf082b1ca179648407e7caa4fb0306ce5eba4e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_learner.py @@ -0,0 +1,115 @@ +from typing import Dict + +import numpy as np + +from ray.rllib.algorithms.dqn.dqn_learner import DQNLearner +from ray.rllib.core.learner.learner import Learner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict +from ray.rllib.utils.typing import ModuleID, TensorType + +# Now, this is double defined: In `DefaultSACRLModule` and here. I would keep it here +# or push it into the `Learner` as these are recurring keys in RL. +LOGPS_KEY = "logps" +QF_LOSS_KEY = "qf_loss" +QF_MEAN_KEY = "qf_mean" +QF_MAX_KEY = "qf_max" +QF_MIN_KEY = "qf_min" +QF_PREDS = "qf_preds" +QF_TWIN_LOSS_KEY = "qf_twin_loss" +QF_TWIN_PREDS = "qf_twin_preds" +TD_ERROR_MEAN_KEY = "td_error_mean" +CRITIC_TARGET = "critic_target" +ACTION_DIST_INPUTS_NEXT = "action_dist_inputs_next" +QF_TARGET_NEXT = "q_target_next" +ACTION_PROBS_NEXT = "action_probs_next" +ACTION_LOG_PROBS_NEXT = "action_log_probs_next" +ACTION_PROBS = "action_probs" +ACTION_LOG_PROBS = "action_log_probs" + + +class SACLearner(DQNLearner): + @override(Learner) + def build(self) -> None: + # Store the current alpha in log form. We need it during optimization + # in log form. + self.curr_log_alpha: Dict[ModuleID, TensorType] = LambdaDefaultDict( + lambda module_id: self._get_tensor_variable( + # Note, we want to train the temperature parameter. + [ + np.log( + self.config.get_config_for_module(module_id).initial_alpha + ).astype(np.float32) + ], + trainable=True, + ) + ) + + # We need to call the `super()`'s `build()` method here to have the variables + # for the alpha already defined. + super().build() + + self.target_entropy: Dict[ModuleID, TensorType] = LambdaDefaultDict( + lambda module_id: self._get_tensor_variable( + self._get_target_entropy(module_id) + ) + ) + + @override(Learner) + def remove_module(self, module_id: ModuleID) -> None: + """Removes the temperature and target entropy. + + Note, this means that we also need to remove the corresponding + temperature optimizer. + """ + super().remove_module(module_id) + self.curr_log_alpha.pop(module_id, None) + self.target_entropy.pop(module_id, None) + + @override(Learner) + def add_module( + self, + *, + module_id, + module_spec, + config_overrides=None, + new_should_module_be_updated=None + ): + # First call `super`'s `add_module` method. + super().add_module( + module_id=module_id, + module_spec=module_spec, + config_overrides=config_overrides, + new_should_module_be_updated=new_should_module_be_updated, + ) + # Now add the log alpha. + self.curr_log_alpha[module_id] = self._get_tensor_variable( + # Note, we want to train the temperature parameter. + [ + np.log( + self.config.get_config_for_module(module_id).initial_alpha + ).astype(np.float32) + ], + trainable=True, + ) + # Add also the target entropy for the new module. + self.target_entropy[module_id] = self._get_tensor_variable( + self._get_target_entropy(module_id) + ) + + def _get_target_entropy(self, module_id): + """Returns the target entropy to use for the loss. + + Args: + module_id: Module ID for which the target entropy should be + returned. + + Returns: + Target entropy. + """ + target_entropy = self.config.get_config_for_module(module_id).target_entropy + if target_entropy is None or target_entropy == "auto": + target_entropy = -np.prod( + self._module_spec.module_specs[module_id].action_space.shape + ) + return target_entropy diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_tf_model.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_tf_model.py new file mode 100644 index 0000000000000000000000000000000000000000..e3b3479ff684ae71bb2a9b383bc5f9f586f16ef9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_tf_model.py @@ -0,0 +1,322 @@ +from typing import Dict, List, Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.typing import ModelConfigDict, TensorStructType, TensorType + +tf1, tf, tfv = try_import_tf() + + +class SACTFModel(TFModelV2): + """Extension of the standard TFModelV2 for SAC. + + To customize, do one of the following: + - sub-class SACTFModel and override one or more of its methods. + - Use SAC's `q_model_config` and `policy_model` keys to tweak the default model + behaviors (e.g. fcnet_hiddens, conv_filters, etc..). + - Use SAC's `q_model_config->custom_model` and `policy_model->custom_model` keys + to specify your own custom Q-model(s) and policy-models, which will be + created within this SACTFModel (see `build_policy_model` and + `build_q_model`. + + Note: It is not recommended to override the `forward` method for SAC. This + would lead to shared weights (between policy and Q-nets), which will then + not be optimized by either of the critic- or actor-optimizers! + + Data flow: + `obs` -> forward() (should stay a noop method!) -> `model_out` + `model_out` -> get_policy_output() -> pi(actions|obs) + `model_out`, `actions` -> get_q_values() -> Q(s, a) + `model_out`, `actions` -> get_twin_q_values() -> Q_twin(s, a) + """ + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: Optional[int], + model_config: ModelConfigDict, + name: str, + policy_model_config: ModelConfigDict = None, + q_model_config: ModelConfigDict = None, + twin_q: bool = False, + initial_alpha: float = 1.0, + target_entropy: Optional[float] = None, + ): + """Initialize a SACTFModel instance. + + Args: + policy_model_config: The config dict for the + policy network. + q_model_config: The config dict for the + Q-network(s) (2 if twin_q=True). + twin_q: Build twin Q networks (Q-net and target) for more + stable Q-learning. + initial_alpha: The initial value for the to-be-optimized + alpha parameter (default: 1.0). + target_entropy (Optional[float]): A target entropy value for + the to-be-optimized alpha parameter. If None, will use the + defaults described in the papers for SAC (and discrete SAC). + + Note that the core layers for forward() are not defined here, this + only defines the layers for the output heads. Those layers for + forward() should be defined in subclasses of SACModel. + """ + super(SACTFModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + if isinstance(action_space, Discrete): + self.action_dim = action_space.n + self.discrete = True + action_outs = q_outs = self.action_dim + elif isinstance(action_space, Box): + self.action_dim = np.prod(action_space.shape) + self.discrete = False + action_outs = 2 * self.action_dim + q_outs = 1 + else: + assert isinstance(action_space, Simplex) + self.action_dim = np.prod(action_space.shape) + self.discrete = False + action_outs = self.action_dim + q_outs = 1 + + self.action_model = self.build_policy_model( + self.obs_space, action_outs, policy_model_config, "policy_model" + ) + + self.q_net = self.build_q_model( + self.obs_space, self.action_space, q_outs, q_model_config, "q" + ) + if twin_q: + self.twin_q_net = self.build_q_model( + self.obs_space, self.action_space, q_outs, q_model_config, "twin_q" + ) + else: + self.twin_q_net = None + + self.log_alpha = tf.Variable( + np.log(initial_alpha), dtype=tf.float32, name="log_alpha" + ) + self.alpha = tf.exp(self.log_alpha) + + # Auto-calculate the target entropy. + if target_entropy is None or target_entropy == "auto": + # See hyperparams in [2] (README.md). + if self.discrete: + target_entropy = 0.98 * np.array( + -np.log(1.0 / action_space.n), dtype=np.float32 + ) + # See [1] (README.md). + else: + target_entropy = -np.prod(action_space.shape) + self.target_entropy = target_entropy + + @override(TFModelV2) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + """The common (Q-net and policy-net) forward pass. + + NOTE: It is not(!) recommended to override this method as it would + introduce a shared pre-network, which would be updated by both + actor- and critic optimizers. + """ + return input_dict["obs"], state + + def build_policy_model(self, obs_space, num_outputs, policy_model_config, name): + """Builds the policy model used by this SAC. + + Override this method in a sub-class of SACTFModel to implement your + own policy net. Alternatively, simply set `custom_model` within the + top level SAC `policy_model` config key to make this default + implementation of `build_policy_model` use your custom policy network. + + Returns: + TFModelV2: The TFModelV2 policy sub-model. + """ + model = ModelCatalog.get_model_v2( + obs_space, + self.action_space, + num_outputs, + policy_model_config, + framework="tf", + name=name, + ) + return model + + def build_q_model(self, obs_space, action_space, num_outputs, q_model_config, name): + """Builds one of the (twin) Q-nets used by this SAC. + + Override this method in a sub-class of SACTFModel to implement your + own Q-nets. Alternatively, simply set `custom_model` within the + top level SAC `q_model_config` config key to make this default implementation + of `build_q_model` use your custom Q-nets. + + Returns: + TFModelV2: The TFModelV2 Q-net sub-model. + """ + self.concat_obs_and_actions = False + if self.discrete: + input_space = obs_space + else: + orig_space = getattr(obs_space, "original_space", obs_space) + if isinstance(orig_space, Box) and len(orig_space.shape) == 1: + input_space = Box( + float("-inf"), + float("inf"), + shape=(orig_space.shape[0] + action_space.shape[0],), + ) + self.concat_obs_and_actions = True + else: + input_space = gym.spaces.Tuple([orig_space, action_space]) + + model = ModelCatalog.get_model_v2( + input_space, + action_space, + num_outputs, + q_model_config, + framework="tf", + name=name, + ) + return model + + def get_q_values( + self, model_out: TensorType, actions: Optional[TensorType] = None + ) -> TensorType: + """Returns Q-values, given the output of self.__call__(). + + This implements Q(s, a) -> [single Q-value] for the continuous case and + Q(s) -> [Q-values for all actions] for the discrete case. + + Args: + model_out: Feature outputs from the model layers + (result of doing `self.__call__(obs)`). + actions (Optional[TensorType]): Continuous action batch to return + Q-values for. Shape: [BATCH_SIZE, action_dim]. If None + (discrete action case), return Q-values for all actions. + + Returns: + TensorType: Q-values tensor of shape [BATCH_SIZE, 1]. + """ + return self._get_q_value(model_out, actions, self.q_net) + + def get_twin_q_values( + self, model_out: TensorType, actions: Optional[TensorType] = None + ) -> TensorType: + """Same as get_q_values but using the twin Q net. + + This implements the twin Q(s, a). + + Args: + model_out: Feature outputs from the model layers + (result of doing `self.__call__(obs)`). + actions (Optional[Tensor]): Actions to return the Q-values for. + Shape: [BATCH_SIZE, action_dim]. If None (discrete action + case), return Q-values for all actions. + + Returns: + TensorType: Q-values tensor of shape [BATCH_SIZE, 1]. + """ + return self._get_q_value(model_out, actions, self.twin_q_net) + + def _get_q_value(self, model_out, actions, net): + # Model outs may come as original Tuple/Dict observations, concat them + # here if this is the case. + if isinstance(net.obs_space, Box): + if isinstance(model_out, (list, tuple)): + model_out = tf.concat(model_out, axis=-1) + elif isinstance(model_out, dict): + model_out = tf.concat(list(model_out.values()), axis=-1) + + # Continuous case -> concat actions to model_out. + if actions is not None: + if self.concat_obs_and_actions: + input_dict = {"obs": tf.concat([model_out, actions], axis=-1)} + else: + # TODO(junogng) : SampleBatch doesn't support list columns yet. + # Use ModelInputDict. + input_dict = {"obs": (model_out, actions)} + # Discrete case -> return q-vals for all actions. + else: + input_dict = {"obs": model_out} + # Switch on training mode (when getting Q-values, we are usually in + # training). + input_dict["is_training"] = True + + return net(input_dict, [], None) + + def get_action_model_outputs( + self, + model_out: TensorType, + state_in: List[TensorType] = None, + seq_lens: TensorType = None, + ) -> (TensorType, List[TensorType]): + """Returns distribution inputs and states given the output of + policy.model(). + + For continuous action spaces, these will be the mean/stddev + distribution inputs for the (SquashedGaussian) action distribution. + For discrete action spaces, these will be the logits for a categorical + distribution. + + Args: + model_out: Feature outputs from the model layers + (result of doing `model(obs)`). + state_in List(TensorType): State input for recurrent cells + seq_lens: Sequence lengths of input- and state + sequences + + Returns: + TensorType: Distribution inputs for sampling actions. + """ + + def concat_obs_if_necessary(obs: TensorStructType): + """Concat model outs if they are original tuple observations.""" + if isinstance(obs, (list, tuple)): + obs = tf.concat(obs, axis=-1) + elif isinstance(obs, dict): + obs = tf.concat( + [ + tf.expand_dims(val, 1) if len(val.shape) == 1 else val + for val in tree.flatten(obs.values()) + ], + axis=-1, + ) + return obs + + if state_in is None: + state_in = [] + + if isinstance(model_out, dict) and "obs" in model_out: + # Model outs may come as original Tuple observations + if isinstance(self.action_model.obs_space, Box): + model_out["obs"] = concat_obs_if_necessary(model_out["obs"]) + return self.action_model(model_out, state_in, seq_lens) + else: + if isinstance(self.action_model.obs_space, Box): + model_out = concat_obs_if_necessary(model_out) + return self.action_model({"obs": model_out}, state_in, seq_lens) + + def policy_variables(self): + """Return the list of variables for the policy net.""" + + return self.action_model.variables() + + def q_variables(self): + """Return the list of variables for Q / twin Q nets.""" + + return self.q_net.variables() + ( + self.twin_q_net.variables() if self.twin_q_net else [] + ) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_tf_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_tf_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..e4322518e46ad21a6bc5381af823787c00ae6555 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_tf_policy.py @@ -0,0 +1,793 @@ +""" +TensorFlow policy class used for SAC. +""" + +import copy +import logging +from functools import partial +from typing import Dict, List, Optional, Tuple, Type, Union + +import gymnasium as gym +from gymnasium.spaces import Box, Discrete + +import ray +from ray.rllib.algorithms.dqn.dqn_tf_policy import ( + PRIO_WEIGHTS, + postprocess_nstep_and_prio, +) +from ray.rllib.algorithms.sac.sac_tf_model import SACTFModel +from ray.rllib.algorithms.sac.sac_torch_model import SACTorchModel +from ray.rllib.models import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import ( + Beta, + Categorical, + DiagGaussian, + Dirichlet, + SquashedGaussian, + TFActionDistribution, +) +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_mixins import TargetNetworkMixin +from ray.rllib.policy.tf_policy_template import build_tf_policy +from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.utils.framework import get_variable, try_import_tf +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.tf_utils import huber_loss, make_tf_callable +from ray.rllib.utils.typing import ( + AgentID, + AlgorithmConfigDict, + LocalOptimizer, + ModelGradients, + TensorType, +) + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +def build_sac_model( + policy: Policy, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, +) -> ModelV2: + """Constructs the necessary ModelV2 for the Policy and returns it. + + Args: + policy: The TFPolicy that will use the models. + obs_space (gym.spaces.Space): The observation space. + action_space (gym.spaces.Space): The action space. + config: The SACConfig object. + + Returns: + ModelV2: The ModelV2 to be used by the Policy. Note: An additional + target model will be created in this function and assigned to + `policy.target_model`. + """ + # Force-ignore any additionally provided hidden layer sizes. + # Everything should be configured using SAC's `q_model_config` and + # `policy_model_config` config settings. + policy_model_config = copy.deepcopy(config["model"]) + policy_model_config.update(config["policy_model_config"]) + q_model_config = copy.deepcopy(config["model"]) + q_model_config.update(config["q_model_config"]) + + default_model_cls = SACTorchModel if config["framework"] == "torch" else SACTFModel + + model = ModelCatalog.get_model_v2( + obs_space=obs_space, + action_space=action_space, + num_outputs=None, + model_config=config["model"], + framework=config["framework"], + default_model=default_model_cls, + name="sac_model", + policy_model_config=policy_model_config, + q_model_config=q_model_config, + twin_q=config["twin_q"], + initial_alpha=config["initial_alpha"], + target_entropy=config["target_entropy"], + ) + + assert isinstance(model, default_model_cls) + + # Create an exact copy of the model and store it in `policy.target_model`. + # This will be used for tau-synched Q-target models that run behind the + # actual Q-networks and are used for target q-value calculations in the + # loss terms. + policy.target_model = ModelCatalog.get_model_v2( + obs_space=obs_space, + action_space=action_space, + num_outputs=None, + model_config=config["model"], + framework=config["framework"], + default_model=default_model_cls, + name="target_sac_model", + policy_model_config=policy_model_config, + q_model_config=q_model_config, + twin_q=config["twin_q"], + initial_alpha=config["initial_alpha"], + target_entropy=config["target_entropy"], + ) + + assert isinstance(policy.target_model, default_model_cls) + + return model + + +def postprocess_trajectory( + policy: Policy, + sample_batch: SampleBatch, + other_agent_batches: Optional[Dict[AgentID, SampleBatch]] = None, + episode=None, +) -> SampleBatch: + """Postprocesses a trajectory and returns the processed trajectory. + + The trajectory contains only data from one episode and from one agent. + - If `config.batch_mode=truncate_episodes` (default), sample_batch may + contain a truncated (at-the-end) episode, in case the + `config.rollout_fragment_length` was reached by the sampler. + - If `config.batch_mode=complete_episodes`, sample_batch will contain + exactly one episode (no matter how long). + New columns can be added to sample_batch and existing ones may be altered. + + Args: + policy: The Policy used to generate the trajectory + (`sample_batch`) + sample_batch: The SampleBatch to postprocess. + other_agent_batches (Optional[Dict[AgentID, SampleBatch]]): Optional + dict of AgentIDs mapping to other agents' trajectory data (from the + same episode). NOTE: The other agents use the same policy. + episode (Optional[Episode]): Optional multi-agent episode + object in which the agents operated. + + Returns: + SampleBatch: The postprocessed, modified SampleBatch (or a new one). + """ + return postprocess_nstep_and_prio(policy, sample_batch) + + +def _get_dist_class( + policy: Policy, config: AlgorithmConfigDict, action_space: gym.spaces.Space +) -> Type[TFActionDistribution]: + """Helper function to return a dist class based on config and action space. + + Args: + policy: The policy for which to return the action + dist class. + config: The Algorithm's config dict. + action_space (gym.spaces.Space): The action space used. + + Returns: + Type[TFActionDistribution]: A TF distribution class. + """ + if hasattr(policy, "dist_class") and policy.dist_class is not None: + return policy.dist_class + elif config["model"].get("custom_action_dist"): + action_dist_class, _ = ModelCatalog.get_action_dist( + action_space, config["model"], framework="tf" + ) + return action_dist_class + elif isinstance(action_space, Discrete): + return Categorical + elif isinstance(action_space, Simplex): + return Dirichlet + else: + assert isinstance(action_space, Box) + if config["normalize_actions"]: + return SquashedGaussian if not config["_use_beta_distribution"] else Beta + else: + return DiagGaussian + + +def get_distribution_inputs_and_class( + policy: Policy, + model: ModelV2, + obs_batch: TensorType, + *, + explore: bool = True, + **kwargs +) -> Tuple[TensorType, Type[TFActionDistribution], List[TensorType]]: + """The action distribution function to be used the algorithm. + + An action distribution function is used to customize the choice of action + distribution class and the resulting action distribution inputs (to + parameterize the distribution object). + After parameterizing the distribution, a `sample()` call + will be made on it to generate actions. + + Args: + policy: The Policy being queried for actions and calling this + function. + model: The SAC specific Model to use to generate the + distribution inputs (see sac_tf|torch_model.py). Must support the + `get_action_model_outputs` method. + obs_batch: The observations to be used as inputs to the + model. + explore: Whether to activate exploration or not. + + Returns: + Tuple[TensorType, Type[TFActionDistribution], List[TensorType]]: The + dist inputs, dist class, and a list of internal state outputs + (in the RNN case). + """ + # Get base-model (forward) output (this should be a noop call). + forward_out, state_out = model( + SampleBatch(obs=obs_batch, _is_training=policy._get_is_training_placeholder()), + [], + None, + ) + # Use the base output to get the policy outputs from the SAC model's + # policy components. + distribution_inputs, _ = model.get_action_model_outputs(forward_out) + # Get a distribution class to be used with the just calculated dist-inputs. + action_dist_class = _get_dist_class(policy, policy.config, policy.action_space) + + return distribution_inputs, action_dist_class, state_out + + +def sac_actor_critic_loss( + policy: Policy, + model: ModelV2, + dist_class: Type[TFActionDistribution], + train_batch: SampleBatch, +) -> Union[TensorType, List[TensorType]]: + """Constructs the loss for the Soft Actor Critic. + + Args: + policy: The Policy to calculate the loss for. + model (ModelV2): The Model to calculate the loss for. + dist_class (Type[ActionDistribution]: The action distr. class. + train_batch: The training data. + + Returns: + Union[TensorType, List[TensorType]]: A single loss tensor or a list + of loss tensors. + """ + # Should be True only for debugging purposes (e.g. test cases)! + deterministic = policy.config["_deterministic_loss"] + + _is_training = policy._get_is_training_placeholder() + # Get the base model output from the train batch. + model_out_t, _ = model( + SampleBatch(obs=train_batch[SampleBatch.CUR_OBS], _is_training=_is_training), + [], + None, + ) + + # Get the base model output from the next observations in the train batch. + model_out_tp1, _ = model( + SampleBatch(obs=train_batch[SampleBatch.NEXT_OBS], _is_training=_is_training), + [], + None, + ) + + # Get the target model's base outputs from the next observations in the + # train batch. + target_model_out_tp1, _ = policy.target_model( + SampleBatch(obs=train_batch[SampleBatch.NEXT_OBS], _is_training=_is_training), + [], + None, + ) + + # Discrete actions case. + if model.discrete: + # Get all action probs directly from pi and form their logp. + action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t) + log_pis_t = tf.nn.log_softmax(action_dist_inputs_t, -1) + policy_t = tf.math.exp(log_pis_t) + + action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1) + log_pis_tp1 = tf.nn.log_softmax(action_dist_inputs_tp1, -1) + policy_tp1 = tf.math.exp(log_pis_tp1) + + # Q-values. + q_t, _ = model.get_q_values(model_out_t) + # Target Q-values. + q_tp1, _ = policy.target_model.get_q_values(target_model_out_tp1) + if policy.config["twin_q"]: + twin_q_t, _ = model.get_twin_q_values(model_out_t) + twin_q_tp1, _ = policy.target_model.get_twin_q_values(target_model_out_tp1) + q_tp1 = tf.reduce_min((q_tp1, twin_q_tp1), axis=0) + q_tp1 -= model.alpha * log_pis_tp1 + + # Actually selected Q-values (from the actions batch). + one_hot = tf.one_hot( + train_batch[SampleBatch.ACTIONS], depth=q_t.shape.as_list()[-1] + ) + q_t_selected = tf.reduce_sum(q_t * one_hot, axis=-1) + if policy.config["twin_q"]: + twin_q_t_selected = tf.reduce_sum(twin_q_t * one_hot, axis=-1) + # Discrete case: "Best" means weighted by the policy (prob) outputs. + q_tp1_best = tf.reduce_sum(tf.multiply(policy_tp1, q_tp1), axis=-1) + q_tp1_best_masked = ( + 1.0 - tf.cast(train_batch[SampleBatch.TERMINATEDS], tf.float32) + ) * q_tp1_best + # Continuous actions case. + else: + # Sample simgle actions from distribution. + action_dist_class = _get_dist_class(policy, policy.config, policy.action_space) + action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t) + action_dist_t = action_dist_class(action_dist_inputs_t, policy.model) + policy_t = ( + action_dist_t.sample() + if not deterministic + else action_dist_t.deterministic_sample() + ) + log_pis_t = tf.expand_dims(action_dist_t.logp(policy_t), -1) + + action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1) + action_dist_tp1 = action_dist_class(action_dist_inputs_tp1, policy.model) + policy_tp1 = ( + action_dist_tp1.sample() + if not deterministic + else action_dist_tp1.deterministic_sample() + ) + log_pis_tp1 = tf.expand_dims(action_dist_tp1.logp(policy_tp1), -1) + + # Q-values for the actually selected actions. + q_t, _ = model.get_q_values( + model_out_t, tf.cast(train_batch[SampleBatch.ACTIONS], tf.float32) + ) + if policy.config["twin_q"]: + twin_q_t, _ = model.get_twin_q_values( + model_out_t, tf.cast(train_batch[SampleBatch.ACTIONS], tf.float32) + ) + + # Q-values for current policy in given current state. + q_t_det_policy, _ = model.get_q_values(model_out_t, policy_t) + if policy.config["twin_q"]: + twin_q_t_det_policy, _ = model.get_twin_q_values(model_out_t, policy_t) + q_t_det_policy = tf.reduce_min( + (q_t_det_policy, twin_q_t_det_policy), axis=0 + ) + + # target q network evaluation + q_tp1, _ = policy.target_model.get_q_values(target_model_out_tp1, policy_tp1) + if policy.config["twin_q"]: + twin_q_tp1, _ = policy.target_model.get_twin_q_values( + target_model_out_tp1, policy_tp1 + ) + # Take min over both twin-NNs. + q_tp1 = tf.reduce_min((q_tp1, twin_q_tp1), axis=0) + + q_t_selected = tf.squeeze(q_t, axis=len(q_t.shape) - 1) + if policy.config["twin_q"]: + twin_q_t_selected = tf.squeeze(twin_q_t, axis=len(q_t.shape) - 1) + q_tp1 -= model.alpha * log_pis_tp1 + + q_tp1_best = tf.squeeze(input=q_tp1, axis=len(q_tp1.shape) - 1) + q_tp1_best_masked = ( + 1.0 - tf.cast(train_batch[SampleBatch.TERMINATEDS], tf.float32) + ) * q_tp1_best + + # Compute RHS of bellman equation for the Q-loss (critic(s)). + q_t_selected_target = tf.stop_gradient( + tf.cast(train_batch[SampleBatch.REWARDS], tf.float32) + + policy.config["gamma"] ** policy.config["n_step"] * q_tp1_best_masked + ) + + # Compute the TD-error (potentially clipped). + base_td_error = tf.math.abs(q_t_selected - q_t_selected_target) + if policy.config["twin_q"]: + twin_td_error = tf.math.abs(twin_q_t_selected - q_t_selected_target) + td_error = 0.5 * (base_td_error + twin_td_error) + else: + td_error = base_td_error + + # Calculate one or two critic losses (2 in the twin_q case). + prio_weights = tf.cast(train_batch[PRIO_WEIGHTS], tf.float32) + critic_loss = [tf.reduce_mean(prio_weights * huber_loss(base_td_error))] + if policy.config["twin_q"]: + critic_loss.append(tf.reduce_mean(prio_weights * huber_loss(twin_td_error))) + + # Alpha- and actor losses. + # Note: In the papers, alpha is used directly, here we take the log. + # Discrete case: Multiply the action probs as weights with the original + # loss terms (no expectations needed). + if model.discrete: + alpha_loss = tf.reduce_mean( + tf.reduce_sum( + tf.multiply( + tf.stop_gradient(policy_t), + -model.log_alpha + * tf.stop_gradient(log_pis_t + model.target_entropy), + ), + axis=-1, + ) + ) + actor_loss = tf.reduce_mean( + tf.reduce_sum( + tf.multiply( + # NOTE: No stop_grad around policy output here + # (compare with q_t_det_policy for continuous case). + policy_t, + model.alpha * log_pis_t - tf.stop_gradient(q_t), + ), + axis=-1, + ) + ) + else: + alpha_loss = -tf.reduce_mean( + model.log_alpha * tf.stop_gradient(log_pis_t + model.target_entropy) + ) + actor_loss = tf.reduce_mean(model.alpha * log_pis_t - q_t_det_policy) + + # Save for stats function. + policy.policy_t = policy_t + policy.q_t = q_t + policy.td_error = td_error + policy.actor_loss = actor_loss + policy.critic_loss = critic_loss + policy.alpha_loss = alpha_loss + policy.alpha_value = model.alpha + policy.target_entropy = model.target_entropy + + # In a custom apply op we handle the losses separately, but return them + # combined in one loss here. + return actor_loss + tf.math.add_n(critic_loss) + alpha_loss + + +def compute_and_clip_gradients( + policy: Policy, optimizer: LocalOptimizer, loss: TensorType +) -> ModelGradients: + """Gradients computing function (from loss tensor, using local optimizer). + + Note: For SAC, optimizer and loss are ignored b/c we have 3 + losses and 3 local optimizers (all stored in policy). + `optimizer` will be used, though, in the tf-eager case b/c it is then a + fake optimizer (OptimizerWrapper) object with a `tape` property to + generate a GradientTape object for gradient recording. + + Args: + policy: The Policy object that generated the loss tensor and + that holds the given local optimizer. + optimizer: The tf (local) optimizer object to + calculate the gradients with. + loss: The loss tensor for which gradients should be + calculated. + + Returns: + ModelGradients: List of the possibly clipped gradients- and variable + tuples. + """ + # Eager: Use GradientTape (which is a property of the `optimizer` object + # (an OptimizerWrapper): see rllib/policy/eager_tf_policy.py). + if policy.config["framework"] == "tf2": + tape = optimizer.tape + pol_weights = policy.model.policy_variables() + actor_grads_and_vars = list( + zip(tape.gradient(policy.actor_loss, pol_weights), pol_weights) + ) + q_weights = policy.model.q_variables() + if policy.config["twin_q"]: + half_cutoff = len(q_weights) // 2 + grads_1 = tape.gradient(policy.critic_loss[0], q_weights[:half_cutoff]) + grads_2 = tape.gradient(policy.critic_loss[1], q_weights[half_cutoff:]) + critic_grads_and_vars = list(zip(grads_1, q_weights[:half_cutoff])) + list( + zip(grads_2, q_weights[half_cutoff:]) + ) + else: + critic_grads_and_vars = list( + zip(tape.gradient(policy.critic_loss[0], q_weights), q_weights) + ) + + alpha_vars = [policy.model.log_alpha] + alpha_grads_and_vars = list( + zip(tape.gradient(policy.alpha_loss, alpha_vars), alpha_vars) + ) + # Tf1.x: Use optimizer.compute_gradients() + else: + actor_grads_and_vars = policy._actor_optimizer.compute_gradients( + policy.actor_loss, var_list=policy.model.policy_variables() + ) + + q_weights = policy.model.q_variables() + if policy.config["twin_q"]: + half_cutoff = len(q_weights) // 2 + base_q_optimizer, twin_q_optimizer = policy._critic_optimizer + critic_grads_and_vars = base_q_optimizer.compute_gradients( + policy.critic_loss[0], var_list=q_weights[:half_cutoff] + ) + twin_q_optimizer.compute_gradients( + policy.critic_loss[1], var_list=q_weights[half_cutoff:] + ) + else: + critic_grads_and_vars = policy._critic_optimizer[0].compute_gradients( + policy.critic_loss[0], var_list=q_weights + ) + alpha_grads_and_vars = policy._alpha_optimizer.compute_gradients( + policy.alpha_loss, var_list=[policy.model.log_alpha] + ) + + # Clip if necessary. + if policy.config["grad_clip"]: + clip_func = partial(tf.clip_by_norm, clip_norm=policy.config["grad_clip"]) + else: + clip_func = tf.identity + + # Save grads and vars for later use in `build_apply_op`. + policy._actor_grads_and_vars = [ + (clip_func(g), v) for (g, v) in actor_grads_and_vars if g is not None + ] + policy._critic_grads_and_vars = [ + (clip_func(g), v) for (g, v) in critic_grads_and_vars if g is not None + ] + policy._alpha_grads_and_vars = [ + (clip_func(g), v) for (g, v) in alpha_grads_and_vars if g is not None + ] + + grads_and_vars = ( + policy._actor_grads_and_vars + + policy._critic_grads_and_vars + + policy._alpha_grads_and_vars + ) + return grads_and_vars + + +def apply_gradients( + policy: Policy, optimizer: LocalOptimizer, grads_and_vars: ModelGradients +) -> Union["tf.Operation", None]: + """Gradients applying function (from list of "grad_and_var" tuples). + + Note: For SAC, optimizer and grads_and_vars are ignored b/c we have 3 + losses and optimizers (stored in policy). + + Args: + policy: The Policy object whose Model(s) the given gradients + should be applied to. + optimizer: The tf (local) optimizer object through + which to apply the gradients. + grads_and_vars: The list of grad_and_var tuples to + apply via the given optimizer. + + Returns: + Union[tf.Operation, None]: The tf op to be used to run the apply + operation. None for eager mode. + """ + actor_apply_ops = policy._actor_optimizer.apply_gradients( + policy._actor_grads_and_vars + ) + + cgrads = policy._critic_grads_and_vars + half_cutoff = len(cgrads) // 2 + if policy.config["twin_q"]: + critic_apply_ops = [ + policy._critic_optimizer[0].apply_gradients(cgrads[:half_cutoff]), + policy._critic_optimizer[1].apply_gradients(cgrads[half_cutoff:]), + ] + else: + critic_apply_ops = [policy._critic_optimizer[0].apply_gradients(cgrads)] + + # Eager mode -> Just apply and return None. + if policy.config["framework"] == "tf2": + policy._alpha_optimizer.apply_gradients(policy._alpha_grads_and_vars) + return + # Tf static graph -> Return op. + else: + alpha_apply_ops = policy._alpha_optimizer.apply_gradients( + policy._alpha_grads_and_vars, + global_step=tf1.train.get_or_create_global_step(), + ) + return tf.group([actor_apply_ops, alpha_apply_ops] + critic_apply_ops) + + +def stats(policy: Policy, train_batch: SampleBatch) -> Dict[str, TensorType]: + """Stats function for SAC. Returns a dict with important loss stats. + + Args: + policy: The Policy to generate stats for. + train_batch: The SampleBatch (already) used for training. + + Returns: + Dict[str, TensorType]: The stats dict. + """ + return { + "mean_td_error": tf.reduce_mean(policy.td_error), + "actor_loss": tf.reduce_mean(policy.actor_loss), + "critic_loss": tf.reduce_mean(policy.critic_loss), + "alpha_loss": tf.reduce_mean(policy.alpha_loss), + "alpha_value": tf.reduce_mean(policy.alpha_value), + "target_entropy": tf.constant(policy.target_entropy), + "mean_q": tf.reduce_mean(policy.q_t), + "max_q": tf.reduce_max(policy.q_t), + "min_q": tf.reduce_min(policy.q_t), + } + + +class ActorCriticOptimizerMixin: + """Mixin class to generate the necessary optimizers for actor-critic algos. + + - Creates global step for counting the number of update operations. + - Creates separate optimizers for actor, critic, and alpha. + """ + + def __init__(self, config): + # Eager mode. + if config["framework"] == "tf2": + self.global_step = get_variable(0, tf_name="global_step") + self._actor_optimizer = tf.keras.optimizers.Adam( + learning_rate=config["optimization"]["actor_learning_rate"] + ) + self._critic_optimizer = [ + tf.keras.optimizers.Adam( + learning_rate=config["optimization"]["critic_learning_rate"] + ) + ] + if config["twin_q"]: + self._critic_optimizer.append( + tf.keras.optimizers.Adam( + learning_rate=config["optimization"]["critic_learning_rate"] + ) + ) + self._alpha_optimizer = tf.keras.optimizers.Adam( + learning_rate=config["optimization"]["entropy_learning_rate"] + ) + # Static graph mode. + else: + self.global_step = tf1.train.get_or_create_global_step() + self._actor_optimizer = tf1.train.AdamOptimizer( + learning_rate=config["optimization"]["actor_learning_rate"] + ) + self._critic_optimizer = [ + tf1.train.AdamOptimizer( + learning_rate=config["optimization"]["critic_learning_rate"] + ) + ] + if config["twin_q"]: + self._critic_optimizer.append( + tf1.train.AdamOptimizer( + learning_rate=config["optimization"]["critic_learning_rate"] + ) + ) + self._alpha_optimizer = tf1.train.AdamOptimizer( + learning_rate=config["optimization"]["entropy_learning_rate"] + ) + + +def setup_early_mixins( + policy: Policy, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, +) -> None: + """Call mixin classes' constructors before Policy's initialization. + + Adds the necessary optimizers to the given Policy. + + Args: + policy: The Policy object. + obs_space (gym.spaces.Space): The Policy's observation space. + action_space (gym.spaces.Space): The Policy's action space. + config: The Policy's config. + """ + ActorCriticOptimizerMixin.__init__(policy, config) + + +# TODO: Unify with DDPG's ComputeTDErrorMixin when SAC policy subclasses PolicyV2 +class ComputeTDErrorMixin: + def __init__(self, loss_fn): + @make_tf_callable(self.get_session(), dynamic_shape=True) + def compute_td_error( + obs_t, act_t, rew_t, obs_tp1, terminateds_mask, importance_weights + ): + # Do forward pass on loss to update td errors attribute + # (one TD-error value per item in batch to update PR weights). + loss_fn( + self, + self.model, + None, + { + SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_t), + SampleBatch.ACTIONS: tf.convert_to_tensor(act_t), + SampleBatch.REWARDS: tf.convert_to_tensor(rew_t), + SampleBatch.NEXT_OBS: tf.convert_to_tensor(obs_tp1), + SampleBatch.TERMINATEDS: tf.convert_to_tensor(terminateds_mask), + PRIO_WEIGHTS: tf.convert_to_tensor(importance_weights), + }, + ) + # `self.td_error` is set in loss_fn. + return self.td_error + + self.compute_td_error = compute_td_error + + +def setup_mid_mixins( + policy: Policy, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, +) -> None: + """Call mixin classes' constructors before Policy's loss initialization. + + Adds the `compute_td_error` method to the given policy. + Calling `compute_td_error` with batch data will re-calculate the loss + on that batch AND return the per-batch-item TD-error for prioritized + replay buffer record weight updating (in case a prioritized replay buffer + is used). + + Args: + policy: The Policy object. + obs_space (gym.spaces.Space): The Policy's observation space. + action_space (gym.spaces.Space): The Policy's action space. + config: The Policy's config. + """ + ComputeTDErrorMixin.__init__(policy, sac_actor_critic_loss) + + +def setup_late_mixins( + policy: Policy, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, +) -> None: + """Call mixin classes' constructors after Policy initialization. + + Adds the `update_target` method to the given policy. + Calling `update_target` updates all target Q-networks' weights from their + respective "main" Q-metworks, based on tau (smooth, partial updating). + + Args: + policy: The Policy object. + obs_space (gym.spaces.Space): The Policy's observation space. + action_space (gym.spaces.Space): The Policy's action space. + config: The Policy's config. + """ + TargetNetworkMixin.__init__(policy) + + +def validate_spaces( + policy: Policy, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, +) -> None: + """Validates the observation- and action spaces used for the Policy. + + Args: + policy: The policy, whose spaces are being validated. + observation_space (gym.spaces.Space): The observation space to + validate. + action_space (gym.spaces.Space): The action space to validate. + config: The Policy's config dict. + + Raises: + UnsupportedSpaceException: If one of the spaces is not supported. + """ + # Only support single Box or single Discrete spaces. + if not isinstance(action_space, (Box, Discrete, Simplex)): + raise UnsupportedSpaceException( + "Action space ({}) of {} is not supported for " + "SAC. Must be [Box|Discrete|Simplex].".format(action_space, policy) + ) + # If Box, make sure it's a 1D vector space. + elif isinstance(action_space, (Box, Simplex)) and len(action_space.shape) > 1: + raise UnsupportedSpaceException( + "Action space ({}) of {} has multiple dimensions " + "{}. ".format(action_space, policy, action_space.shape) + + "Consider reshaping this into a single dimension, " + "using a Tuple action space, or the multi-agent API." + ) + + +# Build a child class of `DynamicTFPolicy`, given the custom functions defined +# above. +SACTFPolicy = build_tf_policy( + name="SACTFPolicy", + get_default_config=lambda: ray.rllib.algorithms.sac.sac.SACConfig(), + make_model=build_sac_model, + postprocess_fn=postprocess_trajectory, + action_distribution_fn=get_distribution_inputs_and_class, + loss_fn=sac_actor_critic_loss, + stats_fn=stats, + compute_gradients_fn=compute_and_clip_gradients, + apply_gradients_fn=apply_gradients, + extra_learn_fetches_fn=lambda policy: {"td_error": policy.td_error}, + mixins=[TargetNetworkMixin, ActorCriticOptimizerMixin, ComputeTDErrorMixin], + validate_spaces=validate_spaces, + before_init=setup_early_mixins, + before_loss_init=setup_mid_mixins, + after_init=setup_late_mixins, +) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_torch_model.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_torch_model.py new file mode 100644 index 0000000000000000000000000000000000000000..8c2fcd5b530c8948fe1e1b985357f1645b0a6e14 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_torch_model.py @@ -0,0 +1,330 @@ +from typing import Dict, List, Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.typing import ModelConfigDict, TensorStructType, TensorType + +torch, nn = try_import_torch() + + +class SACTorchModel(TorchModelV2, nn.Module): + """Extension of the standard TorchModelV2 for SAC. + + To customize, do one of the following: + - sub-class SACTorchModel and override one or more of its methods. + - Use SAC's `q_model_config` and `policy_model` keys to tweak the default model + behaviors (e.g. fcnet_hiddens, conv_filters, etc..). + - Use SAC's `q_model_config->custom_model` and `policy_model->custom_model` keys + to specify your own custom Q-model(s) and policy-models, which will be + created within this SACTFModel (see `build_policy_model` and + `build_q_model`. + + Note: It is not recommended to override the `forward` method for SAC. This + would lead to shared weights (between policy and Q-nets), which will then + not be optimized by either of the critic- or actor-optimizers! + + Data flow: + `obs` -> forward() (should stay a noop method!) -> `model_out` + `model_out` -> get_policy_output() -> pi(actions|obs) + `model_out`, `actions` -> get_q_values() -> Q(s, a) + `model_out`, `actions` -> get_twin_q_values() -> Q_twin(s, a) + """ + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: Optional[int], + model_config: ModelConfigDict, + name: str, + policy_model_config: ModelConfigDict = None, + q_model_config: ModelConfigDict = None, + twin_q: bool = False, + initial_alpha: float = 1.0, + target_entropy: Optional[float] = None, + ): + """Initializes a SACTorchModel instance. + 7 + Args: + policy_model_config: The config dict for the + policy network. + q_model_config: The config dict for the + Q-network(s) (2 if twin_q=True). + twin_q: Build twin Q networks (Q-net and target) for more + stable Q-learning. + initial_alpha: The initial value for the to-be-optimized + alpha parameter (default: 1.0). + target_entropy (Optional[float]): A target entropy value for + the to-be-optimized alpha parameter. If None, will use the + defaults described in the papers for SAC (and discrete SAC). + + Note that the core layers for forward() are not defined here, this + only defines the layers for the output heads. Those layers for + forward() should be defined in subclasses of SACModel. + """ + nn.Module.__init__(self) + super(SACTorchModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + + if isinstance(action_space, Discrete): + self.action_dim = action_space.n + self.discrete = True + action_outs = q_outs = self.action_dim + elif isinstance(action_space, Box): + self.action_dim = np.prod(action_space.shape) + self.discrete = False + action_outs = 2 * self.action_dim + q_outs = 1 + else: + assert isinstance(action_space, Simplex) + self.action_dim = np.prod(action_space.shape) + self.discrete = False + action_outs = self.action_dim + q_outs = 1 + + # Build the policy network. + self.action_model = self.build_policy_model( + self.obs_space, action_outs, policy_model_config, "policy_model" + ) + + # Build the Q-network(s). + self.q_net = self.build_q_model( + self.obs_space, self.action_space, q_outs, q_model_config, "q" + ) + if twin_q: + self.twin_q_net = self.build_q_model( + self.obs_space, self.action_space, q_outs, q_model_config, "twin_q" + ) + else: + self.twin_q_net = None + + log_alpha = nn.Parameter( + torch.from_numpy(np.array([np.log(initial_alpha)])).float() + ) + self.register_parameter("log_alpha", log_alpha) + + # Auto-calculate the target entropy. + if target_entropy is None or target_entropy == "auto": + # See hyperparams in [2] (README.md). + if self.discrete: + target_entropy = 0.98 * np.array( + -np.log(1.0 / action_space.n), dtype=np.float32 + ) + # See [1] (README.md). + else: + target_entropy = -np.prod(action_space.shape) + + target_entropy = nn.Parameter( + torch.from_numpy(np.array([target_entropy])).float(), requires_grad=False + ) + self.register_parameter("target_entropy", target_entropy) + + @override(TorchModelV2) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + """The common (Q-net and policy-net) forward pass. + + NOTE: It is not(!) recommended to override this method as it would + introduce a shared pre-network, which would be updated by both + actor- and critic optimizers. + """ + return input_dict["obs"], state + + def build_policy_model(self, obs_space, num_outputs, policy_model_config, name): + """Builds the policy model used by this SAC. + + Override this method in a sub-class of SACTFModel to implement your + own policy net. Alternatively, simply set `custom_model` within the + top level SAC `policy_model` config key to make this default + implementation of `build_policy_model` use your custom policy network. + + Returns: + TorchModelV2: The TorchModelV2 policy sub-model. + """ + model = ModelCatalog.get_model_v2( + obs_space, + self.action_space, + num_outputs, + policy_model_config, + framework="torch", + name=name, + ) + return model + + def build_q_model(self, obs_space, action_space, num_outputs, q_model_config, name): + """Builds one of the (twin) Q-nets used by this SAC. + + Override this method in a sub-class of SACTFModel to implement your + own Q-nets. Alternatively, simply set `custom_model` within the + top level SAC `q_model_config` config key to make this default implementation + of `build_q_model` use your custom Q-nets. + + Returns: + TorchModelV2: The TorchModelV2 Q-net sub-model. + """ + self.concat_obs_and_actions = False + if self.discrete: + input_space = obs_space + else: + orig_space = getattr(obs_space, "original_space", obs_space) + if isinstance(orig_space, Box) and len(orig_space.shape) == 1: + input_space = Box( + float("-inf"), + float("inf"), + shape=(orig_space.shape[0] + action_space.shape[0],), + ) + self.concat_obs_and_actions = True + else: + input_space = gym.spaces.Tuple([orig_space, action_space]) + + model = ModelCatalog.get_model_v2( + input_space, + action_space, + num_outputs, + q_model_config, + framework="torch", + name=name, + ) + return model + + def get_q_values( + self, model_out: TensorType, actions: Optional[TensorType] = None + ) -> TensorType: + """Returns Q-values, given the output of self.__call__(). + + This implements Q(s, a) -> [single Q-value] for the continuous case and + Q(s) -> [Q-values for all actions] for the discrete case. + + Args: + model_out: Feature outputs from the model layers + (result of doing `self.__call__(obs)`). + actions (Optional[TensorType]): Continuous action batch to return + Q-values for. Shape: [BATCH_SIZE, action_dim]. If None + (discrete action case), return Q-values for all actions. + + Returns: + TensorType: Q-values tensor of shape [BATCH_SIZE, 1]. + """ + return self._get_q_value(model_out, actions, self.q_net) + + def get_twin_q_values( + self, model_out: TensorType, actions: Optional[TensorType] = None + ) -> TensorType: + """Same as get_q_values but using the twin Q net. + + This implements the twin Q(s, a). + + Args: + model_out: Feature outputs from the model layers + (result of doing `self.__call__(obs)`). + actions (Optional[Tensor]): Actions to return the Q-values for. + Shape: [BATCH_SIZE, action_dim]. If None (discrete action + case), return Q-values for all actions. + + Returns: + TensorType: Q-values tensor of shape [BATCH_SIZE, 1]. + """ + return self._get_q_value(model_out, actions, self.twin_q_net) + + def _get_q_value(self, model_out, actions, net): + # Model outs may come as original Tuple observations, concat them + # here if this is the case. + if isinstance(net.obs_space, Box): + if isinstance(model_out, (list, tuple)): + model_out = torch.cat(model_out, dim=-1) + elif isinstance(model_out, dict): + model_out = torch.cat(list(model_out.values()), dim=-1) + + # Continuous case -> concat actions to model_out. + if actions is not None: + if self.concat_obs_and_actions: + input_dict = {"obs": torch.cat([model_out, actions], dim=-1)} + else: + # TODO(junogng) : SampleBatch doesn't support list columns yet. + # Use ModelInputDict. + input_dict = {"obs": (model_out, actions)} + # Discrete case -> return q-vals for all actions. + else: + input_dict = {"obs": model_out} + # Switch on training mode (when getting Q-values, we are usually in + # training). + input_dict["is_training"] = True + + return net(input_dict, [], None) + + def get_action_model_outputs( + self, + model_out: TensorType, + state_in: List[TensorType] = None, + seq_lens: TensorType = None, + ) -> (TensorType, List[TensorType]): + """Returns distribution inputs and states given the output of + policy.model(). + + For continuous action spaces, these will be the mean/stddev + distribution inputs for the (SquashedGaussian) action distribution. + For discrete action spaces, these will be the logits for a categorical + distribution. + + Args: + model_out: Feature outputs from the model layers + (result of doing `model(obs)`). + state_in List(TensorType): State input for recurrent cells + seq_lens: Sequence lengths of input- and state + sequences + + Returns: + TensorType: Distribution inputs for sampling actions. + """ + + def concat_obs_if_necessary(obs: TensorStructType): + """Concat model outs if they come as original tuple observations.""" + if isinstance(obs, (list, tuple)): + obs = torch.cat(obs, dim=-1) + elif isinstance(obs, dict): + obs = torch.cat( + [ + torch.unsqueeze(val, 1) if len(val.shape) == 1 else val + for val in tree.flatten(obs.values()) + ], + dim=-1, + ) + return obs + + if state_in is None: + state_in = [] + + if isinstance(model_out, dict) and "obs" in model_out: + # Model outs may come as original Tuple observations + if isinstance(self.action_model.obs_space, Box): + model_out["obs"] = concat_obs_if_necessary(model_out["obs"]) + return self.action_model(model_out, state_in, seq_lens) + else: + if isinstance(self.action_model.obs_space, Box): + model_out = concat_obs_if_necessary(model_out) + return self.action_model({"obs": model_out}, state_in, seq_lens) + + def policy_variables(self): + """Return the list of variables for the policy net.""" + + return self.action_model.variables() + + def q_variables(self): + """Return the list of variables for Q / twin Q nets.""" + + return self.q_net.variables() + ( + self.twin_q_net.variables() if self.twin_q_net else [] + ) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_torch_policy.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_torch_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..b105b856ed0b7f1ce2489a8aebfccb28fc0c7e76 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/sac_torch_policy.py @@ -0,0 +1,517 @@ +""" +PyTorch policy class used for SAC. +""" + +import logging +from typing import Dict, List, Optional, Tuple, Type, Union + +import gymnasium as gym +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete + +import ray +from ray.rllib.algorithms.dqn.dqn_tf_policy import PRIO_WEIGHTS +from ray.rllib.algorithms.sac.sac_tf_policy import ( + build_sac_model, + postprocess_trajectory, + validate_spaces, +) +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.torch_action_dist import ( + TorchBeta, + TorchCategorical, + TorchDiagGaussian, + TorchDirichlet, + TorchDistributionWrapper, + TorchSquashedGaussian, +) +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.policy_template import build_policy_class +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_mixins import TargetNetworkMixin +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.torch_utils import ( + apply_grad_clipping, + concat_multi_gpu_td_errors, + huber_loss, +) +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + LocalOptimizer, + ModelInputDict, + TensorType, +) + +torch, nn = try_import_torch() +F = nn.functional + +logger = logging.getLogger(__name__) + + +def _get_dist_class( + policy: Policy, config: AlgorithmConfigDict, action_space: gym.spaces.Space +) -> Type[TorchDistributionWrapper]: + """Helper function to return a dist class based on config and action space. + + Args: + policy: The policy for which to return the action + dist class. + config: The Algorithm's config dict. + action_space (gym.spaces.Space): The action space used. + + Returns: + Type[TFActionDistribution]: A TF distribution class. + """ + if hasattr(policy, "dist_class") and policy.dist_class is not None: + return policy.dist_class + elif config["model"].get("custom_action_dist"): + action_dist_class, _ = ModelCatalog.get_action_dist( + action_space, config["model"], framework="torch" + ) + return action_dist_class + elif isinstance(action_space, Discrete): + return TorchCategorical + elif isinstance(action_space, Simplex): + return TorchDirichlet + else: + assert isinstance(action_space, Box) + if config["normalize_actions"]: + return ( + TorchSquashedGaussian + if not config["_use_beta_distribution"] + else TorchBeta + ) + else: + return TorchDiagGaussian + + +def build_sac_model_and_action_dist( + policy: Policy, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, +) -> Tuple[ModelV2, Type[TorchDistributionWrapper]]: + """Constructs the necessary ModelV2 and action dist class for the Policy. + + Args: + policy: The TFPolicy that will use the models. + obs_space (gym.spaces.Space): The observation space. + action_space (gym.spaces.Space): The action space. + config: The SACConfig object. + + Returns: + ModelV2: The ModelV2 to be used by the Policy. Note: An additional + target model will be created in this function and assigned to + `policy.target_model`. + """ + model = build_sac_model(policy, obs_space, action_space, config) + action_dist_class = _get_dist_class(policy, config, action_space) + return model, action_dist_class + + +def action_distribution_fn( + policy: Policy, + model: ModelV2, + input_dict: ModelInputDict, + *, + state_batches: Optional[List[TensorType]] = None, + seq_lens: Optional[TensorType] = None, + prev_action_batch: Optional[TensorType] = None, + prev_reward_batch=None, + explore: Optional[bool] = None, + timestep: Optional[int] = None, + is_training: Optional[bool] = None +) -> Tuple[TensorType, Type[TorchDistributionWrapper], List[TensorType]]: + """The action distribution function to be used the algorithm. + + An action distribution function is used to customize the choice of action + distribution class and the resulting action distribution inputs (to + parameterize the distribution object). + After parameterizing the distribution, a `sample()` call + will be made on it to generate actions. + + Args: + policy: The Policy being queried for actions and calling this + function. + model (TorchModelV2): The SAC specific model to use to generate the + distribution inputs (see sac_tf|torch_model.py). Must support the + `get_action_model_outputs` method. + input_dict: The input-dict to be used for the model + call. + state_batches (Optional[List[TensorType]]): The list of internal state + tensor batches. + seq_lens (Optional[TensorType]): The tensor of sequence lengths used + in RNNs. + prev_action_batch (Optional[TensorType]): Optional batch of prev + actions used by the model. + prev_reward_batch (Optional[TensorType]): Optional batch of prev + rewards used by the model. + explore (Optional[bool]): Whether to activate exploration or not. If + None, use value of `config.explore`. + timestep (Optional[int]): An optional timestep. + is_training (Optional[bool]): An optional is-training flag. + + Returns: + Tuple[TensorType, Type[TorchDistributionWrapper], List[TensorType]]: + The dist inputs, dist class, and a list of internal state outputs + (in the RNN case). + """ + # Get base-model output (w/o the SAC specific parts of the network). + model_out, _ = model(input_dict, [], None) + # Use the base output to get the policy outputs from the SAC model's + # policy components. + action_dist_inputs, _ = model.get_action_model_outputs(model_out) + # Get a distribution class to be used with the just calculated dist-inputs. + action_dist_class = _get_dist_class(policy, policy.config, policy.action_space) + + return action_dist_inputs, action_dist_class, [] + + +def actor_critic_loss( + policy: Policy, + model: ModelV2, + dist_class: Type[TorchDistributionWrapper], + train_batch: SampleBatch, +) -> Union[TensorType, List[TensorType]]: + """Constructs the loss for the Soft Actor Critic. + + Args: + policy: The Policy to calculate the loss for. + model (ModelV2): The Model to calculate the loss for. + dist_class (Type[TorchDistributionWrapper]: The action distr. class. + train_batch: The training data. + + Returns: + Union[TensorType, List[TensorType]]: A single loss tensor or a list + of loss tensors. + """ + # Look up the target model (tower) using the model tower. + target_model = policy.target_models[model] + + # Should be True only for debugging purposes (e.g. test cases)! + deterministic = policy.config["_deterministic_loss"] + + model_out_t, _ = model( + SampleBatch(obs=train_batch[SampleBatch.CUR_OBS], _is_training=True), [], None + ) + + model_out_tp1, _ = model( + SampleBatch(obs=train_batch[SampleBatch.NEXT_OBS], _is_training=True), [], None + ) + + target_model_out_tp1, _ = target_model( + SampleBatch(obs=train_batch[SampleBatch.NEXT_OBS], _is_training=True), [], None + ) + + alpha = torch.exp(model.log_alpha) + + # Discrete case. + if model.discrete: + # Get all action probs directly from pi and form their logp. + action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t) + log_pis_t = F.log_softmax(action_dist_inputs_t, dim=-1) + policy_t = torch.exp(log_pis_t) + action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1) + log_pis_tp1 = F.log_softmax(action_dist_inputs_tp1, -1) + policy_tp1 = torch.exp(log_pis_tp1) + # Q-values. + q_t, _ = model.get_q_values(model_out_t) + # Target Q-values. + q_tp1, _ = target_model.get_q_values(target_model_out_tp1) + if policy.config["twin_q"]: + twin_q_t, _ = model.get_twin_q_values(model_out_t) + twin_q_tp1, _ = target_model.get_twin_q_values(target_model_out_tp1) + q_tp1 = torch.min(q_tp1, twin_q_tp1) + q_tp1 -= alpha * log_pis_tp1 + + # Actually selected Q-values (from the actions batch). + one_hot = F.one_hot( + train_batch[SampleBatch.ACTIONS].long(), num_classes=q_t.size()[-1] + ) + q_t_selected = torch.sum(q_t * one_hot, dim=-1) + if policy.config["twin_q"]: + twin_q_t_selected = torch.sum(twin_q_t * one_hot, dim=-1) + # Discrete case: "Best" means weighted by the policy (prob) outputs. + q_tp1_best = torch.sum(torch.mul(policy_tp1, q_tp1), dim=-1) + q_tp1_best_masked = ( + 1.0 - train_batch[SampleBatch.TERMINATEDS].float() + ) * q_tp1_best + # Continuous actions case. + else: + # Sample single actions from distribution. + action_dist_class = _get_dist_class(policy, policy.config, policy.action_space) + action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t) + action_dist_t = action_dist_class(action_dist_inputs_t, model) + policy_t = ( + action_dist_t.sample() + if not deterministic + else action_dist_t.deterministic_sample() + ) + log_pis_t = torch.unsqueeze(action_dist_t.logp(policy_t), -1) + action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1) + action_dist_tp1 = action_dist_class(action_dist_inputs_tp1, model) + policy_tp1 = ( + action_dist_tp1.sample() + if not deterministic + else action_dist_tp1.deterministic_sample() + ) + log_pis_tp1 = torch.unsqueeze(action_dist_tp1.logp(policy_tp1), -1) + + # Q-values for the actually selected actions. + q_t, _ = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS]) + if policy.config["twin_q"]: + twin_q_t, _ = model.get_twin_q_values( + model_out_t, train_batch[SampleBatch.ACTIONS] + ) + + # Q-values for current policy in given current state. + q_t_det_policy, _ = model.get_q_values(model_out_t, policy_t) + if policy.config["twin_q"]: + twin_q_t_det_policy, _ = model.get_twin_q_values(model_out_t, policy_t) + q_t_det_policy = torch.min(q_t_det_policy, twin_q_t_det_policy) + + # Target q network evaluation. + q_tp1, _ = target_model.get_q_values(target_model_out_tp1, policy_tp1) + if policy.config["twin_q"]: + twin_q_tp1, _ = target_model.get_twin_q_values( + target_model_out_tp1, policy_tp1 + ) + # Take min over both twin-NNs. + q_tp1 = torch.min(q_tp1, twin_q_tp1) + + q_t_selected = torch.squeeze(q_t, dim=-1) + if policy.config["twin_q"]: + twin_q_t_selected = torch.squeeze(twin_q_t, dim=-1) + q_tp1 -= alpha * log_pis_tp1 + + q_tp1_best = torch.squeeze(input=q_tp1, dim=-1) + q_tp1_best_masked = ( + 1.0 - train_batch[SampleBatch.TERMINATEDS].float() + ) * q_tp1_best + + # compute RHS of bellman equation + q_t_selected_target = ( + train_batch[SampleBatch.REWARDS] + + (policy.config["gamma"] ** policy.config["n_step"]) * q_tp1_best_masked + ).detach() + + # Compute the TD-error (potentially clipped). + base_td_error = torch.abs(q_t_selected - q_t_selected_target) + if policy.config["twin_q"]: + twin_td_error = torch.abs(twin_q_t_selected - q_t_selected_target) + td_error = 0.5 * (base_td_error + twin_td_error) + else: + td_error = base_td_error + + critic_loss = [torch.mean(train_batch[PRIO_WEIGHTS] * huber_loss(base_td_error))] + if policy.config["twin_q"]: + critic_loss.append( + torch.mean(train_batch[PRIO_WEIGHTS] * huber_loss(twin_td_error)) + ) + + # Alpha- and actor losses. + # Note: In the papers, alpha is used directly, here we take the log. + # Discrete case: Multiply the action probs as weights with the original + # loss terms (no expectations needed). + if model.discrete: + weighted_log_alpha_loss = policy_t.detach() * ( + -model.log_alpha * (log_pis_t + model.target_entropy).detach() + ) + # Sum up weighted terms and mean over all batch items. + alpha_loss = torch.mean(torch.sum(weighted_log_alpha_loss, dim=-1)) + # Actor loss. + actor_loss = torch.mean( + torch.sum( + torch.mul( + # NOTE: No stop_grad around policy output here + # (compare with q_t_det_policy for continuous case). + policy_t, + alpha.detach() * log_pis_t - q_t.detach(), + ), + dim=-1, + ) + ) + else: + alpha_loss = -torch.mean( + model.log_alpha * (log_pis_t + model.target_entropy).detach() + ) + # Note: Do not detach q_t_det_policy here b/c is depends partly + # on the policy vars (policy sample pushed through Q-net). + # However, we must make sure `actor_loss` is not used to update + # the Q-net(s)' variables. + actor_loss = torch.mean(alpha.detach() * log_pis_t - q_t_det_policy) + + # Store values for stats function in model (tower), such that for + # multi-GPU, we do not override them during the parallel loss phase. + model.tower_stats["q_t"] = q_t + model.tower_stats["policy_t"] = policy_t + model.tower_stats["log_pis_t"] = log_pis_t + model.tower_stats["actor_loss"] = actor_loss + model.tower_stats["critic_loss"] = critic_loss + model.tower_stats["alpha_loss"] = alpha_loss + + # TD-error tensor in final stats + # will be concatenated and retrieved for each individual batch item. + model.tower_stats["td_error"] = td_error + + # Return all loss terms corresponding to our optimizers. + return tuple([actor_loss] + critic_loss + [alpha_loss]) + + +def stats(policy: Policy, train_batch: SampleBatch) -> Dict[str, TensorType]: + """Stats function for SAC. Returns a dict with important loss stats. + + Args: + policy: The Policy to generate stats for. + train_batch: The SampleBatch (already) used for training. + + Returns: + Dict[str, TensorType]: The stats dict. + """ + q_t = torch.stack(policy.get_tower_stats("q_t")) + + return { + "actor_loss": torch.mean(torch.stack(policy.get_tower_stats("actor_loss"))), + "critic_loss": torch.mean( + torch.stack(tree.flatten(policy.get_tower_stats("critic_loss"))) + ), + "alpha_loss": torch.mean(torch.stack(policy.get_tower_stats("alpha_loss"))), + "alpha_value": torch.exp(policy.model.log_alpha), + "log_alpha_value": policy.model.log_alpha, + "target_entropy": policy.model.target_entropy, + "policy_t": torch.mean(torch.stack(policy.get_tower_stats("policy_t"))), + "mean_q": torch.mean(q_t), + "max_q": torch.max(q_t), + "min_q": torch.min(q_t), + } + + +def optimizer_fn(policy: Policy, config: AlgorithmConfigDict) -> Tuple[LocalOptimizer]: + """Creates all necessary optimizers for SAC learning. + + The 3 or 4 (twin_q=True) optimizers returned here correspond to the + number of loss terms returned by the loss function. + + Args: + policy: The policy object to be trained. + config: The Algorithm's config dict. + + Returns: + Tuple[LocalOptimizer]: The local optimizers to use for policy training. + """ + policy.actor_optim = torch.optim.Adam( + params=policy.model.policy_variables(), + lr=config["optimization"]["actor_learning_rate"], + eps=1e-7, # to match tf.keras.optimizers.Adam's epsilon default + ) + + critic_split = len(policy.model.q_variables()) + if config["twin_q"]: + critic_split //= 2 + + policy.critic_optims = [ + torch.optim.Adam( + params=policy.model.q_variables()[:critic_split], + lr=config["optimization"]["critic_learning_rate"], + eps=1e-7, # to match tf.keras.optimizers.Adam's epsilon default + ) + ] + if config["twin_q"]: + policy.critic_optims.append( + torch.optim.Adam( + params=policy.model.q_variables()[critic_split:], + lr=config["optimization"]["critic_learning_rate"], + eps=1e-7, # to match tf.keras.optimizers.Adam's eps default + ) + ) + policy.alpha_optim = torch.optim.Adam( + params=[policy.model.log_alpha], + lr=config["optimization"]["entropy_learning_rate"], + eps=1e-7, # to match tf.keras.optimizers.Adam's epsilon default + ) + + return tuple([policy.actor_optim] + policy.critic_optims + [policy.alpha_optim]) + + +# TODO: Unify with DDPG's ComputeTDErrorMixin when SAC policy subclasses PolicyV2 +class ComputeTDErrorMixin: + """Mixin class calculating TD-error (part of critic loss) per batch item. + + - Adds `policy.compute_td_error()` method for TD-error calculation from a + batch of observations/actions/rewards/etc.. + """ + + def __init__(self): + def compute_td_error( + obs_t, act_t, rew_t, obs_tp1, terminateds_mask, importance_weights + ): + input_dict = self._lazy_tensor_dict( + { + SampleBatch.CUR_OBS: obs_t, + SampleBatch.ACTIONS: act_t, + SampleBatch.REWARDS: rew_t, + SampleBatch.NEXT_OBS: obs_tp1, + SampleBatch.TERMINATEDS: terminateds_mask, + PRIO_WEIGHTS: importance_weights, + } + ) + # Do forward pass on loss to update td errors attribute + # (one TD-error value per item in batch to update PR weights). + actor_critic_loss(self, self.model, None, input_dict) + + # `self.model.td_error` is set within actor_critic_loss call. + # Return its updated value here. + return self.model.tower_stats["td_error"] + + # Assign the method to policy (self) for later usage. + self.compute_td_error = compute_td_error + + +def setup_late_mixins( + policy: Policy, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: AlgorithmConfigDict, +) -> None: + """Call mixin classes' constructors after Policy initialization. + + - Moves the target model(s) to the GPU, if necessary. + - Adds the `compute_td_error` method to the given policy. + Calling `compute_td_error` with batch data will re-calculate the loss + on that batch AND return the per-batch-item TD-error for prioritized + replay buffer record weight updating (in case a prioritized replay buffer + is used). + - Also adds the `update_target` method to the given policy. + Calling `update_target` updates all target Q-networks' weights from their + respective "main" Q-metworks, based on tau (smooth, partial updating). + + Args: + policy: The Policy object. + obs_space (gym.spaces.Space): The Policy's observation space. + action_space (gym.spaces.Space): The Policy's action space. + config: The Policy's config. + """ + ComputeTDErrorMixin.__init__(policy) + TargetNetworkMixin.__init__(policy) + + +# Build a child class of `TorchPolicy`, given the custom functions defined +# above. +SACTorchPolicy = build_policy_class( + name="SACTorchPolicy", + framework="torch", + loss_fn=actor_critic_loss, + get_default_config=lambda: ray.rllib.algorithms.sac.sac.SACConfig(), + stats_fn=stats, + postprocess_fn=postprocess_trajectory, + extra_grad_process_fn=apply_grad_clipping, + optimizer_fn=optimizer_fn, + validate_spaces=validate_spaces, + before_loss_init=setup_late_mixins, + make_model_and_action_dist=build_sac_model_and_action_dist, + extra_learn_fetches_fn=concat_multi_gpu_td_errors, + mixins=[TargetNetworkMixin, ComputeTDErrorMixin], + action_distribution_fn=action_distribution_fn, +) diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32179c31ba4217180fc1e2a8e5c58984ae589b2a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/default_sac_torch_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/default_sac_torch_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eee37780573a5a03fbf866def149439a5fbc3d8f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/default_sac_torch_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/sac_torch_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/sac_torch_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac919be444935bb3d48323ae4ff6a16d4fc57f24 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/__pycache__/sac_torch_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/default_sac_torch_rl_module.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/default_sac_torch_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..09a0f6091ab1dc4d6bd6382aa2f7bb7d69993fc5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/default_sac_torch_rl_module.py @@ -0,0 +1,280 @@ +from typing import Any, Dict + +import gymnasium as gym + +from ray.rllib.algorithms.sac.default_sac_rl_module import DefaultSACRLModule +from ray.rllib.algorithms.sac.sac_catalog import SACCatalog +from ray.rllib.algorithms.sac.sac_learner import ( + ACTION_DIST_INPUTS_NEXT, + ACTION_LOG_PROBS, + ACTION_LOG_PROBS_NEXT, + ACTION_PROBS, + ACTION_PROBS_NEXT, + QF_PREDS, + QF_TARGET_NEXT, + QF_TWIN_PREDS, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.models.base import ENCODER_OUT, Encoder, Model +from ray.rllib.core.rl_module.apis import QNetAPI, TargetNetworkAPI +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.util.annotations import DeveloperAPI + +torch, nn = try_import_torch() + + +@DeveloperAPI +class DefaultSACTorchRLModule(TorchRLModule, DefaultSACRLModule): + framework: str = "torch" + + def __init__(self, *args, **kwargs): + catalog_class = kwargs.pop("catalog_class", None) + if catalog_class is None: + catalog_class = SACCatalog + super().__init__(*args, **kwargs, catalog_class=catalog_class) + + @override(RLModule) + def _forward_inference(self, batch: Dict) -> Dict[str, Any]: + output = {} + + # Pi encoder forward pass. + pi_encoder_outs = self.pi_encoder(batch) + + # Pi head. + # assume action space is either discrete or continuous. + output[Columns.ACTION_DIST_INPUTS] = self.pi(pi_encoder_outs[ENCODER_OUT]) + + return output + + @override(RLModule) + def _forward_exploration(self, batch: Dict, **kwargs) -> Dict[str, Any]: + return self._forward_inference(batch) + + @override(RLModule) + def _forward_train(self, batch: Dict) -> Dict[str, Any]: + # Call the `super`'s `forward_train` + super()._forward_train(batch) + if isinstance(self.action_space, gym.spaces.Discrete): + return self._forward_train_discrete(batch) + elif isinstance(self.action_space, gym.spaces.Box): + return self._forward_train_continuous(batch) + else: + raise ValueError( + f"Unsupported action space type: {type(self.action_space)}. " + "Only discrete and continuous action spaces are supported." + ) + + def _forward_train_discrete(self, batch: Dict[str, Any]) -> Dict[str, Any]: + output = {} + + # SAC needs also Q function values and action logits for next observations. + batch_curr = {Columns.OBS: batch[Columns.OBS]} + batch_next = {Columns.OBS: batch[Columns.NEXT_OBS]} + + ## calculate values for the Q target ## + # Also encode the next observations (and next actions for the Q net). + pi_encoder_next_outs = self.pi_encoder(batch_next) + action_logits_next = self.pi(pi_encoder_next_outs[ENCODER_OUT]) + # TODO(inyoung): get the action dist class and use that. But currently TorchCategorical + # does not get the prob value of the actual torch distribution. So we use softmax directly + # for now. + action_probs_next = torch.nn.functional.softmax(action_logits_next, dim=-1) + + output[ACTION_PROBS_NEXT] = action_probs_next + output[ACTION_LOG_PROBS_NEXT] = action_probs_next.log() + + # (B, action_dim) + qf_target_next = self.forward_target(batch_next, squeeze=False) + output[QF_TARGET_NEXT] = qf_target_next + + qf_preds = self._qf_forward_train_helper( + batch_curr, self.qf_encoder, self.qf, squeeze=False + ) + # we don't need straight-through gradient here + output[QF_PREDS] = qf_preds + if self.twin_q: + qf_twin_preds = self._qf_forward_train_helper( + batch_curr, self.qf_twin_encoder, self.qf_twin, squeeze=False + ) + output[QF_TWIN_PREDS] = qf_twin_preds + + ## calculate values for gradient ## + pi_encoder_outs = self.pi_encoder(batch_curr) + action_logits = self.pi(pi_encoder_outs[ENCODER_OUT]) + action_probs = torch.nn.functional.softmax(action_logits, dim=-1) + output[ACTION_PROBS] = action_probs + output[ACTION_LOG_PROBS] = action_probs.log() + + return output + + def _forward_train_continuous(self, batch: Dict[str, Any]) -> Dict[str, Any]: + output = {} + + # SAC needs also Q function values and action logits for next observations. + batch_curr = {Columns.OBS: batch[Columns.OBS]} + batch_next = {Columns.OBS: batch[Columns.NEXT_OBS]} + + # Encoder forward passes. + pi_encoder_outs = self.pi_encoder(batch_curr) + + # Also encode the next observations (and next actions for the Q net). + pi_encoder_next_outs = self.pi_encoder(batch_next) + + # Q-network(s) forward passes. + batch_curr.update({Columns.ACTIONS: batch[Columns.ACTIONS]}) + output[QF_PREDS] = self._qf_forward_train_helper( + batch_curr, self.qf_encoder, self.qf + ) # self._qf_forward_train(batch_curr)[QF_PREDS] + # If necessary make a forward pass through the twin Q network. + if self.twin_q: + output[QF_TWIN_PREDS] = self._qf_forward_train_helper( + batch_curr, self.qf_twin_encoder, self.qf_twin + ) + + # Policy head. + action_logits = self.pi(pi_encoder_outs[ENCODER_OUT]) + # Also get the action logits for the next observations. + action_logits_next = self.pi(pi_encoder_next_outs[ENCODER_OUT]) + output[Columns.ACTION_DIST_INPUTS] = action_logits + output[ACTION_DIST_INPUTS_NEXT] = action_logits_next + + # Get the train action distribution for the current policy and current state. + # This is needed for the policy (actor) loss in SAC. + action_dist_class = self.get_train_action_dist_cls() + action_dist_curr = action_dist_class.from_logits(action_logits) + # Get the train action distribution for the current policy and next state. + # For the Q (critic) loss in SAC, we need to sample from the current policy at + # the next state. + action_dist_next = action_dist_class.from_logits(action_logits_next) + + # Sample actions for the current state. Note that we need to apply the + # reparameterization trick (`rsample()` instead of `sample()`) to avoid the + # expectation over actions. + actions_resampled = action_dist_curr.rsample() + # Compute the log probabilities for the current state (for the critic loss). + output["logp_resampled"] = action_dist_curr.logp(actions_resampled) + + # Sample actions for the next state. + actions_next_resampled = action_dist_next.sample().detach() + # Compute the log probabilities for the next state. + output["logp_next_resampled"] = ( + action_dist_next.logp(actions_next_resampled) + ).detach() + + # Compute Q-values for the current policy in the current state with + # the sampled actions. + q_batch_curr = { + Columns.OBS: batch[Columns.OBS], + Columns.ACTIONS: actions_resampled, + } + # Make sure we perform a "straight-through gradient" pass here, + # ignoring the gradients of the q-net, however, still recording + # the gradients of the policy net (which was used to rsample the actions used + # here). This is different from doing `.detach()` or `with torch.no_grads()`, + # as these two methds would fully block all gradient recordings, including + # the needed policy ones. + all_params = list(self.qf.parameters()) + list(self.qf_encoder.parameters()) + if self.twin_q: + all_params += list(self.qf_twin.parameters()) + list( + self.qf_twin_encoder.parameters() + ) + + for param in all_params: + param.requires_grad = False + output["q_curr"] = self.compute_q_values(q_batch_curr) + for param in all_params: + param.requires_grad = True + + # Compute Q-values from the target Q network for the next state with the + # sampled actions for the next state. + q_batch_next = { + Columns.OBS: batch[Columns.NEXT_OBS], + Columns.ACTIONS: actions_next_resampled, + } + output["q_target_next"] = self.forward_target(q_batch_next).detach() + + # Return the network outputs. + return output + + @override(TargetNetworkAPI) + def forward_target( + self, batch: Dict[str, Any], squeeze: bool = True + ) -> Dict[str, Any]: + target_qvs = self._qf_forward_train_helper( + batch, self.target_qf_encoder, self.target_qf, squeeze=squeeze + ) + + # If a twin Q network should be used, calculate twin Q-values and use the + # minimum. + if self.twin_q: + target_qvs = torch.min( + target_qvs, + self._qf_forward_train_helper( + batch, + self.target_qf_twin_encoder, + self.target_qf_twin, + squeeze=squeeze, + ), + ) + + return target_qvs + + @override(QNetAPI) + def compute_q_values( + self, batch: Dict[str, Any], squeeze: bool = True + ) -> Dict[str, Any]: + qvs = self._qf_forward_train_helper( + batch, self.qf_encoder, self.qf, squeeze=squeeze + ) + # If a twin Q network should be used, calculate twin Q-values and use the + # minimum. + if self.twin_q: + qvs = torch.min( + qvs, + self._qf_forward_train_helper( + batch, self.qf_twin_encoder, self.qf_twin, squeeze=squeeze + ), + ) + return qvs + + @override(DefaultSACRLModule) + def _qf_forward_train_helper( + self, batch: Dict[str, Any], encoder: Encoder, head: Model, squeeze: bool = True + ) -> Dict[str, Any]: + """Executes the forward pass for Q networks. + + Args: + batch: Dict containing a concatenated tensor with observations + and actions under the key `Columns.OBS`. + encoder: An `Encoder` model for the Q state-action encoder. + head: A `Model` for the Q head. + squeeze: If True, squeezes the last dimension of the output if it is 1. Used for continuous action spaces. + + Returns: + The estimated Q-value for the input action for continuous action spaces. + Or the Q-values for all actions for discrete action spaces. + """ + # Construct batch. Note, we need to feed observations and actions. + if isinstance(self.action_space, gym.spaces.Box): + actions = batch[Columns.ACTIONS] + qf_batch = { + Columns.OBS: torch.concat((batch[Columns.OBS], actions), dim=-1) + } + else: + # For discrete action spaces, we don't need to include the actions + # in the batch, as the Q function outputs the Q-values for each action + qf_batch = {Columns.OBS: batch[Columns.OBS]} + + # Encoder forward pass. + qf_encoder_outs = encoder(qf_batch) + + # Q head forward pass. + # (B,latent_size) -> (B, 1|action_dim) + qf_out = head(qf_encoder_outs[ENCODER_OUT]) + if squeeze: + # Squeeze the last dimension if it is 1. + qf_out = qf_out.squeeze(-1) + return qf_out diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/sac_torch_learner.py b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/sac_torch_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..478970795d853ae9282e2ed8b22d30b3eb81e163 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/sac/torch/sac_torch_learner.py @@ -0,0 +1,430 @@ +from typing import Any, Dict + +import gymnasium as gym + +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.algorithms.dqn.torch.dqn_torch_learner import DQNTorchLearner +from ray.rllib.algorithms.sac.sac import SACConfig +from ray.rllib.algorithms.sac.sac_learner import ( + ACTION_LOG_PROBS, + ACTION_LOG_PROBS_NEXT, + ACTION_PROBS, + ACTION_PROBS_NEXT, + LOGPS_KEY, + QF_LOSS_KEY, + QF_MAX_KEY, + QF_MEAN_KEY, + QF_MIN_KEY, + QF_PREDS, + QF_TARGET_NEXT, + QF_TWIN_LOSS_KEY, + QF_TWIN_PREDS, + TD_ERROR_MEAN_KEY, + SACLearner, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.learner import ( + POLICY_LOSS_KEY, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ALL_MODULES, TD_ERROR_KEY +from ray.rllib.utils.typing import ModuleID, ParamDict, TensorType + +torch, nn = try_import_torch() + + +class SACTorchLearner(DQNTorchLearner, SACLearner): + """Implements `torch`-specific SAC loss logic on top of `SACLearner` + + This ' Learner' class implements the loss in its + `self.compute_loss_for_module()` method. In addition, it updates + the target networks of the RLModule(s). + """ + + def build(self) -> None: + super().build() + + # Store loss tensors here temporarily inside the loss function for (exact) + # consumption later by the compute gradients function. + # Keys=(module_id, optimizer_name), values=loss tensors (in-graph). + self._temp_losses = {} + + # TODO (simon): Set different learning rates for optimizers. + @override(DQNTorchLearner) + def configure_optimizers_for_module( + self, module_id: ModuleID, config: AlgorithmConfig = None + ) -> None: + # Receive the module. + module = self._module[module_id] + + # Define the optimizer for the critic. + # TODO (sven): Maybe we change here naming to `qf` for unification. + params_critic = self.get_parameters(module.qf_encoder) + self.get_parameters( + module.qf + ) + optim_critic = torch.optim.Adam(params_critic, eps=1e-7) + + self.register_optimizer( + module_id=module_id, + optimizer_name="qf", + optimizer=optim_critic, + params=params_critic, + lr_or_lr_schedule=config.critic_lr, + ) + # If necessary register also an optimizer for a twin Q network. + if config.twin_q: + params_twin_critic = self.get_parameters( + module.qf_twin_encoder + ) + self.get_parameters(module.qf_twin) + optim_twin_critic = torch.optim.Adam(params_twin_critic, eps=1e-7) + + self.register_optimizer( + module_id=module_id, + optimizer_name="qf_twin", + optimizer=optim_twin_critic, + params=params_twin_critic, + lr_or_lr_schedule=config.critic_lr, + ) + + # Define the optimizer for the actor. + params_actor = self.get_parameters(module.pi_encoder) + self.get_parameters( + module.pi + ) + optim_actor = torch.optim.Adam(params_actor, eps=1e-7) + + self.register_optimizer( + module_id=module_id, + optimizer_name="policy", + optimizer=optim_actor, + params=params_actor, + lr_or_lr_schedule=config.actor_lr, + ) + + # Define the optimizer for the temperature. + temperature = self.curr_log_alpha[module_id] + optim_temperature = torch.optim.Adam([temperature], eps=1e-7) + self.register_optimizer( + module_id=module_id, + optimizer_name="alpha", + optimizer=optim_temperature, + params=[temperature], + lr_or_lr_schedule=config.alpha_lr, + ) + + @override(DQNTorchLearner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: SACConfig, + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + + module = self._module[module_id] + if isinstance(module.action_space, gym.spaces.Discrete): + # Discrete action space: Use the discrete loss function. + return self._compute_loss_for_module_discrete( + module_id=module_id, + config=config, + batch=batch, + fwd_out=fwd_out, + ) + elif isinstance(module.action_space, gym.spaces.Box): + # Continuous action space: Use the continuous loss function. + return self._compute_loss_for_module_continuous( + module_id=module_id, + config=config, + batch=batch, + fwd_out=fwd_out, + ) + else: + raise ValueError( + f"Unsupported action space type: {type(module.action_space)}. " + "Only Discrete and Box action spaces are supported." + ) + + def _compute_loss_for_module_discrete( + self, + *, + module_id: ModuleID, + config: SACConfig, + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + # Receive the current alpha hyperparameter. + alpha = torch.exp(self.curr_log_alpha[module_id]) + + ## Calculate Q value targets + action_probs_next = fwd_out[ACTION_PROBS_NEXT] + action_log_probs_next = fwd_out[ACTION_LOG_PROBS_NEXT] + next_q = fwd_out[QF_TARGET_NEXT] + next_v = ( + (action_probs_next * (next_q - alpha.detach() * action_log_probs_next)) + .sum(-1) + .squeeze(-1) + ) + next_v_masked = (1.0 - batch[Columns.TERMINATEDS].float()) * next_v + target_q = ( + batch[Columns.REWARDS] + (config.gamma ** batch["n_step"]) * next_v_masked + ).detach() + + # Get Q-values for the actually selected actions during rollout. + actions = batch[Columns.ACTIONS].to(dtype=torch.int64).unsqueeze(-1) + qf_pred = fwd_out[QF_PREDS].gather(dim=-1, index=actions).squeeze(-1) + if config.twin_q: + qf_twin_pred = ( + fwd_out[QF_TWIN_PREDS].gather(dim=-1, index=actions).squeeze(-1) + ) + + # Calculate the TD-error. Note, this is needed for the priority weights in + # the replay buffer. + td_error = torch.abs(qf_pred - target_q) + # If a twin Q network should be used, add the TD error of the twin Q network. + if config.twin_q: + td_error += torch.abs(qf_twin_pred - target_q) + # Rescale the TD error. + td_error *= 0.5 + + # MSBE loss for the critic(s) (i.e. Q, see eqs. (7-8) Haarnoja et al. (2018)). + # Note, this needs a sample from the current policy given the next state. + # Note further, we use here the Huber loss instead of the mean squared error + # as it improves training performance. + critic_loss = torch.mean( + batch["weights"] + * torch.nn.HuberLoss(reduction="none", delta=1.0)(qf_pred, target_q) + ) + # If a twin Q network should be used, add the critic loss of the twin Q network. + if config.twin_q: + critic_twin_loss = torch.mean( + batch["weights"] + * torch.nn.HuberLoss(reduction="none", delta=1.0)( + qf_twin_pred, target_q + ) + ) + + ## Calculate the actor loss ## + action_probs = fwd_out[ACTION_PROBS] + action_log_probs = fwd_out[ACTION_LOG_PROBS] + qf = torch.min(fwd_out[QF_PREDS], fwd_out[QF_TWIN_PREDS]).detach() + policy_loss = ( + (action_probs * (alpha.detach() * action_log_probs - qf)).sum(-1).mean() + ) + + ## Calculate the alpha loss ## + entropy = (action_log_probs * action_probs).sum(-1) + alpha_loss = -torch.mean( + self.curr_log_alpha[module_id] + * (entropy.detach() + self.target_entropy[module_id]) + ) + + total_loss = policy_loss + critic_loss + alpha_loss + if config.twin_q: + total_loss += critic_twin_loss + + # Log the TD-error with reduce=None, such that - in case we have n parallel + # Learners - we will re-concatenate the produced TD-error tensors to yield + # a 1:1 representation of the original batch. + self.metrics.log_value( + key=(module_id, TD_ERROR_KEY), + value=td_error, + reduce=None, + clear_on_reduce=True, + ) + # Log other important loss stats (reduce=mean (default), but with window=1 + # in order to keep them history free). + self.metrics.log_dict( + { + POLICY_LOSS_KEY: policy_loss, + QF_LOSS_KEY: critic_loss, + "alpha_loss": alpha_loss, + "alpha_value": alpha[0], + "log_alpha_value": torch.log(alpha)[0], + "target_entropy": self.target_entropy[module_id], + LOGPS_KEY: torch.mean(fwd_out[ACTION_LOG_PROBS]), + QF_MEAN_KEY: torch.mean(fwd_out[QF_PREDS]), + QF_MAX_KEY: torch.max(fwd_out[QF_PREDS]), + QF_MIN_KEY: torch.min(fwd_out[QF_PREDS]), + TD_ERROR_MEAN_KEY: torch.mean(td_error), + }, + key=module_id, + window=1, # <- single items (should not be mean/ema-reduced over time). + ) + + self._temp_losses[(module_id, POLICY_LOSS_KEY)] = policy_loss + self._temp_losses[(module_id, QF_LOSS_KEY)] = critic_loss + self._temp_losses[(module_id, "alpha_loss")] = alpha_loss + + # If twin Q networks should be used add a critic loss for the twin Q network. + # Note, we need this in the `self.compute_gradients()` to optimize. + if config.twin_q: + self.metrics.log_value( + key=(module_id, QF_TWIN_LOSS_KEY), + value=critic_twin_loss, + window=1, # <- single items (should not be mean/ema-reduced over time). + ) + self._temp_losses[(module_id, QF_TWIN_LOSS_KEY)] = critic_twin_loss + + return total_loss + + def _compute_loss_for_module_continuous( + self, + *, + module_id: ModuleID, + config: SACConfig, + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + # Receive the current alpha hyperparameter. + alpha = torch.exp(self.curr_log_alpha[module_id]) + + # Get Q-values for the actually selected actions during rollout. + # In the critic loss we use these as predictions. + q_selected = fwd_out[QF_PREDS] + if config.twin_q: + q_twin_selected = fwd_out[QF_TWIN_PREDS] + + # Compute value function for next state (see eq. (3) in Haarnoja et al. (2018)). + # Note, we use here the sampled actions in the log probabilities. + q_target_next = ( + fwd_out["q_target_next"] - alpha.detach() * fwd_out["logp_next_resampled"] + ) + # Now mask all Q-values with terminated next states in the targets. + q_next_masked = (1.0 - batch[Columns.TERMINATEDS].float()) * q_target_next + + # Compute the right hand side of the Bellman equation. + # Detach this node from the computation graph as we do not want to + # backpropagate through the target network when optimizing the Q loss. + q_selected_target = ( + batch[Columns.REWARDS] + (config.gamma ** batch["n_step"]) * q_next_masked + ).detach() + + # Calculate the TD-error. Note, this is needed for the priority weights in + # the replay buffer. + td_error = torch.abs(q_selected - q_selected_target) + # If a twin Q network should be used, add the TD error of the twin Q network. + if config.twin_q: + td_error += torch.abs(q_twin_selected - q_selected_target) + # Rescale the TD error. + td_error *= 0.5 + + # MSBE loss for the critic(s) (i.e. Q, see eqs. (7-8) Haarnoja et al. (2018)). + # Note, this needs a sample from the current policy given the next state. + # Note further, we use here the Huber loss instead of the mean squared error + # as it improves training performance. + critic_loss = torch.mean( + batch["weights"] + * torch.nn.HuberLoss(reduction="none", delta=1.0)( + q_selected, q_selected_target + ) + ) + # If a twin Q network should be used, add the critic loss of the twin Q network. + if config.twin_q: + critic_twin_loss = torch.mean( + batch["weights"] + * torch.nn.HuberLoss(reduction="none", delta=1.0)( + q_twin_selected, q_selected_target + ) + ) + + # For the actor (policy) loss we need sampled actions from the current policy + # evaluated at the current observations. + # Note that the `q_curr` tensor below has the q-net's gradients ignored, while + # having the policy's gradients registered. The policy net was used to rsample + # actions used to compute `q_curr` (by passing these actions through the q-net). + # Hence, we can't do `fwd_out[q_curr].detach()`! + # Note further, we minimize here, while the original equation in Haarnoja et + # al. (2018) considers maximization. + # TODO (simon): Rename to `resampled` to `current`. + actor_loss = torch.mean( + alpha.detach() * fwd_out["logp_resampled"] - fwd_out["q_curr"] + ) + + # Optimize also the hyperparameter alpha by using the current policy + # evaluated at the current state (sampled values). + # TODO (simon): Check, why log(alpha) is used, prob. just better + # to optimize and monotonic function. Original equation uses alpha. + alpha_loss = -torch.mean( + self.curr_log_alpha[module_id] + * (fwd_out["logp_resampled"].detach() + self.target_entropy[module_id]) + ) + + total_loss = actor_loss + critic_loss + alpha_loss + # If twin Q networks should be used, add the critic loss of the twin Q network. + if config.twin_q: + # TODO (simon): Check, if we need to multiply the critic_loss then with 0.5. + total_loss += critic_twin_loss + + # Log the TD-error with reduce=None, such that - in case we have n parallel + # Learners - we will re-concatenate the produced TD-error tensors to yield + # a 1:1 representation of the original batch. + self.metrics.log_value( + key=(module_id, TD_ERROR_KEY), + value=td_error, + reduce=None, + clear_on_reduce=True, + ) + # Log other important loss stats (reduce=mean (default), but with window=1 + # in order to keep them history free). + self.metrics.log_dict( + { + POLICY_LOSS_KEY: actor_loss, + QF_LOSS_KEY: critic_loss, + "alpha_loss": alpha_loss, + "alpha_value": alpha[0], + "log_alpha_value": torch.log(alpha)[0], + "target_entropy": self.target_entropy[module_id], + LOGPS_KEY: torch.mean(fwd_out["logp_resampled"]), + QF_MEAN_KEY: torch.mean(fwd_out["q_curr"]), + QF_MAX_KEY: torch.max(fwd_out["q_curr"]), + QF_MIN_KEY: torch.min(fwd_out["q_curr"]), + TD_ERROR_MEAN_KEY: torch.mean(td_error), + }, + key=module_id, + window=1, # <- single items (should not be mean/ema-reduced over time). + ) + + self._temp_losses[(module_id, POLICY_LOSS_KEY)] = actor_loss + self._temp_losses[(module_id, QF_LOSS_KEY)] = critic_loss + self._temp_losses[(module_id, "alpha_loss")] = alpha_loss + + # If twin Q networks should be used add a critic loss for the twin Q network. + # Note, we need this in the `self.compute_gradients()` to optimize. + if config.twin_q: + self.metrics.log_value( + key=(module_id, QF_TWIN_LOSS_KEY), + value=critic_twin_loss, + window=1, # <- single items (should not be mean/ema-reduced over time). + ) + self._temp_losses[(module_id, QF_TWIN_LOSS_KEY)] = critic_twin_loss + + return total_loss + + @override(DQNTorchLearner) + def compute_gradients( + self, loss_per_module: Dict[ModuleID, TensorType], **kwargs + ) -> ParamDict: + grads = {} + for module_id in set(loss_per_module.keys()) - {ALL_MODULES}: + # Loop through optimizers registered for this module. + for optim_name, optim in self.get_optimizers_for_module(module_id): + # Zero the gradients. Note, we need to reset the gradients b/c + # each component for a module operates on the same graph. + optim.zero_grad(set_to_none=True) + + # Compute the gradients for the component and module. + loss_tensor = self._temp_losses.pop((module_id, optim_name + "_loss")) + loss_tensor.backward(retain_graph=True) + # Store the gradients for the component and module. + grads.update( + { + pid: p.grad + for pid, p in self.filter_param_dict_for_optimizer( + self._params, optim + ).items() + } + ) + + assert not self._temp_losses + return grads diff --git a/lib/python3.12/site-packages/ray/rllib/algorithms/utils.py b/lib/python3.12/site-packages/ray/rllib/algorithms/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..25f9644512c22efb356c5a8f8cc3d07d9a967e8f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/algorithms/utils.py @@ -0,0 +1,227 @@ +import platform +from typing import List + +import tree # pip install dm_tree + +import ray +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.actor_manager import FaultAwareApply +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.metrics.ray_metrics import ( + DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + TimerAndPrometheusLogger, +) +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import DeveloperAPI +from ray.util.metrics import Counter, Histogram + +torch, _ = try_import_torch() + + +@DeveloperAPI(stability="alpha") +class AggregatorActor(FaultAwareApply): + """Runs episode lists through ConnectorV2 pipeline and creates train batches. + + The actor should be co-located with a Learner worker. Ideally, there should be one + or two aggregator actors per Learner worker (having even more per Learner probably + won't help. Then the main process driving the RL algo can perform the following + execution logic: + - query n EnvRunners to sample the environment and return n lists of episodes as + Ray.ObjectRefs. + - remote call the set of aggregator actors (in round-robin fashion) with these + list[episodes] refs in async fashion. + - gather the results asynchronously, as each actor returns refs pointing to + ready-to-go train batches. + - as soon as we have at least one train batch per Learner, call the LearnerGroup + with the (already sharded) refs. + - an aggregator actor - when receiving p refs to List[EpisodeType] - does: + -- ray.get() the actual p lists and concatenate the p lists into one + List[EpisodeType]. + -- pass the lists of episodes through its LearnerConnector pipeline + -- buffer the output batches of this pipeline until enough batches have been + collected for creating one train batch (matching the config's + `train_batch_size_per_learner`). + -- concatenate q batches into a train batch and return that train batch. + - the algo main process then passes the ray.ObjectRef to the ready-to-go train batch + to the LearnerGroup for calling each Learner with one train batch. + """ + + def __init__(self, config: AlgorithmConfig, rl_module_spec): + self.config = config + + # Set device and node. + self._node = platform.node() + self._device = torch.device("cpu") + self.metrics = MetricsLogger() + + # Create the RLModule. + # TODO (sven): For now, this RLModule (its weights) never gets updated. + # The reason the module is needed is for the connector to know, which + # sub-modules are stateful (and what their initial state tensors are), and + # which IDs the submodules have (to figure out, whether its multi-agent or + # not). + self._module = rl_module_spec.build() + self._module = self._module.as_multi_rl_module() + + # Create the Learner connector pipeline. + self._learner_connector = self.config.build_learner_connector( + input_observation_space=None, + input_action_space=None, + device=self._device, + ) + + # Ray metrics + self._metrics_get_batch_time = Histogram( + name="rllib_utils_aggregator_actor_get_batch_time", + description="Time spent in AggregatorActor.get_batch()", + boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS, + tag_keys=("rllib",), + ) + self._metrics_get_batch_time.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_episode_owner_died = Counter( + name="rllib_utils_aggregator_actor_episode_owner_died_counter", + description="N times ray.get() on an episode ref failed ", + tag_keys=("rllib",), + ) + self._metrics_episode_owner_died.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_get_batch_input_episode_refs = Counter( + name="rllib_utils_aggregator_actor_get_batch_input_episode_refs_counter", + description="Number of episode refs received as input to get_batch()", + tag_keys=("rllib",), + ) + self._metrics_get_batch_input_episode_refs.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_get_batch_output_batches = Counter( + name="rllib_utils_aggregator_actor_get_batch_output_batches_counter", + description="Number of policy batches output by get_batch()", + tag_keys=("rllib",), + ) + self._metrics_get_batch_output_batches.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + def get_batch(self, episode_refs: List[ray.ObjectRef]): + with TimerAndPrometheusLogger(self._metrics_get_batch_time): + if len(episode_refs) > 0: + self._metrics_get_batch_input_episode_refs.inc(value=len(episode_refs)) + + episodes: List[EpisodeType] = [] + # It's possible that individual refs are invalid due to the EnvRunner + # that produced the ref has crashed or had its entire node go down. + # In this case, try each ref individually and collect only valid results. + try: + episodes = tree.flatten(ray.get(episode_refs)) + except ray.exceptions.OwnerDiedError: + for ref in episode_refs: + try: + episodes.extend(ray.get(ref)) + except ray.exceptions.OwnerDiedError: + self._metrics_episode_owner_died.inc(value=1) + + env_steps = sum(len(e) for e in episodes) + + # If we have enough episodes collected to create a single train batch, pass + # them at once through the connector to receive a single train batch. + batch = self._learner_connector( + episodes=episodes, + rl_module=self._module, + metrics=self.metrics, + ) + # Convert to a dict into a `MultiAgentBatch`. + # TODO (sven): Try to get rid of dependency on MultiAgentBatch (once our mini- + # batch iterators support splitting over a dict). + ma_batch = MultiAgentBatch( + policy_batches={ + pid: SampleBatch(pol_batch) for pid, pol_batch in batch.items() + }, + env_steps=env_steps, + ) + self._metrics_get_batch_output_batches.inc(value=1) + return ma_batch + + def get_metrics(self): + return self.metrics.reduce() + + +def _get_env_runner_bundles(config): + return [ + { + "CPU": config.num_cpus_per_env_runner, + "GPU": config.num_gpus_per_env_runner, + **config.custom_resources_per_env_runner, + } + for _ in range(config.num_env_runners) + ] + + +def _get_offline_eval_runner_bundles(config): + return [ + { + "CPU": config.num_cpus_per_offline_eval_runner, + "GPU": config.num_gpus_per_offline_eval_runner, + **config.custom_resources_per_offline_eval_runner, + } + for _ in range(config.num_offline_eval_runners) + ] + + +def _get_learner_bundles(config): + try: + from ray.rllib.extensions.algorithm_utils import _get_learner_bundles as func + + return func(config) + except Exception: + pass + + if config.num_learners == 0: + if config.num_aggregator_actors_per_learner > 0: + return [{"CPU": config.num_aggregator_actors_per_learner}] + else: + return [] + + num_cpus_per_learner = ( + config.num_cpus_per_learner + if config.num_cpus_per_learner != "auto" + else 1 + if config.num_gpus_per_learner == 0 + else 0 + ) + + bundles = [ + { + "CPU": config.num_learners + * (num_cpus_per_learner + config.num_aggregator_actors_per_learner), + "GPU": config.num_learners * config.num_gpus_per_learner, + } + ] + + return bundles + + +def _get_main_process_bundle(config): + if config.num_learners == 0: + num_cpus_per_learner = ( + config.num_cpus_per_learner + if config.num_cpus_per_learner != "auto" + else 1 + if config.num_gpus_per_learner == 0 + else 0 + ) + bundle = { + "CPU": max(num_cpus_per_learner, config.num_cpus_for_main_process), + "GPU": config.num_gpus_per_learner, + } + else: + bundle = {"CPU": config.num_cpus_for_main_process, "GPU": 0} + + return bundle diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/__init__.py b/lib/python3.12/site-packages/ray/rllib/connectors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8aae9a9d2ae1155d9f37729295736dc3869adbce Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b66d4a5252cd40f4f74117864ffdfb59a879ca7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector_pipeline_v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector_pipeline_v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b428566699c54e00f8bd228a3b599c489e7c6bf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector_pipeline_v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector_v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector_v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7123ec98a1d2a69cfa593d2f54febd7b5852c57 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/connector_v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/registry.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/registry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea21b33c83d99380279c13fb8eedde7cef6be04e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/registry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab9ba1e3e524a2047f69298f9023d9b9c33b361c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/__init__.py b/lib/python3.12/site-packages/ray/rllib/connectors/action/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b55a1da43f8324c272df2927da3973b1d093671 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/clip.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/clip.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbaddcde582d3653afc53bf9a4b8784d0a3c8aee Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/clip.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/immutable.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/immutable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..565652c38946175139682708d10b66ac73207715 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/immutable.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/lambdas.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/lambdas.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dec99458e79a45a9a0b2da7d974ae466ccac3685 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/lambdas.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/normalize.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/normalize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ecf95e1f3d78efb1ef00c4d315f5598c26c895fe Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/normalize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/pipeline.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70fa0b58da1a4a8ee18c77f3e3e3597eb273600b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/action/__pycache__/pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/clip.py b/lib/python3.12/site-packages/ray/rllib/connectors/action/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..4450ea4442381d0b9f2b2aa0fb1ee8024325b390 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/action/clip.py @@ -0,0 +1,41 @@ +from typing import Any + +from ray.rllib.connectors.connector import ( + ActionConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.spaces.space_utils import clip_action, get_base_struct_from_space +from ray.rllib.utils.typing import ActionConnectorDataType + + +@OldAPIStack +class ClipActionsConnector(ActionConnector): + def __init__(self, ctx: ConnectorContext): + super().__init__(ctx) + + self._action_space_struct = get_base_struct_from_space(ctx.action_space) + + def transform(self, ac_data: ActionConnectorDataType) -> ActionConnectorDataType: + assert isinstance( + ac_data.output, tuple + ), "Action connector requires PolicyOutputType data." + + actions, states, fetches = ac_data.output + return ActionConnectorDataType( + ac_data.env_id, + ac_data.agent_id, + ac_data.input_dict, + (clip_action(actions, self._action_space_struct), states, fetches), + ) + + def to_state(self): + return ClipActionsConnector.__name__, None + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return ClipActionsConnector(ctx) + + +register_connector(ClipActionsConnector.__name__, ClipActionsConnector) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/immutable.py b/lib/python3.12/site-packages/ray/rllib/connectors/action/immutable.py new file mode 100644 index 0000000000000000000000000000000000000000..13f3cbc3c722bd67c70cbb8d516e384a3049abbc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/action/immutable.py @@ -0,0 +1,40 @@ +from typing import Any + +import tree # pip install dm_tree + +from ray.rllib.connectors.connector import ( + ActionConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.numpy import make_action_immutable +from ray.rllib.utils.typing import ActionConnectorDataType + + +@OldAPIStack +class ImmutableActionsConnector(ActionConnector): + def transform(self, ac_data: ActionConnectorDataType) -> ActionConnectorDataType: + assert isinstance( + ac_data.output, tuple + ), "Action connector requires PolicyOutputType data." + + actions, states, fetches = ac_data.output + tree.traverse(make_action_immutable, actions, top_down=False) + + return ActionConnectorDataType( + ac_data.env_id, + ac_data.agent_id, + ac_data.input_dict, + (actions, states, fetches), + ) + + def to_state(self): + return ImmutableActionsConnector.__name__, None + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return ImmutableActionsConnector(ctx) + + +register_connector(ImmutableActionsConnector.__name__, ImmutableActionsConnector) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/lambdas.py b/lib/python3.12/site-packages/ray/rllib/connectors/action/lambdas.py new file mode 100644 index 0000000000000000000000000000000000000000..8501e66ba3c295a1dae6a2bb8acc1943f21a3a3a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/action/lambdas.py @@ -0,0 +1,76 @@ +from typing import Any, Callable, Dict, Type + +from ray.rllib.connectors.connector import ( + ActionConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import ( + ActionConnectorDataType, + PolicyOutputType, + StateBatches, + TensorStructType, +) + + +@OldAPIStack +def register_lambda_action_connector( + name: str, fn: Callable[[TensorStructType, StateBatches, Dict], PolicyOutputType] +) -> Type[ActionConnector]: + """A util to register any function transforming PolicyOutputType as an ActionConnector. + + The only requirement is that fn should take actions, states, and fetches as input, + and return transformed actions, states, and fetches. + + Args: + name: Name of the resulting actor connector. + fn: The function that transforms PolicyOutputType. + + Returns: + A new ActionConnector class that transforms PolicyOutputType using fn. + """ + + class LambdaActionConnector(ActionConnector): + def transform( + self, ac_data: ActionConnectorDataType + ) -> ActionConnectorDataType: + assert isinstance( + ac_data.output, tuple + ), "Action connector requires PolicyOutputType data." + + actions, states, fetches = ac_data.output + return ActionConnectorDataType( + ac_data.env_id, + ac_data.agent_id, + ac_data.input_dict, + fn(actions, states, fetches), + ) + + def to_state(self): + return name, None + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return LambdaActionConnector(ctx) + + LambdaActionConnector.__name__ = name + LambdaActionConnector.__qualname__ = name + + register_connector(name, LambdaActionConnector) + + return LambdaActionConnector + + +# Convert actions and states into numpy arrays if necessary. +ConvertToNumpyConnector = OldAPIStack( + register_lambda_action_connector( + "ConvertToNumpyConnector", + lambda actions, states, fetches: ( + convert_to_numpy(actions), + convert_to_numpy(states), + fetches, + ), + ), +) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/normalize.py b/lib/python3.12/site-packages/ray/rllib/connectors/action/normalize.py new file mode 100644 index 0000000000000000000000000000000000000000..de20adc09e99b3bbb6d6f9feb7a07766db22c307 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/action/normalize.py @@ -0,0 +1,44 @@ +from typing import Any + +from ray.rllib.connectors.connector import ( + ActionConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.spaces.space_utils import ( + get_base_struct_from_space, + unsquash_action, +) +from ray.rllib.utils.typing import ActionConnectorDataType + + +@OldAPIStack +class NormalizeActionsConnector(ActionConnector): + def __init__(self, ctx: ConnectorContext): + super().__init__(ctx) + + self._action_space_struct = get_base_struct_from_space(ctx.action_space) + + def transform(self, ac_data: ActionConnectorDataType) -> ActionConnectorDataType: + assert isinstance( + ac_data.output, tuple + ), "Action connector requires PolicyOutputType data." + + actions, states, fetches = ac_data.output + return ActionConnectorDataType( + ac_data.env_id, + ac_data.agent_id, + ac_data.input_dict, + (unsquash_action(actions, self._action_space_struct), states, fetches), + ) + + def to_state(self): + return NormalizeActionsConnector.__name__, None + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return NormalizeActionsConnector(ctx) + + +register_connector(NormalizeActionsConnector.__name__, NormalizeActionsConnector) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/action/pipeline.py b/lib/python3.12/site-packages/ray/rllib/connectors/action/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..8f78da44fb48cc413230dc97a1454c013015625a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/action/pipeline.py @@ -0,0 +1,60 @@ +import logging +from collections import defaultdict +from typing import Any, List + +from ray.rllib.connectors.connector import ( + ActionConnector, + Connector, + ConnectorContext, + ConnectorPipeline, +) +from ray.rllib.connectors.registry import get_connector, register_connector +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import ActionConnectorDataType +from ray.util.timer import _Timer + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class ActionConnectorPipeline(ConnectorPipeline, ActionConnector): + def __init__(self, ctx: ConnectorContext, connectors: List[Connector]): + super().__init__(ctx, connectors) + self.timers = defaultdict(_Timer) + + def __call__(self, ac_data: ActionConnectorDataType) -> ActionConnectorDataType: + for c in self.connectors: + timer = self.timers[str(c)] + with timer: + ac_data = c(ac_data) + return ac_data + + def to_state(self): + children = [] + for c in self.connectors: + state = c.to_state() + assert isinstance(state, tuple) and len(state) == 2, ( + "Serialized connector state must be in the format of " + f"Tuple[name: str, params: Any]. Instead we got {state}" + f"for connector {c.__name__}." + ) + children.append(state) + return ActionConnectorPipeline.__name__, children + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + assert ( + type(params) is list + ), "ActionConnectorPipeline takes a list of connector params." + connectors = [] + for state in params: + try: + name, subparams = state + connectors.append(get_connector(name, ctx, subparams)) + except Exception as e: + logger.error(f"Failed to de-serialize connector state: {state}") + raise e + return ActionConnectorPipeline(ctx, connectors) + + +register_connector(ActionConnectorPipeline.__name__, ActionConnectorPipeline) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__init__.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11abb4832c07044a7401e8c3a12682ba399b52f1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/clip_reward.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/clip_reward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4e468deb3f05f9b2e3b1d27c4a453545f1e7621 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/clip_reward.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/env_sampling.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/env_sampling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36f0cc7acb47c2700e8b8f2e6530c9950d493797 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/env_sampling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/lambdas.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/lambdas.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cea0ab5ebc72e3ed790497979f28a33e2e93c473 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/lambdas.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/mean_std_filter.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/mean_std_filter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce8106cea64d06da3928354b03f84adec05fda36 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/mean_std_filter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/obs_preproc.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/obs_preproc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad1eb9b9ca60b80dbe1971d6963be49c07c99f3c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/obs_preproc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/pipeline.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e64e69296dcef5c5dae4e731cbf08200f6f4a017 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/state_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/state_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c8f31ea5d83f3a9dbef4e87537d3d7c1f14766d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/state_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/synced_filter.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/synced_filter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..99a3b46391d8e54c2452a1526b3ae668ff257171 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/synced_filter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/view_requirement.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/view_requirement.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ca5485adfce826345791d6a3842ed66d4635149 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/agent/__pycache__/view_requirement.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/clip_reward.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/clip_reward.py new file mode 100644 index 0000000000000000000000000000000000000000..b5ce8b5cee71c30fbe926afab59872278d603063 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/clip_reward.py @@ -0,0 +1,56 @@ +from typing import Any + +import numpy as np + +from ray.rllib.connectors.connector import ( + AgentConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import AgentConnectorDataType + + +@OldAPIStack +class ClipRewardAgentConnector(AgentConnector): + def __init__(self, ctx: ConnectorContext, sign=False, limit=None): + super().__init__(ctx) + assert ( + not sign or not limit + ), "should not enable both sign and limit reward clipping." + self.sign = sign + self.limit = limit + + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + d = ac_data.data + assert ( + type(d) is dict + ), "Single agent data must be of type Dict[str, TensorStructType]" + + if SampleBatch.REWARDS not in d: + # Nothing to clip. May happen for initial obs. + return ac_data + + if self.sign: + d[SampleBatch.REWARDS] = np.sign(d[SampleBatch.REWARDS]) + elif self.limit: + d[SampleBatch.REWARDS] = np.clip( + d[SampleBatch.REWARDS], + a_min=-self.limit, + a_max=self.limit, + ) + return ac_data + + def to_state(self): + return ClipRewardAgentConnector.__name__, { + "sign": self.sign, + "limit": self.limit, + } + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return ClipRewardAgentConnector(ctx, **params) + + +register_connector(ClipRewardAgentConnector.__name__, ClipRewardAgentConnector) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/env_sampling.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/env_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..0621e100c7f0328dfb74bc0d1576b437e34f1c4b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/env_sampling.py @@ -0,0 +1,30 @@ +from typing import Any + +from ray.rllib.connectors.connector import ( + AgentConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import AgentConnectorDataType + + +@OldAPIStack +class EnvSamplingAgentConnector(AgentConnector): + def __init__(self, ctx: ConnectorContext, sign=False, limit=None): + super().__init__(ctx) + self.observation_space = ctx.observation_space + + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + # EnvSamplingAgentConnector is a no-op connector. + return ac_data + + def to_state(self): + return EnvSamplingAgentConnector.__name__, {} + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return EnvSamplingAgentConnector(ctx, **params) + + +register_connector(EnvSamplingAgentConnector.__name__, EnvSamplingAgentConnector) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/lambdas.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/lambdas.py new file mode 100644 index 0000000000000000000000000000000000000000..643f5e4ec351b9ea21f0d322f53f6d5b4e0bbd48 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/lambdas.py @@ -0,0 +1,86 @@ +from typing import Any, Callable, Type + +import numpy as np +import tree # dm_tree + +from ray.rllib.connectors.connector import ( + AgentConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import ( + AgentConnectorDataType, + AgentConnectorsOutput, +) + + +@OldAPIStack +def register_lambda_agent_connector( + name: str, fn: Callable[[Any], Any] +) -> Type[AgentConnector]: + """A util to register any simple transforming function as an AgentConnector + + The only requirement is that fn should take a single data object and return + a single data object. + + Args: + name: Name of the resulting actor connector. + fn: The function that transforms env / agent data. + + Returns: + A new AgentConnector class that transforms data using fn. + """ + + class LambdaAgentConnector(AgentConnector): + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + return AgentConnectorDataType( + ac_data.env_id, ac_data.agent_id, fn(ac_data.data) + ) + + def to_state(self): + return name, None + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return LambdaAgentConnector(ctx) + + LambdaAgentConnector.__name__ = name + LambdaAgentConnector.__qualname__ = name + + register_connector(name, LambdaAgentConnector) + + return LambdaAgentConnector + + +@OldAPIStack +def flatten_data(data: AgentConnectorsOutput): + assert isinstance( + data, AgentConnectorsOutput + ), "Single agent data must be of type AgentConnectorsOutput" + + raw_dict = data.raw_dict + sample_batch = data.sample_batch + + flattened = {} + for k, v in sample_batch.items(): + if k in [SampleBatch.INFOS, SampleBatch.ACTIONS] or k.startswith("state_out_"): + # Do not flatten infos, actions, and state_out_ columns. + flattened[k] = v + continue + if v is None: + # Keep the same column shape. + flattened[k] = None + continue + flattened[k] = np.array(tree.flatten(v)) + flattened = SampleBatch(flattened, is_training=False) + + return AgentConnectorsOutput(raw_dict, flattened) + + +# Agent connector to build and return a flattened observation SampleBatch +# in addition to the original input dict. +FlattenDataAgentConnector = OldAPIStack( + register_lambda_agent_connector("FlattenDataAgentConnector", flatten_data) +) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/mean_std_filter.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/mean_std_filter.py new file mode 100644 index 0000000000000000000000000000000000000000..b22abb0867efa922f0657ac6c4f031edc8b28156 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/mean_std_filter.py @@ -0,0 +1,190 @@ +from typing import Any, List + +import numpy as np +import tree +from gymnasium.spaces import Discrete, MultiDiscrete + +from ray.rllib.connectors.agent.synced_filter import SyncedFilterAgentConnector +from ray.rllib.connectors.connector import ( + AgentConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.filter import ( + ConcurrentMeanStdFilter, + Filter, + MeanStdFilter, + RunningStat, +) +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.typing import AgentConnectorDataType + + +@OldAPIStack +class MeanStdObservationFilterAgentConnector(SyncedFilterAgentConnector): + """A connector used to mean-std-filter observations. + + Incoming observations are filtered such that the output of this filter is on + average zero and has a standard deviation of 1. This filtering is applied + separately per element of the observation space. + """ + + def __init__( + self, + ctx: ConnectorContext, + demean: bool = True, + destd: bool = True, + clip: float = 10.0, + ): + SyncedFilterAgentConnector.__init__(self, ctx) + # We simply use the old MeanStdFilter until non-connector env_runner is fully + # deprecated to avoid duplicate code + + filter_shape = tree.map_structure( + lambda s: ( + None + if isinstance(s, (Discrete, MultiDiscrete)) # noqa + else np.array(s.shape) + ), + get_base_struct_from_space(ctx.observation_space), + ) + self.filter = MeanStdFilter(filter_shape, demean=demean, destd=destd, clip=clip) + + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + d = ac_data.data + assert ( + type(d) is dict + ), "Single agent data must be of type Dict[str, TensorStructType]" + if SampleBatch.OBS in d: + d[SampleBatch.OBS] = self.filter( + d[SampleBatch.OBS], update=self._is_training + ) + if SampleBatch.NEXT_OBS in d: + d[SampleBatch.NEXT_OBS] = self.filter( + d[SampleBatch.NEXT_OBS], update=self._is_training + ) + + return ac_data + + def to_state(self): + # Flattening is deterministic + flattened_rs = tree.flatten(self.filter.running_stats) + flattened_buffer = tree.flatten(self.filter.buffer) + return MeanStdObservationFilterAgentConnector.__name__, { + "shape": self.filter.shape, + "no_preprocessor": self.filter.no_preprocessor, + "demean": self.filter.demean, + "destd": self.filter.destd, + "clip": self.filter.clip, + "running_stats": [s.to_state() for s in flattened_rs], + "buffer": [s.to_state() for s in flattened_buffer], + } + + # demean, destd, clip, and a state dict + @staticmethod + def from_state( + ctx: ConnectorContext, + params: List[Any] = None, + demean: bool = True, + destd: bool = True, + clip: float = 10.0, + ): + connector = MeanStdObservationFilterAgentConnector(ctx, demean, destd, clip) + if params: + connector.filter.shape = params["shape"] + connector.filter.no_preprocessor = params["no_preprocessor"] + connector.filter.demean = params["demean"] + connector.filter.destd = params["destd"] + connector.filter.clip = params["clip"] + + # Unflattening is deterministic + running_stats = [RunningStat.from_state(s) for s in params["running_stats"]] + connector.filter.running_stats = tree.unflatten_as( + connector.filter.shape, running_stats + ) + + # Unflattening is deterministic + buffer = [RunningStat.from_state(s) for s in params["buffer"]] + connector.filter.buffer = tree.unflatten_as(connector.filter.shape, buffer) + + return connector + + def reset_state(self) -> None: + """Creates copy of current state and resets accumulated state""" + if not self._is_training: + raise ValueError( + "State of {} can only be changed when trainin.".format(self.__name__) + ) + self.filter.reset_buffer() + + def apply_changes(self, other: "Filter", *args, **kwargs) -> None: + """Updates self with state from other filter.""" + # inline this as soon as we deprecate ordinary filter with non-connector + # env_runner + if not self._is_training: + raise ValueError( + "Changes can only be applied to {} when trainin.".format(self.__name__) + ) + return self.filter.apply_changes(other, *args, **kwargs) + + def copy(self) -> "Filter": + """Creates a new object with same state as self. + + This is a legacy Filter method that we need to keep around for now + + Returns: + A copy of self. + """ + # inline this as soon as we deprecate ordinary filter with non-connector + # env_runner + return self.filter.copy() + + def sync(self, other: "AgentConnector") -> None: + """Copies all state from other filter to self.""" + # inline this as soon as we deprecate ordinary filter with non-connector + # env_runner + if not self._is_training: + raise ValueError( + "{} can only be synced when trainin.".format(self.__name__) + ) + return self.filter.sync(other.filter) + + +@OldAPIStack +class ConcurrentMeanStdObservationFilterAgentConnector( + MeanStdObservationFilterAgentConnector +): + """A concurrent version of the MeanStdObservationFilterAgentConnector. + + This version's filter has all operations wrapped by a threading.RLock. + It can therefore be safely used by multiple threads. + """ + + def __init__(self, ctx: ConnectorContext, demean=True, destd=True, clip=10.0): + SyncedFilterAgentConnector.__init__(self, ctx) + # We simply use the old MeanStdFilter until non-connector env_runner is fully + # deprecated to avoid duplicate code + + filter_shape = tree.map_structure( + lambda s: ( + None + if isinstance(s, (Discrete, MultiDiscrete)) # noqa + else np.array(s.shape) + ), + get_base_struct_from_space(ctx.observation_space), + ) + self.filter = ConcurrentMeanStdFilter( + filter_shape, demean=True, destd=True, clip=10.0 + ) + + +register_connector( + MeanStdObservationFilterAgentConnector.__name__, + MeanStdObservationFilterAgentConnector, +) +register_connector( + ConcurrentMeanStdObservationFilterAgentConnector.__name__, + ConcurrentMeanStdObservationFilterAgentConnector, +) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/obs_preproc.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/obs_preproc.py new file mode 100644 index 0000000000000000000000000000000000000000..300fee9f7d5b6f9611794062851bc56a5e9a02c7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/obs_preproc.py @@ -0,0 +1,69 @@ +from typing import Any + +from ray.rllib.connectors.connector import ( + AgentConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.models.preprocessors import NoPreprocessor, get_preprocessor +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import AgentConnectorDataType + + +@OldAPIStack +class ObsPreprocessorConnector(AgentConnector): + """A connector that wraps around existing RLlib observation preprocessors. + + This includes: + - OneHotPreprocessor for Discrete and Multi-Discrete spaces. + - GenericPixelPreprocessor and AtariRamPreprocessor for Atari spaces. + - TupleFlatteningPreprocessor and DictFlatteningPreprocessor for flattening + arbitrary nested input observations. + - RepeatedValuesPreprocessor for padding observations from RLlib Repeated + observation space. + """ + + def __init__(self, ctx: ConnectorContext): + super().__init__(ctx) + + if hasattr(ctx.observation_space, "original_space"): + # ctx.observation_space is the space this Policy deals with. + # We need to preprocess data from the original observation space here. + obs_space = ctx.observation_space.original_space + else: + obs_space = ctx.observation_space + + self._preprocessor = get_preprocessor(obs_space)( + obs_space, ctx.config.get("model", {}) + ) + + def is_identity(self): + """Returns whether this preprocessor connector is a no-op preprocessor.""" + return isinstance(self._preprocessor, NoPreprocessor) + + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + d = ac_data.data + assert type(d) is dict, ( + "Single agent data must be of type Dict[str, TensorStructType] but is of " + "type {}".format(type(d)) + ) + + if SampleBatch.OBS in d: + d[SampleBatch.OBS] = self._preprocessor.transform(d[SampleBatch.OBS]) + if SampleBatch.NEXT_OBS in d: + d[SampleBatch.NEXT_OBS] = self._preprocessor.transform( + d[SampleBatch.NEXT_OBS] + ) + + return ac_data + + def to_state(self): + return ObsPreprocessorConnector.__name__, None + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return ObsPreprocessorConnector(ctx) + + +register_connector(ObsPreprocessorConnector.__name__, ObsPreprocessorConnector) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/pipeline.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..26655386094a1e7b321620d7cfc7229c783ac042 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/pipeline.py @@ -0,0 +1,71 @@ +import logging +from collections import defaultdict +from typing import Any, List + +from ray.rllib.connectors.connector import ( + AgentConnector, + Connector, + ConnectorContext, + ConnectorPipeline, +) +from ray.rllib.connectors.registry import get_connector, register_connector +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import ActionConnectorDataType, AgentConnectorDataType +from ray.util.timer import _Timer + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class AgentConnectorPipeline(ConnectorPipeline, AgentConnector): + def __init__(self, ctx: ConnectorContext, connectors: List[Connector]): + super().__init__(ctx, connectors) + self.timers = defaultdict(_Timer) + + def reset(self, env_id: str): + for c in self.connectors: + c.reset(env_id) + + def on_policy_output(self, output: ActionConnectorDataType): + for c in self.connectors: + c.on_policy_output(output) + + def __call__( + self, acd_list: List[AgentConnectorDataType] + ) -> List[AgentConnectorDataType]: + ret = acd_list + for c in self.connectors: + timer = self.timers[str(c)] + with timer: + ret = c(ret) + return ret + + def to_state(self): + children = [] + for c in self.connectors: + state = c.to_state() + assert isinstance(state, tuple) and len(state) == 2, ( + "Serialized connector state must be in the format of " + f"Tuple[name: str, params: Any]. Instead we got {state}" + f"for connector {c.__name__}." + ) + children.append(state) + return AgentConnectorPipeline.__name__, children + + @staticmethod + def from_state(ctx: ConnectorContext, params: List[Any]): + assert ( + type(params) is list + ), "AgentConnectorPipeline takes a list of connector params." + connectors = [] + for state in params: + try: + name, subparams = state + connectors.append(get_connector(name, ctx, subparams)) + except Exception as e: + logger.error(f"Failed to de-serialize connector state: {state}") + raise e + return AgentConnectorPipeline(ctx, connectors) + + +register_connector(AgentConnectorPipeline.__name__, AgentConnectorPipeline) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/state_buffer.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/state_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..54ff17160ca5b78c19243a792bd29cbe10ec6d43 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/state_buffer.py @@ -0,0 +1,118 @@ +import logging +import pickle +from collections import defaultdict +from typing import Any + +import numpy as np +import tree # dm_tree + +from ray import cloudpickle +from ray.rllib.connectors.connector import ( + AgentConnector, + Connector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.core.columns import Columns +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.typing import ActionConnectorDataType, AgentConnectorDataType + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class StateBufferConnector(AgentConnector): + def __init__(self, ctx: ConnectorContext, states: Any = None): + super().__init__(ctx) + + self._initial_states = ctx.initial_states + self._action_space_struct = get_base_struct_from_space(ctx.action_space) + + self._states = defaultdict(lambda: defaultdict(lambda: (None, None, None))) + self._enable_new_api_stack = False + # TODO(jungong) : we would not need this if policies are never stashed + # during the rollout of a single episode. + if states: + try: + self._states = cloudpickle.loads(states) + except pickle.UnpicklingError: + # StateBufferConnector states are only needed for rare cases + # like stashing then restoring a policy during the rollout of + # a single episode. + # It is ok to ignore the error for most of the cases here. + logger.info( + "Can not restore StateBufferConnector states. This warning can " + "usually be ignore, unless it is from restoring a stashed policy." + ) + + @override(Connector) + def in_eval(self): + super().in_eval() + + def reset(self, env_id: str): + # States should not be carried over between episodes. + if env_id in self._states: + del self._states[env_id] + + def on_policy_output(self, ac_data: ActionConnectorDataType): + # Buffer latest output states for next input __call__. + self._states[ac_data.env_id][ac_data.agent_id] = ac_data.output + + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + d = ac_data.data + assert ( + type(d) is dict + ), "Single agent data must be of type Dict[str, TensorStructType]" + + env_id = ac_data.env_id + agent_id = ac_data.agent_id + assert ( + env_id is not None and agent_id is not None + ), f"StateBufferConnector requires env_id(f{env_id}) and agent_id(f{agent_id})" + + action, states, fetches = self._states[env_id][agent_id] + + if action is not None: + d[SampleBatch.ACTIONS] = action # Last action + else: + # Default zero action. + d[SampleBatch.ACTIONS] = tree.map_structure( + lambda s: np.zeros_like(s.sample(), s.dtype) + if hasattr(s, "dtype") + else np.zeros_like(s.sample()), + self._action_space_struct, + ) + + if states is None: + states = self._initial_states + if self._enable_new_api_stack: + if states: + d[Columns.STATE_OUT] = states + else: + for i, v in enumerate(states): + d["state_out_{}".format(i)] = v + + # Also add extra fetches if available. + if fetches: + d.update(fetches) + + return ac_data + + def to_state(self): + # Note(jungong) : it is ok to use cloudpickle here for stats because: + # 1. self._states may contain arbitary data objects, and will be hard + # to serialize otherwise. + # 2. seriazlized states are only useful if a policy is stashed and + # restored during the rollout of a single episode. So it is ok to + # use cloudpickle for such non-persistent data bits. + states = cloudpickle.dumps(self._states) + return StateBufferConnector.__name__, states + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return StateBufferConnector(ctx, params) + + +register_connector(StateBufferConnector.__name__, StateBufferConnector) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/synced_filter.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/synced_filter.py new file mode 100644 index 0000000000000000000000000000000000000000..51c625d9aceec9401599ef87b42cf522704c5ea3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/synced_filter.py @@ -0,0 +1,52 @@ +from ray.rllib.connectors.connector import ( + AgentConnector, + ConnectorContext, +) +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.filter import Filter + + +@OldAPIStack +class SyncedFilterAgentConnector(AgentConnector): + """An agent connector that filters with synchronized parameters.""" + + def __init__(self, ctx: ConnectorContext, *args, **kwargs): + super().__init__(ctx) + if args or kwargs: + raise ValueError( + "SyncedFilterAgentConnector does not take any additional arguments, " + "but got args=`{}` and kwargs={}.".format(args, kwargs) + ) + + def apply_changes(self, other: "Filter", *args, **kwargs) -> None: + """Updates self with state from other filter.""" + # TODO: (artur) inline this as soon as we deprecate ordinary filter with + # non-connecto env_runner + return self.filter.apply_changes(other, *args, **kwargs) + + def copy(self) -> "Filter": + """Creates a new object with same state as self. + + This is a legacy Filter method that we need to keep around for now + + Returns: + A copy of self. + """ + # inline this as soon as we deprecate ordinary filter with non-connector + # env_runner + return self.filter.copy() + + def sync(self, other: "AgentConnector") -> None: + """Copies all state from other filter to self.""" + # TODO: (artur) inline this as soon as we deprecate ordinary filter with + # non-connector env_runner + return self.filter.sync(other.filter) + + def reset_state(self) -> None: + """Creates copy of current state and resets accumulated state""" + raise NotImplementedError + + def as_serializable(self) -> "Filter": + # TODO: (artur) inline this as soon as we deprecate ordinary filter with + # non-connector env_runner + return self.filter.as_serializable() diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/agent/view_requirement.py b/lib/python3.12/site-packages/ray/rllib/connectors/agent/view_requirement.py new file mode 100644 index 0000000000000000000000000000000000000000..c24ef373fa496e4abac0f9d605765895e07e0c3a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/agent/view_requirement.py @@ -0,0 +1,135 @@ +from collections import defaultdict +from typing import Any + +from ray.rllib.connectors.connector import ( + AgentConnector, + ConnectorContext, +) +from ray.rllib.connectors.registry import register_connector +from ray.rllib.evaluation.collectors.agent_collector import AgentCollector +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import ( + AgentConnectorDataType, + AgentConnectorsOutput, +) + + +@OldAPIStack +class ViewRequirementAgentConnector(AgentConnector): + """This connector does 2 things: + 1. It filters data columns based on view_requirements for training and inference. + 2. It buffers the right amount of history for computing the sample batch for + action computation. + The output of this connector is AgentConnectorsOut, which basically is + a tuple of 2 things: + { + "raw_dict": {"obs": ...} + "sample_batch": SampleBatch + } + raw_dict, which contains raw data for the latest time slice, + can be used to construct a complete episode by Sampler for training purpose. + The "for_action" SampleBatch can be used to directly call the policy. + """ + + def __init__(self, ctx: ConnectorContext): + super().__init__(ctx) + + self._view_requirements = ctx.view_requirements + _enable_new_api_stack = False + + # a dict of env_id to a dict of agent_id to a list of agent_collector objects + self.agent_collectors = defaultdict( + lambda: defaultdict( + lambda: AgentCollector( + self._view_requirements, + max_seq_len=ctx.config["model"]["max_seq_len"], + intial_states=ctx.initial_states, + disable_action_flattening=ctx.config.get( + "_disable_action_flattening", False + ), + is_policy_recurrent=ctx.is_policy_recurrent, + # Note(jungong): We only leverage AgentCollector for building sample + # batches for computing actions. + # So regardless of whether this ViewRequirement connector is in + # training or inference mode, we should tell these AgentCollectors + # to behave in inference mode, so they don't accumulate episode data + # that is not useful for inference. + is_training=False, + _enable_new_api_stack=_enable_new_api_stack, + ) + ) + ) + + def reset(self, env_id: str): + if env_id in self.agent_collectors: + del self.agent_collectors[env_id] + + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + d = ac_data.data + assert ( + type(d) is dict + ), "Single agent data must be of type Dict[str, TensorStructType]" + + env_id = ac_data.env_id + agent_id = ac_data.agent_id + # TODO: we don't keep episode_id around so use env_id as episode_id ? + episode_id = env_id if SampleBatch.EPS_ID not in d else d[SampleBatch.EPS_ID] + + assert env_id is not None and agent_id is not None, ( + f"ViewRequirementAgentConnector requires env_id({env_id}) " + "and agent_id({agent_id})" + ) + + assert ( + self._view_requirements + ), "ViewRequirements required by ViewRequirementAgentConnector" + + # Note(jungong) : we need to keep the entire input dict here. + # A column may be used by postprocessing (GAE) even if its + # view_requirement.used_for_training is False. + training_dict = d + + agent_collector = self.agent_collectors[env_id][agent_id] + + if SampleBatch.NEXT_OBS not in d: + raise ValueError(f"connector data {d} should contain next_obs.") + # TODO(avnishn; kourosh) Unsure how agent_index is necessary downstream + # since there is no mapping from agent_index to agent_id that exists. + # need to remove this from the SampleBatch later. + # fall back to using dummy index if no index is available + if SampleBatch.AGENT_INDEX in d: + agent_index = d[SampleBatch.AGENT_INDEX] + else: + try: + agent_index = float(agent_id) + except ValueError: + agent_index = -1 + if agent_collector.is_empty(): + agent_collector.add_init_obs( + episode_id=episode_id, + agent_index=agent_index, + env_id=env_id, + init_obs=d[SampleBatch.NEXT_OBS], + init_infos=d.get(SampleBatch.INFOS), + ) + else: + agent_collector.add_action_reward_next_obs(d) + sample_batch = agent_collector.build_for_inference() + + return_data = AgentConnectorDataType( + env_id, agent_id, AgentConnectorsOutput(training_dict, sample_batch) + ) + return return_data + + def to_state(self): + return ViewRequirementAgentConnector.__name__, None + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any): + return ViewRequirementAgentConnector(ctx) + + +register_connector( + ViewRequirementAgentConnector.__name__, ViewRequirementAgentConnector +) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__init__.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b11ec7b77263d875e4fae999baaba11273f9c504 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/__init__.py @@ -0,0 +1,23 @@ +from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import ( + AddObservationsFromEpisodesToBatch, +) +from ray.rllib.connectors.common.add_states_from_episodes_to_batch import ( + AddStatesFromEpisodesToBatch, +) +from ray.rllib.connectors.common.add_time_dim_to_batch_and_zero_pad import ( + AddTimeDimToBatchAndZeroPad, +) +from ray.rllib.connectors.common.agent_to_module_mapping import AgentToModuleMapping +from ray.rllib.connectors.common.batch_individual_items import BatchIndividualItems +from ray.rllib.connectors.common.numpy_to_tensor import NumpyToTensor +from ray.rllib.connectors.common.tensor_to_numpy import TensorToNumpy + +__all__ = [ + "AddObservationsFromEpisodesToBatch", + "AddStatesFromEpisodesToBatch", + "AddTimeDimToBatchAndZeroPad", + "AgentToModuleMapping", + "BatchIndividualItems", + "NumpyToTensor", + "TensorToNumpy", +] diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..364516352c3ee878b402ae8ade4817c775e35e4a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_observations_from_episodes_to_batch.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_observations_from_episodes_to_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4bc88e842cfac82dec7d011fac2b79ce344387cc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_observations_from_episodes_to_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_states_from_episodes_to_batch.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_states_from_episodes_to_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5d9cfb05e4eddf63ae86ba88049a02ee9e8f98e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_states_from_episodes_to_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_time_dim_to_batch_and_zero_pad.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_time_dim_to_batch_and_zero_pad.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d3c79720e0719dd5887fe42c6af2b97ad0dd90c8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/add_time_dim_to_batch_and_zero_pad.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/agent_to_module_mapping.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/agent_to_module_mapping.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e4c884ee553a76545df259c193d630f6c5c1ce2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/agent_to_module_mapping.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/batch_individual_items.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/batch_individual_items.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad91bd1840e44ba1c3b216a4c929956b485292b0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/batch_individual_items.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/flatten_observations.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/flatten_observations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..396d149f10205afdbd6436f6fc24ba6d68605a72 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/flatten_observations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/frame_stacking.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/frame_stacking.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f25dff7cfd5a5dcb16637a4be95b908d6d93f5d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/frame_stacking.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/module_to_agent_unmapping.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/module_to_agent_unmapping.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7966e106525e43f721ef4f9df7d494642ccd919f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/module_to_agent_unmapping.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/numpy_to_tensor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/numpy_to_tensor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5be6358a25bee6f37427514ce23bbc82e593f3ad Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/numpy_to_tensor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/tensor_to_numpy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/tensor_to_numpy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc500b6a60ed9439e316de3c5d71b70b8c5c5a9c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/common/__pycache__/tensor_to_numpy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/add_observations_from_episodes_to_batch.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/add_observations_from_episodes_to_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..189684ba97d4dab97e5edbfc40f8dbc84e897fa4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/add_observations_from_episodes_to_batch.py @@ -0,0 +1,180 @@ +from typing import Any, Dict, List, Optional + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AddObservationsFromEpisodesToBatch(ConnectorV2): + """Gets the last observation from a running episode and adds it to the batch. + + Note: This is one of the default env-to-module or Learner ConnectorV2 pieces that + are added automatically by RLlib into every env-to-module/Learner connector + pipeline, unless `config.add_default_connectors_to_env_to_module_pipeline` or + `config.add_default_connectors_to_learner_pipeline ` are set to + False. + + The default env-to-module connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + The default Learner connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddColumnsFromEpisodesToTrainBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + This ConnectorV2: + - Operates on a list of Episode objects (single- or multi-agent). + - Gets the most recent observation(s) from all the given episodes and adds them + to the batch under construction (as a list of individual observations). + - Does NOT alter any observations (or other data) in the given episodes. + - Can be used in EnvToModule and Learner connector pipelines. + + .. testcode:: + + import gymnasium as gym + import numpy as np + + from ray.rllib.connectors.common import AddObservationsFromEpisodesToBatch + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Create two dummy SingleAgentEpisodes, each containing 2 observations, + # 1 action and 1 reward (both are length=1). + obs_space = gym.spaces.Box(-1.0, 1.0, (2,), np.float32) + act_space = gym.spaces.Discrete(2) + + episodes = [SingleAgentEpisode( + observations=[obs_space.sample(), obs_space.sample()], + actions=[act_space.sample()], + rewards=[1.0], + len_lookback_buffer=0, + ) for _ in range(2)] + eps_1_last_obs = episodes[0].get_observations(-1) + eps_2_last_obs = episodes[1].get_observations(-1) + print(f"1st Episode's last obs is {eps_1_last_obs}") + print(f"2nd Episode's last obs is {eps_2_last_obs}") + + # Create an instance of this class. + connector = AddObservationsFromEpisodesToBatch() + + # Call the connector with the two created episodes. + # Note that this particular connector works without an RLModule, so we + # simplify here for the sake of this example. + output_batch = connector( + rl_module=None, + batch={}, + episodes=episodes, + explore=True, + shared_data={}, + ) + # The output data should now contain the last observations of both episodes, + # in a "per-episode organized" fashion. + check( + output_batch, + { + "obs": { + (episodes[0].id_,): [eps_1_last_obs], + (episodes[1].id_,): [eps_2_last_obs], + }, + }, + ) + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + as_learner_connector: bool = False, + **kwargs, + ): + """Initializes a AddObservationsFromEpisodesToBatch instance. + + Args: + as_learner_connector: Whether this connector is part of a Learner connector + pipeline, as opposed to a env-to-module pipeline. As a Learner + connector, it will add an entire Episode's observations (each timestep) + to the batch. + """ + super().__init__( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + **kwargs, + ) + + self._as_learner_connector = as_learner_connector + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # If "obs" already in data, early out. + if Columns.OBS in batch: + return batch + for i, sa_episode in enumerate( + self.single_agent_episode_iterator( + episodes, + # If Learner connector, get all episodes (for train batch). + # If EnvToModule, get only those ongoing episodes that just had their + # agent step (b/c those are the ones we need to compute actions for + # next). + agents_that_stepped_only=not self._as_learner_connector, + ) + ): + if self._as_learner_connector: + # TODO (sven): Resolve this hack by adding a new connector piece that + # performs this very task. + if "_" not in sa_episode.id_: + sa_episode.id_ += "_" + str(i) + + self.add_n_batch_items( + batch, + Columns.OBS, + # Add all observations, except the very last one. + # For a terminated episode, this is the terminal observation that + # has no value for training. + # For a truncated episode, algorithms either add an extra NEXT_OBS + # column to the batch (ex. DQN) or extend the episode length by one + # (using a separate connector piece and this truncated last obs), + # then bootstrap the value estimation for that extra timestep. + items_to_add=sa_episode.get_observations(slice(0, len(sa_episode))), + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + else: + assert not sa_episode.is_numpy + self.add_batch_item( + batch, + Columns.OBS, + item_to_add=sa_episode.get_observations(-1), + single_agent_episode=sa_episode, + ) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/add_states_from_episodes_to_batch.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/add_states_from_episodes_to_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..6cfe284d347a25daf4269150f4244a7f6148c235 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/add_states_from_episodes_to_batch.py @@ -0,0 +1,353 @@ +import math +from typing import Any, Dict, List, Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AddStatesFromEpisodesToBatch(ConnectorV2): + """Gets last STATE_OUT from running episode and adds it as STATE_IN to the batch. + + Note: This is one of the default env-to-module or Learner ConnectorV2 pieces that + are added automatically by RLlib into every env-to-module/Learner connector + pipeline, unless `config.add_default_connectors_to_env_to_module_pipeline` or + `config.add_default_connectors_to_learner_pipeline ` are set to + False. + + The default env-to-module connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + The default Learner connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddColumnsFromEpisodesToTrainBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + If the RLModule is stateful, the episodes' STATE_OUTS will be extracted + and restructured under a new STATE_IN key. + As a Learner connector, the resulting STATE_IN batch has the shape (B', ...). + Here, B' is the sum of splits we have to do over the given episodes, such that each + chunk is at most `max_seq_len` long (T-axis). + As a EnvToModule connector, the resulting STATE_IN batch simply consists of n + states coming from n vectorized environments/episodes. + + Also, all other data (observations, rewards, etc.. if applicable) will be properly + reshaped into (B, T=max_seq_len (learner) or 1 (env-to-module), ...) and will be + zero-padded, if necessary. + + This ConnectorV2: + - Operates on a list of Episode objects. + - Gets the most recent STATE_OUT from all the given episodes and adds them under + the STATE_IN key to the batch under construction. + - Does NOT alter any data in the given episodes. + - Can be used in EnvToModule and Learner connector pipelines. + + .. testcode:: + + from ray.rllib.connectors.common import AddStatesFromEpisodesToBatch + from ray.rllib.core.columns import Columns + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Create a simple dummy class, pretending to be an RLModule with + # `get_initial_state`, `is_stateful` and `model_config` property defined: + class MyStateModule: + # dummy config + model_config = {"max_seq_len": 2} + + def is_stateful(self): + return True + + def get_initial_state(self): + return 0.0 + + + # Create an empty episode. The connector should use the RLModule's initial state + # to populate STATE_IN for the next forward pass. + episode = SingleAgentEpisode() + + rl_module = MyStateModule() + rl_module_init_state = rl_module.get_initial_state() + + # Create an instance of this class (as a env-to-module connector). + connector = AddStatesFromEpisodesToBatch(as_learner_connector=False) + + # Call the connector. + output_batch = connector( + rl_module=rl_module, + batch={}, + episodes=[episode], + shared_data={}, + ) + # The output data's STATE_IN key should now contain the RLModule's initial state + # plus the one state out found in the episode in a "per-episode organized" + # fashion. + check( + output_batch[Columns.STATE_IN], + { + (episode.id_,): [rl_module_init_state], + }, + ) + + # Create a SingleAgentEpisodes containing 5 observations, + # 4 actions and 4 rewards, and 4 STATE_OUTs. + # The same connector should now use the episode-stored last STATE_OUT as + # STATE_IN for the next forward pass. + episode = SingleAgentEpisode( + observations=[0, 1, 2, 3, 4], + actions=[1, 2, 3, 4], + rewards=[1.0, 2.0, 3.0, 4.0], + # STATE_OUT in episode will show up under STATE_IN in the batch. + extra_model_outputs={ + Columns.STATE_OUT: [-4.0, -3.0, -2.0, -1.0], + }, + len_lookback_buffer = 0, + ) + + # Call the connector. + output_batch = connector( + rl_module=rl_module, + batch={}, + episodes=[episode], + shared_data={}, + ) + # The output data's STATE_IN key should now contain the episode's last + # STATE_OUT, NOT the RLModule's initial state in a "per-episode organized" + # fashion. + check( + output_batch[Columns.STATE_IN], + { + # Expect the episode's last STATE_OUT. + (episode.id_,): [-1.0], + }, + ) + + # Create a new connector as a learner connector with a RNN seq len of 2 (for + # testing purposes only). Passing the same data through this learner connector, + # we expect the STATE_IN data to contain a) the initial module state and then + # every 2nd STATE_OUT stored in the episode. + connector = AddStatesFromEpisodesToBatch(as_learner_connector=True) + + # Call the connector. + output_batch = connector( + rl_module=rl_module, + batch={}, + episodes=[episode], + shared_data={}, + ) + check( + output_batch[Columns.STATE_IN], + { + # Expect initial module state + every 2nd STATE_OUT from episode, but + # not the very last one (just like the very last observation, this data + # is NOT passed through the forward_train, b/c there is nothing to learn + # at that timestep, unless we need to compute e.g. bootstrap value + # predictions). + # Also note that the different STATE_IN timesteps are already present + # as one batched item per episode in the list. + (episode.id_,): [rl_module_init_state, -3.0], + }, + ) + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + as_learner_connector: bool = False, + **kwargs, + ): + """Initializes a AddObservationsFromEpisodesToBatch instance. + + Args: + as_learner_connector: Whether this connector is part of a Learner connector + pipeline, as opposed to a env-to-module pipeline. As a Learner + connector, it will add an entire Episode's observations (each timestep) + to the batch. + """ + super().__init__( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + **kwargs, + ) + + self._as_learner_connector = as_learner_connector + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # If not stateful OR STATE_IN already in data, early out. + if ( + rl_module is None + or not rl_module.is_stateful() + or Columns.STATE_IN in batch + ): + return batch + + for sa_episode in self.single_agent_episode_iterator( + episodes, + # If Learner connector, get all episodes (for train batch). + # If EnvToModule, get only those ongoing episodes that just had their + # agent step (b/c those are the ones we need to compute actions for next). + agents_that_stepped_only=not self._as_learner_connector, + ): + if self._as_learner_connector: + # Multi-agent case: Extract correct single agent RLModule (to get its + # individual state). + if sa_episode.module_id is not None: + sa_module = rl_module[sa_episode.module_id] + else: + sa_module = ( + rl_module[DEFAULT_MODULE_ID] + if isinstance(rl_module, MultiRLModule) + else rl_module + ) + # This single-agent RLModule is NOT stateful -> Skip. + if not sa_module.is_stateful(): + continue + + max_seq_len = sa_module.model_config["max_seq_len"] + + # look_back_state.shape=([state-dim],) + look_back_state = ( + # Episode has a (reset) beginning -> Prepend initial + # state. + convert_to_numpy(sa_module.get_initial_state()) + if sa_episode.t_started == 0 + or (Columns.STATE_OUT not in sa_episode.extra_model_outputs) + # Episode starts somewhere in the middle (is a cut + # continuation chunk) -> Use previous chunk's last + # STATE_OUT as initial state. + else sa_episode.get_extra_model_outputs( + key=Columns.STATE_OUT, + indices=-1, + neg_index_as_lookback=True, + ) + ) + # If we have `"state_out"`s (e.g. from rollouts) use them for the + # `"state_in"`s. + if Columns.STATE_OUT in sa_episode.extra_model_outputs: + # state_outs.shape=(T,[state-dim]) T=episode len + state_outs = sa_episode.get_extra_model_outputs( + key=Columns.STATE_OUT + ) + # Otherwise, we have no `"state_out"` (e.g. because we are sampling + # from offline data and the expert policy was not stateful). + else: + # Then simply use the `look_back_state`, i.e. in this case the + # initial state as `"state_in` in training. + if sa_episode.is_numpy: + state_outs = tree.map_structure( + lambda a, _sae=sa_episode: np.repeat( + a[np.newaxis, ...], len(_sae), axis=0 + ), + look_back_state, + ) + else: + state_outs = [look_back_state for _ in range(len(sa_episode))] + # Explanation: + # B=episode len // max_seq_len + # [::max_seq_len]: only keep every Tth state. + # [:-1]: Shift state outs by one; ignore very last + # STATE_OUT, but therefore add the lookback/init state at + # the beginning. + items_to_add = ( + tree.map_structure( + lambda i, o, m=max_seq_len: np.concatenate([[i], o[:-1]])[::m], + look_back_state, + state_outs, + ) + if sa_episode.is_numpy + else ([look_back_state] + state_outs[:-1])[::max_seq_len] + ) + self.add_n_batch_items( + batch=batch, + column=Columns.STATE_IN, + items_to_add=items_to_add, + num_items=int(math.ceil(len(sa_episode) / max_seq_len)), + single_agent_episode=sa_episode, + ) + if Columns.NEXT_OBS in batch: + items_to_add = ( + tree.map_structure( + lambda i, m=max_seq_len: i[::m], + state_outs, + ) + if sa_episode.is_numpy + else state_outs[::max_seq_len] + ) + self.add_n_batch_items( + batch=batch, + column=Columns.NEXT_STATE_IN, + items_to_add=items_to_add, + num_items=int(math.ceil(len(sa_episode) / max_seq_len)), + single_agent_episode=sa_episode, + ) + + else: + assert not sa_episode.is_numpy + + # Multi-agent case: Extract correct single agent RLModule (to get the + # state for individually). + sa_module = rl_module + if sa_episode.module_id is not None: + sa_module = rl_module[sa_episode.module_id] + # This single-agent RLModule is NOT stateful -> Skip. + if not sa_module.is_stateful(): + continue + + # Episode just started or has no `"state_out"` (e.g. in offline + # sampling) -> Get initial state from our RLModule. + if (sa_episode.t_started == 0 and len(sa_episode) == 0) or ( + Columns.STATE_OUT not in sa_episode.extra_model_outputs + ): + state = sa_module.get_initial_state() + state = convert_to_numpy(state) + # Episode is already ongoing -> Use most recent STATE_OUT. + else: + state = sa_episode.get_extra_model_outputs( + key=Columns.STATE_OUT, + indices=-1, + ) + self.add_batch_item( + batch, + Columns.STATE_IN, + item_to_add=state, + single_agent_episode=sa_episode, + ) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/add_time_dim_to_batch_and_zero_pad.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/add_time_dim_to_batch_and_zero_pad.py new file mode 100644 index 0000000000000000000000000000000000000000..8358e44ae28183cbdfd72c475a38caa1908db5b3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/add_time_dim_to_batch_and_zero_pad.py @@ -0,0 +1,306 @@ +from typing import Any, Dict, List, Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.postprocessing.zero_padding import ( + create_mask_and_seq_lens, + split_and_zero_pad, +) +from ray.rllib.utils.spaces.space_utils import BatchedNdArray +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AddTimeDimToBatchAndZeroPad(ConnectorV2): + """Adds an extra time dim (axis=1) to all data currently in the batch. + + Note: This is one of the default env-to-module or Learner ConnectorV2 pieces that + are added automatically by RLlib into every env-to-module/Learner connector + pipeline, unless `config.add_default_connectors_to_env_to_module_pipeline` or + `config.add_default_connectors_to_learner_pipeline ` are set to + False. + + The default env-to-module connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + The default Learner connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddColumnsFromEpisodesToTrainBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + If the RLModule is stateful, an extra time dim at axis=1 is added to all data in the + batch. + + Also, all data (observations, rewards, etc.. if applicable) will be properly + reshaped into (B, T=max_seq_len (learner) or 1 (env-to-module), ...) and will be + zero-padded, if necessary. + + This ConnectorV2: + - Operates on a list of Episode objects. + - Adds a time dim at axis=1 to all columns already in the batch. + - In case of a learner connector pipeline, zero-pads the data according to the + module's `self.model_config["max_seq_len"]` setting and reshapes all data to + (B, T, ...). The connector also adds SEQ_LENS information and loss mask + information to the batch based on the added zero-padding. + - Does NOT alter any data in the given episodes. + - Can be used in EnvToModule and Learner connector pipelines. + + .. testcode:: + + from ray.rllib.connectors.common import AddTimeDimToBatchAndZeroPad + from ray.rllib.core.columns import Columns + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + + # Create a simple dummy class, pretending to be an RLModule with + # `get_initial_state`, `is_stateful` and `model_config` property defined: + class MyStateModule: + # dummy config + model_config = {"max_seq_len": 3} + + def is_stateful(self): + return True + + def get_initial_state(self): + return 0.0 + + + # Create an already reset episode. Expect the connector to add a time-dim to the + # reset observation. + episode = SingleAgentEpisode(observations=[0]) + + rl_module = MyStateModule() + + # Create an instance of this class (as an env-to-module connector). + connector = AddTimeDimToBatchAndZeroPad(as_learner_connector=False) + + # Call the connector. + output_batch = connector( + rl_module=rl_module, + batch={Columns.OBS: [0]}, + episodes=[episode], + shared_data={}, + ) + # The output data's OBS key should now be reshaped to (B, T) + check(output_batch[Columns.OBS], [[0]]) + + # Create a SingleAgentEpisodes containing 5 observations, + # 4 actions and 4 rewards. + episode = SingleAgentEpisode( + observations=[0, 1, 2, 3, 4], + actions=[1, 2, 3, 4], + rewards=[1.0, 2.0, 3.0, 4.0], + len_lookback_buffer=0, + ) + + # Call the connector. + output_batch = connector( + rl_module=rl_module, + batch={Columns.OBS: [4]}, + episodes=[episode], + shared_data={}, + ) + # The output data's OBS, ACTIONS, and REWARDS keys should now all have a time + # rank. + check( + # Expect the episode's last OBS. + output_batch[Columns.OBS], [[4]], + ) + + # Create a new connector as a learner connector with a RNN seq len of 4 (for + # testing purposes only). Passing the same data through this learner connector, + # we expect the data to also be zero-padded. + connector = AddTimeDimToBatchAndZeroPad(as_learner_connector=True) + + # Call the connector. + output_batch = connector( + rl_module=rl_module, + batch={Columns.OBS: {(episode.id_,): [0, 1, 2, 3]}}, + episodes=[episode], + shared_data={}, + ) + check(output_batch[Columns.OBS], {(episode.id_,): [[0, 1, 2], [3, 0, 0]]}) + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + as_learner_connector: bool = False, + **kwargs, + ): + """Initializes a AddObservationsFromEpisodesToBatch instance. + + Args: + as_learner_connector: Whether this connector is part of a Learner connector + pipeline, as opposed to a env-to-module pipeline. As a Learner + connector, it will add an entire Episode's observations (each timestep) + to the batch. + """ + super().__init__( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + **kwargs, + ) + + self._as_learner_connector = as_learner_connector + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + + # If not stateful OR STATE_IN already in data, early out. + if ( + rl_module is None + or not rl_module.is_stateful() + or Columns.STATE_IN in batch + ): + return batch + + # Make all inputs (other than STATE_IN) have an additional T-axis. + # Since data has not been batched yet (we are still operating on lists in the + # batch), we add this time axis as 0 (not 1). When we batch, the batch axis will + # be 0 and the time axis will be 1. + # Also, let module-to-env pipeline know that we had added a single timestep + # time rank to the data (to remove it again). + if not self._as_learner_connector: + for column in batch.keys(): + self.foreach_batch_item_change_in_place( + batch=batch, + column=column, + func=lambda item, eps_id, aid, mid: ( + item + if mid is not None and not rl_module[mid].is_stateful() + # Expand on axis 0 (the to-be-time-dim) if item has not been + # batched yet, otherwise axis=1 (the time-dim). + else tree.map_structure( + lambda s: np.expand_dims( + s, axis=(1 if isinstance(s, BatchedNdArray) else 0) + ), + item, + ) + ), + ) + shared_data["_added_single_ts_time_rank"] = True + else: + # Before adding STATE_IN to the `data`, zero-pad existing data and batch + # into max_seq_len chunks. + for column, column_data in batch.copy().items(): + # Do not zero-pad INFOS column. + if column == Columns.INFOS: + continue + for key, item_list in column_data.items(): + # Multi-agent case AND RLModule is not stateful -> Do not zero-pad + # for this model. + assert isinstance(key, tuple) + mid = None + if len(key) == 3: + eps_id, aid, mid = key + if not rl_module[mid].is_stateful(): + continue + column_data[key] = split_and_zero_pad( + item_list, + max_seq_len=self._get_max_seq_len(rl_module, module_id=mid), + ) + # TODO (sven): Remove this hint/hack once we are not relying on + # SampleBatch anymore (which has to set its property + # zero_padded=True when shuffling). + shared_data[ + ( + "_zero_padded_for_mid=" + f"{mid if mid is not None else DEFAULT_MODULE_ID}" + ) + ] = True + + for sa_episode in self.single_agent_episode_iterator( + # If Learner connector, get all episodes (for train batch). + # If EnvToModule, get only those ongoing episodes that just had their + # agent step (b/c those are the ones we need to compute actions for next). + episodes, + agents_that_stepped_only=False, + ): + # Multi-agent case: Extract correct single agent RLModule (to get its + # individual state). + if sa_episode.module_id is not None: + sa_module = rl_module[sa_episode.module_id] + else: + sa_module = ( + rl_module[DEFAULT_MODULE_ID] + if isinstance(rl_module, MultiRLModule) + else rl_module + ) + # This single-agent RLModule is NOT stateful -> Skip. + if not sa_module.is_stateful(): + continue + + max_seq_len = sa_module.model_config["max_seq_len"] + + # Also, create the loss mask (b/c of our now possibly zero-padded data) + # as well as the seq_lens array and add these to `data` as well. + mask, seq_lens = create_mask_and_seq_lens(len(sa_episode), max_seq_len) + self.add_n_batch_items( + batch=batch, + column=Columns.SEQ_LENS, + items_to_add=seq_lens, + num_items=len(seq_lens), + single_agent_episode=sa_episode, + ) + if not shared_data.get("_added_loss_mask_for_valid_episode_ts"): + self.add_n_batch_items( + batch=batch, + column=Columns.LOSS_MASK, + items_to_add=mask, + num_items=len(mask), + single_agent_episode=sa_episode, + ) + + return batch + + def _get_max_seq_len(self, rl_module, module_id=None): + if not isinstance(rl_module, MultiRLModule): + mod = rl_module + elif module_id: + mod = rl_module[module_id] + else: + mod = next(iter(rl_module.values())) + if "max_seq_len" not in mod.model_config: + raise ValueError( + "You are using a stateful RLModule and are not providing a " + "'max_seq_len' key inside your `model_config`. You can set this " + "dict and/or override keys in it via `config.rl_module(" + "model_config={'max_seq_len': [some int]})`." + ) + return mod.model_config["max_seq_len"] diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/agent_to_module_mapping.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/agent_to_module_mapping.py new file mode 100644 index 0000000000000000000000000000000000000000..63bd78cf9b783804aedb63563e5c5fa80f79de50 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/agent_to_module_mapping.py @@ -0,0 +1,291 @@ +from collections import defaultdict +from typing import Any, Dict, List, Optional + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType, ModuleID +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AgentToModuleMapping(ConnectorV2): + """ConnectorV2 that performs mapping of data from AgentID based to ModuleID based. + + Note: This is one of the default env-to-module or Learner ConnectorV2 pieces that + are added automatically by RLlib into every env-to-module/Learner connector + pipeline, unless `config.add_default_connectors_to_env_to_module_pipeline` or + `config.add_default_connectors_to_learner_pipeline ` are set to + False. + + The default env-to-module connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + The default Learner connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddColumnsFromEpisodesToTrainBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + This connector piece is only used by RLlib (as a default connector piece) in a + multi-agent setup. + + Note that before the mapping, `data` is expected to have the following + structure: + [col0]: + (eps_id0, ag0, mod0): [list of individual batch items] + (eps_id0, ag1, mod2): [list of individual batch items] + (eps_id1, ag0, mod1): [list of individual batch items] + [col1]: + etc.. + + The target structure of the above `data` would then be: + [mod0]: + [col0]: [batched data -> batch_size_B will be the number of all items in the + input data under col0 that have mod0 as their ModuleID] + [col1]: [batched data] + [mod1]: + [col0]: etc. + + Mapping happens in the following stages: + + 1) Under each column name, sort keys first by EpisodeID, then AgentID. + 2) Add ModuleID keys under each column name (no cost/extra memory) and map these + new keys to empty lists. + [col0] -> [mod0] -> []: Then push items that belong to mod0 into these lists. + 3) Perform batching on the per-module lists under each column: + [col0] -> [mod0]: [...] <- now batched data (numpy array or struct of numpy + arrays). + 4) Flip column names with ModuleIDs (no cost/extra memory): + [mod0]: + [col0]: [batched data] + etc.. + + Note that in order to unmap the resulting batch back into an AgentID based one, + we have to store the env vector index AND AgentID of each module's batch item + in an additionally returned `memorized_map_structure`. + + .. testcode:: + + from ray.rllib.connectors.env_to_module import AgentToModuleMapping + from ray.rllib.utils.test_utils import check + + batch = { + "obs": { + ("MA-EPS0", "agent0", "module0"): [0, 1, 2], + ("MA-EPS0", "agent1", "module1"): [3, 4, 5], + }, + "actions": { + ("MA-EPS1", "agent2", "module0"): [8], + ("MA-EPS0", "agent1", "module1"): [9], + }, + } + + # Create our connector piece. + connector = AgentToModuleMapping( + rl_module_specs={"module0", "module1"}, + agent_to_module_mapping_fn=( + lambda agent_id, eps: "module1" if agent_id == "agent1" else "module0" + ), + ) + + # Call the connector (and thereby flip from AgentID based to ModuleID based + # structure.. + output_batch = connector( + rl_module=None, # This particular connector works without an RLModule. + batch=batch, + episodes=[], # This particular connector works without a list of episodes. + explore=True, + shared_data={}, + ) + + # `data` should now be mapped from ModuleIDs to module data. + check( + output_batch, + { + "module0": { + "obs": [0, 1, 2], + "actions": [8], + }, + "module1": { + "obs": [3, 4, 5], + "actions": [9], + }, + }, + ) + """ + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + return self._map_space_if_necessary(input_observation_space, "obs") + + @override(ConnectorV2) + def recompute_output_action_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + return self._map_space_if_necessary(input_action_space, "act") + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + rl_module_specs: Dict[ModuleID, RLModuleSpec], + agent_to_module_mapping_fn, + ): + super().__init__(input_observation_space, input_action_space) + + self._rl_module_specs = rl_module_specs + self._agent_to_module_mapping_fn = agent_to_module_mapping_fn + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # Current agent to module mapping function. + # agent_to_module_mapping_fn = shared_data.get("agent_to_module_mapping_fn") + # Store in shared data, which module IDs map to which episode/agent, such + # that the module-to-env pipeline can map the data back to agents. + memorized_map_structure = defaultdict(list) + for column, agent_data in batch.items(): + if rl_module is not None and column in rl_module: + continue + for eps_id, agent_id, module_id in agent_data.keys(): + memorized_map_structure[module_id].append((eps_id, agent_id)) + # TODO (sven): We should check that all columns have the same struct. + break + + shared_data["memorized_map_structure"] = dict(memorized_map_structure) + + # Mapping from ModuleID to column data. + data_by_module = {} + + # Iterating over each column in the original data: + for column, agent_data in batch.items(): + if rl_module is not None and column in rl_module: + if column in data_by_module: + data_by_module[column].update(agent_data) + else: + data_by_module[column] = agent_data + continue + for ( + eps_id, + agent_id, + module_id, + ), values_batch_or_list in agent_data.items(): + assert isinstance(values_batch_or_list, list) + for value in values_batch_or_list: + if module_id not in data_by_module: + data_by_module[module_id] = {column: []} + elif column not in data_by_module[module_id]: + data_by_module[module_id][column] = [] + + # Append the data. + data_by_module[module_id][column].append(value) + + return data_by_module + + def _map_space_if_necessary(self, space: gym.Space, which: str = "obs"): + # Analyze input observation space to check, whether the user has already taken + # care of the agent to module mapping. + if set(self._rl_module_specs) == set(space.spaces.keys()): + return space + + # We need to take care of agent to module mapping. Figure out the resulting + # observation space here. + dummy_eps = MultiAgentEpisode() + + ret_space = {} + for module_id in self._rl_module_specs: + # Easy way out, user has provided space in the RLModule spec dict. + if ( + isinstance(self._rl_module_specs, dict) + and module_id in self._rl_module_specs + ): + if ( + which == "obs" + and self._rl_module_specs[module_id].observation_space + ): + ret_space[module_id] = self._rl_module_specs[ + module_id + ].observation_space + continue + elif which == "act" and self._rl_module_specs[module_id].action_space: + ret_space[module_id] = self._rl_module_specs[module_id].action_space + continue + + # Need to reverse map spaces (for the different agents) to certain + # module IDs (using a dummy MultiAgentEpisode). + one_space = next(iter(space.spaces.values())) + # If all obs spaces are the same anyway, just use the first + # single-agent space. + if all(s == one_space for s in space.spaces.values()): + ret_space[module_id] = one_space + # Otherwise, we have to compare the ModuleID with all possible + # AgentIDs and find the agent ID that matches. + else: + match_aid = None + one_agent_for_module_found = False + for aid in space.spaces.keys(): + # Match: Assign spaces for this agentID to the PolicyID. + if self._agent_to_module_mapping_fn(aid, dummy_eps) == module_id: + # Make sure, different agents that map to the same + # policy don't have different spaces. + if ( + module_id in ret_space + and space[aid] != ret_space[module_id] + ): + raise ValueError( + f"Two agents ({aid} and {match_aid}) in your " + "environment map to the same ModuleID (as per your " + "`agent_to_module_mapping_fn`), however, these agents " + "also have different observation spaces as per the env!" + ) + ret_space[module_id] = space[aid] + match_aid = aid + one_agent_for_module_found = True + # Still no space found for this module ID -> Error out. + if not one_agent_for_module_found: + raise ValueError( + f"Could not find or derive any {which}-space for RLModule " + f"{module_id}! This can happen if your `config.rl_module(rl_" + f"module_spec=...)` does NOT contain space information for this" + " particular single-agent module AND your agent-to-module-" + "mapping function is stochastic (such that for some agent A, " + "more than one ModuleID might be returned somewhat randomly). " + f"Fix this error by providing {which}-space information using " + "`config.rl_module(rl_module_spec=MultiRLModuleSpec(" + f"rl_module_specs={{'{module_id}': RLModuleSpec(" + "observation_space=..., action_space=...)}}))" + ) + + return gym.spaces.Dict(ret_space) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/batch_individual_items.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/batch_individual_items.py new file mode 100644 index 0000000000000000000000000000000000000000..8d3bd5a53cd7aa560d139145135f3a1f679b6238 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/batch_individual_items.py @@ -0,0 +1,200 @@ +from typing import Any, Dict, List, Optional + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.spaces.space_utils import batch as batch_fn +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class BatchIndividualItems(ConnectorV2): + """Batches individual data-items (in lists) into tensors (with batch dimension). + + Note: This is one of the default env-to-module or Learner ConnectorV2 pieces that + are added automatically by RLlib into every env-to-module/Learner connector + pipeline, unless `config.add_default_connectors_to_env_to_module_pipeline` or + `config.add_default_connectors_to_learner_pipeline ` are set to + False. + + The default env-to-module connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + The default Learner connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddColumnsFromEpisodesToTrainBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + This ConnectorV2: + - Operates only on the input `data`, NOT the incoming list of episode objects + (ignored). + - In the single-agent case, `data` must already be a dict, structured as follows by + prior connector pieces of the same pipeline: + [col0] -> {[(eps_id,)]: [list of individual batch items]} + - In the multi-agent case, `data` must already be a dict, structured as follows by + prior connector pieces of the same pipeline (in particular the + `AgentToModuleMapping` piece): + [module_id] -> [col0] -> [list of individual batch items] + - Translates the above data under the different columns (e.g. "obs") into final + (batched) structures. For the single-agent case, the output `data` looks like this: + [col0] -> [possibly complex struct of batches (at the leafs)]. + For the multi-agent case, the output `data` looks like this: + [module_id] -> [col0] -> [possibly complex struct of batches (at the leafs)]. + + .. testcode:: + + from ray.rllib.connectors.common import BatchIndividualItems + from ray.rllib.utils.test_utils import check + + single_agent_batch = { + "obs": { + # Note that at this stage, next-obs is not part of the data anymore .. + ("MA-EPS0",): [0, 1], + ("MA-EPS1",): [2, 3], + }, + "actions": { + # .. so we have as many actions per episode as we have observations. + ("MA-EPS0",): [4, 5], + ("MA-EPS1",): [6, 7], + }, + } + + # Create our (single-agent) connector piece. + connector = BatchIndividualItems() + + # Call the connector (and thereby batch the individual items). + output_batch = connector( + rl_module=None, # This particular connector works without an RLModule. + batch=single_agent_batch, + episodes=[], # This particular connector works without a list of episodes. + explore=True, + shared_data={}, + ) + + # `output_batch` should now be batched (episode IDs should have been removed + # from the struct). + check( + output_batch, + {"obs": [0, 1, 2, 3], "actions": [4, 5, 6, 7]}, + ) + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + multi_agent: bool = False, + **kwargs, + ): + """Initializes a BatchIndividualItems instance. + + Args: + multi_agent: Whether this is a connector operating on a multi-agent + observation space mapping AgentIDs to individual agents' observations. + """ + super().__init__( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + **kwargs, + ) + self._multi_agent = multi_agent + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + is_multi_rl_module = isinstance(rl_module, MultiRLModule) + + # Convert lists of individual items into properly batched data. + for column, column_data in batch.copy().items(): + # Multi-agent case: This connector piece should only be used after(!) + # the AgentToModuleMapping connector has already been applied, leading + # to a batch structure of: + # [module_id] -> [col0] -> [list of individual batch items] + if is_multi_rl_module and column in rl_module: + # Case, in which a column has already been properly batched before this + # connector piece is called. + if not self._multi_agent: + continue + # If MA Off-Policy and independent sampling we need to overcome this + # check. + module_data = column_data + for col, col_data in module_data.copy().items(): + if isinstance(col_data, list) and col != Columns.INFOS: + module_data[col] = batch_fn( + col_data, + individual_items_already_have_batch_dim="auto", + ) + + # Simple case: There is a list directly under `column`: + # Batch the list. + elif isinstance(column_data, list): + batch[column] = batch_fn( + column_data, + individual_items_already_have_batch_dim="auto", + ) + + # Single-agent case: There is a dict under `column` mapping + # `eps_id` to lists of items: + # Concat all these lists, then batch. + elif not self._multi_agent: + # TODO: only really need this in non-Learner connector pipeline + memorized_map_structure = [] + list_to_be_batched = [] + for (eps_id,) in column_data.keys(): + for item in column_data[(eps_id,)]: + # Only record structure for OBS column. + if column == Columns.OBS: + memorized_map_structure.append(eps_id) + list_to_be_batched.append(item) + # INFOS should not be batched (remain a list). + batch[column] = ( + list_to_be_batched + if column == Columns.INFOS + else batch_fn( + list_to_be_batched, + individual_items_already_have_batch_dim="auto", + ) + ) + if is_multi_rl_module: + if DEFAULT_MODULE_ID not in batch: + batch[DEFAULT_MODULE_ID] = {} + batch[DEFAULT_MODULE_ID][column] = batch.pop(column) + + # Only record structure for OBS column. + if column == Columns.OBS: + shared_data["memorized_map_structure"] = memorized_map_structure + # Multi-agent case: But Module ID not found in our RLModule -> Ignore this + # `module_id` entirely. + # else: + # pass + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/flatten_observations.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/flatten_observations.py new file mode 100644 index 0000000000000000000000000000000000000000..51dc8e8c038f565edb2b658238acfb0b87d61ff6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/flatten_observations.py @@ -0,0 +1,433 @@ +from typing import Any, Collection, Dict, List, Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.numpy import flatten_inputs_to_1d_tensor +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.typing import AgentID, EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class FlattenObservations(ConnectorV2): + """A connector piece that flattens all observation components into a 1D array. + + - Can be used either in env-to-module or learner pipelines. + - When used in env-to-module pipelines: + - Works directly on the incoming episodes list and changes the last observation + in-place (write the flattened observation back into the episode). + - This connector does NOT alter the incoming batch (`data`) when called. + - When used in learner pipelines: + Works directly on the incoming episodes list and changes all observations + before stacking them into the batch. + + + .. testcode:: + + import gymnasium as gym + import numpy as np + + from ray.rllib.connectors.env_to_module import FlattenObservations + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Some arbitrarily nested, complex observation space. + obs_space = gym.spaces.Dict({ + "a": gym.spaces.Box(-10.0, 10.0, (), np.float32), + "b": gym.spaces.Tuple([ + gym.spaces.Discrete(2), + gym.spaces.Box(-1.0, 1.0, (2, 1), np.float32), + ]), + "c": gym.spaces.MultiDiscrete([2, 3]), + }) + act_space = gym.spaces.Discrete(2) + + # Two example episodes, both with initial (reset) observations coming from the + # above defined observation space. + episode_1 = SingleAgentEpisode( + observations=[ + { + "a": np.array(-10.0, np.float32), + "b": (1, np.array([[-1.0], [-1.0]], np.float32)), + "c": np.array([0, 2]), + }, + ], + ) + episode_2 = SingleAgentEpisode( + observations=[ + { + "a": np.array(10.0, np.float32), + "b": (0, np.array([[1.0], [1.0]], np.float32)), + "c": np.array([1, 1]), + }, + ], + ) + + # Construct our connector piece. + connector = FlattenObservations(obs_space, act_space) + + # Call our connector piece with the example data. + output_batch = connector( + rl_module=None, # This connector works without an RLModule. + batch={}, # This connector does not alter the input batch. + episodes=[episode_1, episode_2], + explore=True, + shared_data={}, + ) + + # The connector does not alter the data and acts as pure pass-through. + check(output_batch, {}) + + # The connector has flattened each item in the episodes to a 1D tensor. + check( + episode_1.get_observations(0), + # box() disc(2). box(2, 1). multidisc(2, 3)........ + np.array([-10.0, 0.0, 1.0, -1.0, -1.0, 1.0, 0.0, 0.0, 0.0, 1.0]), + ) + check( + episode_2.get_observations(0), + # box() disc(2). box(2, 1). multidisc(2, 3)........ + np.array([10.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0]), + ) + + # Two example episodes, both with initial (reset) observations coming from the + # above defined observation space. + episode_1 = SingleAgentEpisode( + observations=[ + { + "a": np.array(-10.0, np.float32), + "b": (1, np.array([[-1.0], [-1.0]], np.float32)), + "c": np.array([0, 2]), + }, + ], + ) + episode_2 = SingleAgentEpisode( + observations=[ + { + "a": np.array(10.0, np.float32), + "b": (0, np.array([[1.0], [1.0]], np.float32)), + "c": np.array([1, 1]), + }, + ], + ) + + # Construct our connector piece and remove in it the "a" (Box()) key-value + # pair from the dictionary observations. + connector = FlattenObservations(obs_space, act_space, keys_to_remove=["a"]) + + # Call our connector piece with the example data. + output_batch = connector( + rl_module=None, # This connector works without an RLModule. + batch={}, # This connector does not alter the input batch. + episodes=[episode_1, episode_2], + explore=True, + shared_data={}, + ) + + # The connector has flattened each item in the episodes to a 1D tensor + # and removed the "a" (Box()) key-value pair. + check( + episode_1.get_observations(0), + # disc(2). box(2, 1). multidisc(2, 3)........ + np.array([0.0, 1.0, -1.0, -1.0, 1.0, 0.0, 0.0, 0.0, 1.0]), + ) + check( + episode_2.get_observations(0), + # disc(2). box(2, 1). multidisc(2, 3)........ + np.array([1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0]), + ) + + # Use the connector in a learner pipeline. Note, we need here two + # observations in the episode because the agent has to have stepped + # at least once. + episode_1 = SingleAgentEpisode( + observations=[ + { + "a": np.array(-10.0, np.float32), + "b": (1, np.array([[-1.0], [-1.0]], np.float32)), + "c": np.array([0, 2]), + }, + { + "a": np.array(-10.0, np.float32), + "b": (1, np.array([[-1.0], [-1.0]], np.float32)), + "c": np.array([0, 2]), + }, + ], + actions=[1], + rewards=[0], + # Set the length of the lookback buffer to 0 to read the data as + # from an actual step. + len_lookback_buffer=0, + ) + episode_2 = SingleAgentEpisode( + observations=[ + { + "a": np.array(10.0, np.float32), + "b": (0, np.array([[1.0], [1.0]], np.float32)), + "c": np.array([1, 1]), + }, + { + "a": np.array(10.0, np.float32), + "b": (0, np.array([[1.0], [1.0]], np.float32)), + "c": np.array([1, 1]), + }, + + ], + actions=[1], + rewards=[0], + # Set the length of the lookback buffer to 0 to read the data as + # from an actual step. + len_lookback_buffer=0, + ) + + # Construct our connector piece for a learner pipeline and remove the + # "a" (Box()) key-value pair. + connector = FlattenObservations( + obs_space, + act_space, + as_learner_connector=True, + keys_to_remove=["a"] + ) + + # Call our connector piece with the example data. + output_batch = connector( + rl_module=None, # This connector works without an RLModule. + batch={}, # This connector does not alter the input batch. + episodes=[episode_1, episode_2], + explore=True, + shared_data={}, + ) + + check(list(output_batch.keys()), ["obs"]) + check(list(output_batch["obs"].keys()), [(episode_1.id_,), (episode_2.id_,)]) + + check( + output_batch["obs"][(episode_1.id_,)][0][0], + np.array([0.0, 1.0, -1.0, -1.0, 1.0, 0.0, 0.0, 0.0, 1.0]), + ) + check( + output_batch["obs"][(episode_2.id_,)][0][0], + np.array([1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0]), + ) + """ + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space, + input_action_space, + ) -> gym.Space: + self._input_obs_base_struct = get_base_struct_from_space( + self.input_observation_space + ) + + if self._multi_agent: + spaces = {} + for agent_id, space in self._input_obs_base_struct.items(): + # Remove keys, if necessary. + # TODO (simon): Maybe allow to remove different keys for different agents. + if self._keys_to_remove: + self._input_obs_base_struct[agent_id] = { + k: v + for k, v in self._input_obs_base_struct[agent_id].items() + if k not in self._keys_to_remove + } + if self._agent_ids and agent_id not in self._agent_ids: + spaces[agent_id] = self._input_obs_base_struct[agent_id] + else: + sample = flatten_inputs_to_1d_tensor( + tree.map_structure( + lambda s: s.sample(), + self._input_obs_base_struct[agent_id], + ), + self._input_obs_base_struct[agent_id], + batch_axis=False, + ) + spaces[agent_id] = Box( + float("-inf"), float("inf"), (len(sample),), np.float32 + ) + return gym.spaces.Dict(spaces) + else: + # Remove keys, if necessary. + if self._keys_to_remove: + self._input_obs_base_struct = { + k: v + for k, v in self._input_obs_base_struct.items() + if k not in self._keys_to_remove + } + sample = flatten_inputs_to_1d_tensor( + tree.map_structure( + lambda s: s.sample(), + self._input_obs_base_struct, + ), + self._input_obs_base_struct, + batch_axis=False, + ) + return Box(float("-inf"), float("inf"), (len(sample),), np.float32) + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + multi_agent: bool = False, + agent_ids: Optional[Collection[AgentID]] = None, + as_learner_connector: bool = False, + keys_to_remove: Optional[List[str]] = None, + **kwargs, + ): + """Initializes a FlattenObservations instance. + + Args: + multi_agent: Whether this connector operates on multi-agent observations, + in which case, the top-level of the Dict space (where agent IDs are + mapped to individual agents' observation spaces) is left as-is. + agent_ids: If multi_agent is True, this argument defines a collection of + AgentIDs for which to flatten. AgentIDs not in this collection are + ignored. + If None, flatten observations for all AgentIDs. None is the default. + as_learner_connector: Whether this connector is part of a Learner connector + pipeline, as opposed to an env-to-module pipeline. + Note, this is usually only used for offline rl where the data comes + from an offline dataset instead of a simulator. With a simulator the + data is simply rewritten. + keys_to_remove: Optional keys to remove from the observations. + + """ + self._input_obs_base_struct = None + self._multi_agent = multi_agent + self._agent_ids = agent_ids + self._as_learner_connector = as_learner_connector + + assert keys_to_remove is None or ( + keys_to_remove + and isinstance(input_observation_space, gym.spaces.Dict) + or ( + multi_agent + and any( + isinstance(agent_space, gym.spaces.Dict) + for agent_space in self.input_observation_space + ) + ) + ), "When using `keys_to_remove` the observation space must be of type `gym.spaces.Dict`." + self._keys_to_remove = keys_to_remove or [] + + super().__init__(input_observation_space, input_action_space, **kwargs) + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + if self._as_learner_connector: + for sa_episode in self.single_agent_episode_iterator( + episodes, agents_that_stepped_only=True + ): + + def _map_fn(obs, _sa_episode=sa_episode): + # Remove keys, if necessary. + obs = [self._remove_keys_from_dict(o, sa_episode) for o in obs] + + batch_size = len(sa_episode) + flattened_obs = flatten_inputs_to_1d_tensor( + inputs=obs, + # In the multi-agent case, we need to use the specific agent's + # space struct, not the multi-agent observation space dict. + spaces_struct=self._input_obs_base_struct, + # Our items are individual observations (no batch axis present). + batch_axis=False, + ) + return flattened_obs.reshape(batch_size, -1).copy() + + self.add_n_batch_items( + batch=batch, + column=Columns.OBS, + items_to_add=_map_fn( + sa_episode.get_observations(indices=slice(0, len(sa_episode))), + sa_episode, + ), + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + else: + for sa_episode in self.single_agent_episode_iterator( + episodes, agents_that_stepped_only=True + ): + last_obs = sa_episode.get_observations(-1) + # Remove keys, if necessary. + last_obs = self._remove_keys_from_dict(last_obs, sa_episode) + + if self._multi_agent: + if ( + self._agent_ids is not None + and sa_episode.agent_id not in self._agent_ids + ): + flattened_obs = last_obs + else: + flattened_obs = flatten_inputs_to_1d_tensor( + inputs=last_obs, + # In the multi-agent case, we need to use the specific agent's + # space struct, not the multi-agent observation space dict. + spaces_struct=self._input_obs_base_struct[ + sa_episode.agent_id + ], + # Our items are individual observations (no batch axis present). + batch_axis=False, + ) + else: + flattened_obs = flatten_inputs_to_1d_tensor( + inputs=last_obs, + spaces_struct=self._input_obs_base_struct, + # Our items are individual observations (no batch axis present). + batch_axis=False, + ) + + # Write new observation directly back into the episode. + sa_episode.set_observations(at_indices=-1, new_data=flattened_obs) + # We set the Episode's observation space to ours so that we can safely + # set the last obs to the new value (without causing a space mismatch + # error). + sa_episode.observation_space = self.observation_space + + return batch + + def _remove_keys_from_dict(self, obs, sa_episode): + """Removes keys from dictionary spaces. + + Args: + obs: Observation sample from space. + sa_episode: Single-agent episode. Needs `agent_id` set in multi-agent + setups. + + Returns: + Observation sample `obs` with keys in `self._keys_to_remove` removed. + """ + + # Only remove keys for agents that have a dictionary space. + is_dict_space = False + if self._multi_agent: + is_dict_space = isinstance( + self.input_observation_space[sa_episode.agent_id], gym.spaces.Dict + ) + else: + is_dict_space = isinstance(self.input_observation_space, gym.spaces.Dict) + + # Remove keys, if necessary. + if is_dict_space and self._keys_to_remove: + obs = {k: v for k, v in obs.items() if k not in self._keys_to_remove} + + return obs diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/frame_stacking.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/frame_stacking.py new file mode 100644 index 0000000000000000000000000000000000000000..a99b145bd5f5a02e970ea038bf93e8e8a5a5286a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/frame_stacking.py @@ -0,0 +1,149 @@ +from typing import Any, Dict, List, Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class FrameStacking(ConnectorV2): + """A connector piece that stacks the previous n observations into one.""" + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + # Change our observation space according to the given stacking settings. + if self._multi_agent: + ret = {} + for agent_id, obs_space in input_observation_space.spaces.items(): + ret[agent_id] = self._convert_individual_space(obs_space) + return gym.spaces.Dict(ret) + else: + return self._convert_individual_space(input_observation_space) + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + num_frames: int = 1, + multi_agent: bool = False, + as_learner_connector: bool = False, + **kwargs, + ): + """Initializes a FrameStackingConnector instance. + + Args: + num_frames: The number of observation frames to stack up (into a single + observation) for the RLModule's forward pass. + multi_agent: Whether this is a connector operating on a multi-agent + observation space mapping AgentIDs to individual agents' observations. + as_learner_connector: Whether this connector is part of a Learner connector + pipeline, as opposed to an env-to-module pipeline. + """ + super().__init__( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + **kwargs, + ) + + self._multi_agent = multi_agent + self.num_frames = num_frames + self._as_learner_connector = as_learner_connector + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # Learner connector pipeline. Episodes have been numpy'ized. + if self._as_learner_connector: + for sa_episode in self.single_agent_episode_iterator( + episodes, agents_that_stepped_only=False + ): + + def _map_fn(s, _sa_episode=sa_episode): + # Squeeze out last dim. + s = np.squeeze(s, axis=-1) + # Calculate new shape and strides + new_shape = (len(_sa_episode), self.num_frames) + s.shape[1:] + new_strides = (s.strides[0],) + s.strides + # Create a strided view of the array. + # But return a copy to avoid non-contiguous memory in the object + # store (which is very expensive to deserialize). + return np.transpose( + np.lib.stride_tricks.as_strided( + s, shape=new_shape, strides=new_strides + ), + axes=[0, 2, 3, 1], + ).copy() + + # Get all observations from the episode in one np array (except for + # the very last one, which is the final observation not needed for + # learning). + self.add_n_batch_items( + batch=batch, + column=Columns.OBS, + items_to_add=tree.map_structure( + _map_fn, + sa_episode.get_observations( + indices=slice(-self.num_frames + 1, len(sa_episode)), + neg_index_as_lookback=True, + fill=0.0, + ), + ), + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + # Env-to-module pipeline. Episodes still operate on lists. + else: + for sa_episode in self.single_agent_episode_iterator(episodes): + assert not sa_episode.is_numpy + # Get the list of observations to stack. + obs_stack = sa_episode.get_observations( + indices=slice(-self.num_frames, None), + fill=0.0, + ) + # Observation components are (w, h, 1) + # -> concatenate along axis=-1 to (w, h, [num_frames]). + stacked_obs = tree.map_structure( + lambda *s: np.concatenate(s, axis=2), + *obs_stack, + ) + self.add_batch_item( + batch=batch, + column=Columns.OBS, + item_to_add=stacked_obs, + single_agent_episode=sa_episode, + ) + + return batch + + def _convert_individual_space(self, obs_space): + # Some assumptions: Space is box AND last dim (the stacking one) is 1. + assert isinstance(obs_space, gym.spaces.Box), obs_space + assert obs_space.shape[-1] == 1, obs_space + + return gym.spaces.Box( + low=np.repeat(obs_space.low, repeats=self.num_frames, axis=-1), + high=np.repeat(obs_space.high, repeats=self.num_frames, axis=-1), + shape=list(obs_space.shape)[:-1] + [self.num_frames], + dtype=obs_space.dtype, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/module_to_agent_unmapping.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/module_to_agent_unmapping.py new file mode 100644 index 0000000000000000000000000000000000000000..fb3a2b1e954e98ee89a455c53d70304d2cd9b7f0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/module_to_agent_unmapping.py @@ -0,0 +1,48 @@ +from collections import defaultdict +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class ModuleToAgentUnmapping(ConnectorV2): + """Performs flipping of `data` from ModuleID- to AgentID based mapping. + + Before mapping: + data[module1] -> [col, e.g. ACTIONS] + -> [dict mapping episode-identifying tuples to lists of data] + data[module2] -> ... + + After mapping: + data[ACTIONS]: [dict mapping episode-identifying tuples to lists of data] + + Note that episode-identifying tuples have the form of: (episode_id,) in the + single-agent case and (ma_episode_id, agent_id, module_id) in the multi-agent + case. + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # This Connector should only be used in a multi-agent setting. + assert isinstance(episodes[0], MultiAgentEpisode) + + agent_data = defaultdict(dict) + for module_id, module_data in batch.items(): + for column, values_dict in module_data.items(): + agent_data[column].update(values_dict) + + return dict(agent_data) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/numpy_to_tensor.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/numpy_to_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..2e4c954dd6a783879f1208727a15c6f3948f4bc9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/numpy_to_tensor.py @@ -0,0 +1,125 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.rllib.utils.typing import DeviceType + + +@PublicAPI(stability="alpha") +class NumpyToTensor(ConnectorV2): + """Converts numpy arrays across the entire input data into (framework) tensors. + + The framework information is received via the provided `rl_module` arg in the + `__call__()` method. + + Note: This is one of the default env-to-module or Learner ConnectorV2 pieces that + are added automatically by RLlib into every env-to-module/Learner connector + pipeline, unless `config.add_default_connectors_to_env_to_module_pipeline` or + `config.add_default_connectors_to_learner_pipeline ` are set to + False. + + The default env-to-module connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + The default Learner connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddColumnsFromEpisodesToTrainBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + This ConnectorV2: + - Loops through the input `data` and converts all found numpy arrays into + framework-specific tensors (possibly on a GPU). + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + pin_memory: bool = False, + device: Optional["DeviceType"] = None, + **kwargs, + ): + """Initializes a NumpyToTensor instance. + + Args: + pin_memory: Whether to pin memory when creating (torch) tensors. + If None (default), pins memory if `as_learner_connector` is True, + otherwise doesn't pin memory. + device: An optional device to move the resulting tensors to. If not + provided, all data will be left on the CPU. + **kwargs: + """ + super().__init__( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + **kwargs, + ) + self._pin_memory = pin_memory + self._device = device + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + is_single_agent = False + is_multi_rl_module = isinstance(rl_module, MultiRLModule) + # `data` already a ModuleID to batch mapping format. + if not (is_multi_rl_module and all(c in rl_module._rl_modules for c in batch)): + is_single_agent = True + batch = {DEFAULT_MODULE_ID: batch} + + for module_id, module_data in batch.copy().items(): + # If `rl_module` is None, leave data in numpy format. + if rl_module is not None: + infos = module_data.pop(Columns.INFOS, None) + if rl_module.framework == "torch": + module_data = convert_to_torch_tensor( + module_data, pin_memory=self._pin_memory, device=self._device + ) + else: + raise ValueError( + "`NumpyToTensor`does NOT support frameworks other than torch!" + ) + if infos is not None: + module_data[Columns.INFOS] = infos + + # Early out with data under(!) `DEFAULT_MODULE_ID`, b/c we are in plain + # single-agent mode. + if is_single_agent: + return module_data + batch[module_id] = module_data + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/common/tensor_to_numpy.py b/lib/python3.12/site-packages/ray/rllib/connectors/common/tensor_to_numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..f6bbb2669c7942087e2e9fc1817a9dd3d48e1280 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/common/tensor_to_numpy.py @@ -0,0 +1,26 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class TensorToNumpy(ConnectorV2): + """Converts (framework) tensors across the entire input data into numpy arrays.""" + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + return convert_to_numpy(batch) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/connector.py b/lib/python3.12/site-packages/ray/rllib/connectors/connector.py new file mode 100644 index 0000000000000000000000000000000000000000..2a8a6a51911bfce76219d92e259cbfaaddfad2fb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/connector.py @@ -0,0 +1,478 @@ +"""This file defines base types and common structures for RLlib connectors. +""" + +import abc +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union + +import gymnasium as gym + +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import ( + ActionConnectorDataType, + AgentConnectorDataType, + AlgorithmConfigDict, + TensorType, +) + +if TYPE_CHECKING: + from ray.rllib.policy.policy import Policy + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class ConnectorContext: + """Data bits that may be needed for running connectors. + + Note(jungong) : we need to be really careful with the data fields here. + E.g., everything needs to be serializable, in case we need to fetch them + in a remote setting. + """ + + # TODO(jungong) : figure out how to fetch these in a remote setting. + # Probably from a policy server when initializing a policy client. + + def __init__( + self, + config: AlgorithmConfigDict = None, + initial_states: List[TensorType] = None, + observation_space: gym.Space = None, + action_space: gym.Space = None, + view_requirements: Dict[str, ViewRequirement] = None, + is_policy_recurrent: bool = False, + ): + """Construct a ConnectorContext instance. + + Args: + initial_states: States that are used for constructing + the initial input dict for RNN models. [] if a model is not recurrent. + action_space_struct: a policy's action space, in python + data format. E.g., python dict instead of DictSpace, python tuple + instead of TupleSpace. + """ + self.config = config or {} + self.initial_states = initial_states or [] + self.observation_space = observation_space + self.action_space = action_space + self.view_requirements = view_requirements + self.is_policy_recurrent = is_policy_recurrent + + @staticmethod + def from_policy(policy: "Policy") -> "ConnectorContext": + """Build ConnectorContext from a given policy. + + Args: + policy: Policy + + Returns: + A ConnectorContext instance. + """ + return ConnectorContext( + config=policy.config, + initial_states=policy.get_initial_state(), + observation_space=policy.observation_space, + action_space=policy.action_space, + view_requirements=policy.view_requirements, + is_policy_recurrent=policy.is_recurrent(), + ) + + +@OldAPIStack +class Connector(abc.ABC): + """Connector base class. + + A connector is a step of transformation, of either envrionment data before they + get to a policy, or policy output before it is sent back to the environment. + + Connectors may be training-aware, for example, behave slightly differently + during training and inference. + + All connectors are required to be serializable and implement to_state(). + """ + + def __init__(self, ctx: ConnectorContext): + # Default is training mode. + self._is_training = True + + def in_training(self): + self._is_training = True + + def in_eval(self): + self._is_training = False + + def __str__(self, indentation: int = 0): + return " " * indentation + self.__class__.__name__ + + def to_state(self) -> Tuple[str, Any]: + """Serialize a connector into a JSON serializable Tuple. + + to_state is required, so that all Connectors are serializable. + + Returns: + A tuple of connector's name and its serialized states. + String should match the name used to register the connector, + while state can be any single data structure that contains the + serialized state of the connector. If a connector is stateless, + state can simply be None. + """ + # Must implement by each connector. + return NotImplementedError + + @staticmethod + def from_state(ctx: ConnectorContext, params: Any) -> "Connector": + """De-serialize a JSON params back into a Connector. + + from_state is required, so that all Connectors are serializable. + + Args: + ctx: Context for constructing this connector. + params: Serialized states of the connector to be recovered. + + Returns: + De-serialized connector. + """ + # Must implement by each connector. + return NotImplementedError + + +@OldAPIStack +class AgentConnector(Connector): + """Connector connecting user environments to RLlib policies. + + An agent connector transforms a list of agent data in AgentConnectorDataType + format into a new list in the same AgentConnectorDataTypes format. + The input API is designed so agent connectors can have access to all the + agents assigned to a particular policy. + + AgentConnectorDataTypes can be used to specify arbitrary type of env data, + + Example: + + Represent a list of agent data from one env step() call. + + .. testcode:: + + import numpy as np + ac = AgentConnectorDataType( + env_id="env_1", + agent_id=None, + data={ + "agent_1": np.array([1, 2, 3]), + "agent_2": np.array([4, 5, 6]), + } + ) + + Or a single agent data ready to be preprocessed. + + .. testcode:: + + ac = AgentConnectorDataType( + env_id="env_1", + agent_id="agent_1", + data=np.array([1, 2, 3]), + ) + + We can also adapt a simple stateless function into an agent connector by + using register_lambda_agent_connector: + + .. testcode:: + + import numpy as np + from ray.rllib.connectors.agent.lambdas import ( + register_lambda_agent_connector + ) + TimesTwoAgentConnector = register_lambda_agent_connector( + "TimesTwoAgentConnector", lambda data: data * 2 + ) + + # More complicated agent connectors can be implemented by extending this + # AgentConnector class: + + class FrameSkippingAgentConnector(AgentConnector): + def __init__(self, n): + self._n = n + self._frame_count = default_dict(str, default_dict(str, int)) + + def reset(self, env_id: str): + del self._frame_count[env_id] + + def __call__( + self, ac_data: List[AgentConnectorDataType] + ) -> List[AgentConnectorDataType]: + ret = [] + for d in ac_data: + assert d.env_id and d.agent_id, "Skipping works per agent!" + + count = self._frame_count[ac_data.env_id][ac_data.agent_id] + self._frame_count[ac_data.env_id][ac_data.agent_id] = ( + count + 1 + ) + + if count % self._n == 0: + ret.append(d) + return ret + + As shown, an agent connector may choose to emit an empty list to stop input + observations from being further prosessed. + """ + + def reset(self, env_id: str): + """Reset connector state for a specific environment. + + For example, at the end of an episode. + + Args: + env_id: required. ID of a user environment. Required. + """ + pass + + def on_policy_output(self, output: ActionConnectorDataType): + """Callback on agent connector of policy output. + + This is useful for certain connectors, for example RNN state buffering, + where the agent connect needs to be aware of the output of a policy + forward pass. + + Args: + ctx: Context for running this connector call. + output: Env and agent IDs, plus data output from policy forward pass. + """ + pass + + def __call__( + self, acd_list: List[AgentConnectorDataType] + ) -> List[AgentConnectorDataType]: + """Transform a list of data items from env before they reach policy. + + Args: + ac_data: List of env and agent IDs, plus arbitrary data items from + an environment or upstream agent connectors. + + Returns: + A list of transformed data items in AgentConnectorDataType format. + The shape of a returned list does not have to match that of the input list. + An AgentConnector may choose to derive multiple outputs for a single piece + of input data, for example multi-agent obs -> multiple single agent obs. + Agent connectors may also choose to skip emitting certain inputs, + useful for connectors such as frame skipping. + """ + assert isinstance( + acd_list, (list, tuple) + ), "Input to agent connectors are list of AgentConnectorDataType." + # Default implementation. Simply call transform on each agent connector data. + return [self.transform(d) for d in acd_list] + + def transform(self, ac_data: AgentConnectorDataType) -> AgentConnectorDataType: + """Transform a single agent connector data item. + + Args: + data: Env and agent IDs, plus arbitrary data item from a single agent + of an environment. + + Returns: + A transformed piece of agent connector data. + """ + raise NotImplementedError + + +@OldAPIStack +class ActionConnector(Connector): + """Action connector connects policy outputs including actions, + to user environments. + + An action connector transforms a single piece of policy output in + ActionConnectorDataType format, which is basically PolicyOutputType plus env and + agent IDs. + + Any functions that operate directly on PolicyOutputType can be easily adapted + into an ActionConnector by using register_lambda_action_connector. + + Example: + + .. testcode:: + + from ray.rllib.connectors.action.lambdas import ( + register_lambda_action_connector + ) + ZeroActionConnector = register_lambda_action_connector( + "ZeroActionsConnector", + lambda actions, states, fetches: ( + np.zeros_like(actions), states, fetches + ) + ) + + More complicated action connectors can also be implemented by sub-classing + this ActionConnector class. + """ + + def __call__(self, ac_data: ActionConnectorDataType) -> ActionConnectorDataType: + """Transform policy output before they are sent to a user environment. + + Args: + ac_data: Env and agent IDs, plus policy output. + + Returns: + The processed action connector data. + """ + return self.transform(ac_data) + + def transform(self, ac_data: ActionConnectorDataType) -> ActionConnectorDataType: + """Implementation of the actual transform. + + Users should override transform instead of __call__ directly. + + Args: + ac_data: Env and agent IDs, plus policy output. + + Returns: + The processed action connector data. + """ + raise NotImplementedError + + +@OldAPIStack +class ConnectorPipeline(abc.ABC): + """Utility class for quick manipulation of a connector pipeline.""" + + def __init__(self, ctx: ConnectorContext, connectors: List[Connector]): + self.connectors = connectors + + def in_training(self): + for c in self.connectors: + c.in_training() + + def in_eval(self): + for c in self.connectors: + c.in_eval() + + def remove(self, name: str): + """Remove a connector by + + Args: + name: name of the connector to be removed. + """ + idx = -1 + for i, c in enumerate(self.connectors): + if c.__class__.__name__ == name: + idx = i + break + if idx >= 0: + del self.connectors[idx] + logger.info(f"Removed connector {name} from {self.__class__.__name__}.") + else: + logger.warning(f"Trying to remove a non-existent connector {name}.") + + def insert_before(self, name: str, connector: Connector): + """Insert a new connector before connector + + Args: + name: name of the connector before which a new connector + will get inserted. + connector: a new connector to be inserted. + """ + idx = -1 + for idx, c in enumerate(self.connectors): + if c.__class__.__name__ == name: + break + if idx < 0: + raise ValueError(f"Can not find connector {name}") + self.connectors.insert(idx, connector) + + logger.info( + f"Inserted {connector.__class__.__name__} before {name} " + f"to {self.__class__.__name__}." + ) + + def insert_after(self, name: str, connector: Connector): + """Insert a new connector after connector + + Args: + name: name of the connector after which a new connector + will get inserted. + connector: a new connector to be inserted. + """ + idx = -1 + for idx, c in enumerate(self.connectors): + if c.__class__.__name__ == name: + break + if idx < 0: + raise ValueError(f"Can not find connector {name}") + self.connectors.insert(idx + 1, connector) + + logger.info( + f"Inserted {connector.__class__.__name__} after {name} " + f"to {self.__class__.__name__}." + ) + + def prepend(self, connector: Connector): + """Append a new connector at the beginning of a connector pipeline. + + Args: + connector: a new connector to be appended. + """ + self.connectors.insert(0, connector) + + logger.info( + f"Added {connector.__class__.__name__} to the beginning of " + f"{self.__class__.__name__}." + ) + + def append(self, connector: Connector): + """Append a new connector at the end of a connector pipeline. + + Args: + connector: a new connector to be appended. + """ + self.connectors.append(connector) + + logger.info( + f"Added {connector.__class__.__name__} to the end of " + f"{self.__class__.__name__}." + ) + + def __str__(self, indentation: int = 0): + return "\n".join( + [" " * indentation + self.__class__.__name__] + + [c.__str__(indentation + 4) for c in self.connectors] + ) + + def __getitem__(self, key: Union[str, int, type]): + """Returns a list of connectors that fit 'key'. + + If key is a number n, we return a list with the nth element of this pipeline. + If key is a Connector class or a string matching the class name of a + Connector class, we return a list of all connectors in this pipeline matching + the specified class. + + Args: + key: The key to index by + + Returns: The Connector at index `key`. + """ + # In case key is a class + if not isinstance(key, str): + if isinstance(key, slice): + raise NotImplementedError( + "Slicing of ConnectorPipeline is currently not supported." + ) + elif isinstance(key, int): + return [self.connectors[key]] + elif isinstance(key, type): + results = [] + for c in self.connectors: + if issubclass(c.__class__, key): + results.append(c) + return results + else: + raise NotImplementedError( + "Indexing by {} is currently not supported.".format(type(key)) + ) + + results = [] + for c in self.connectors: + if c.__class__.__name__ == key: + results.append(c) + + return results diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/connector_pipeline_v2.py b/lib/python3.12/site-packages/ray/rllib/connectors/connector_pipeline_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..66ee6a5aed9b076925103ec0315b7b6c143d5808 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/connector_pipeline_v2.py @@ -0,0 +1,415 @@ +import logging +from typing import Any, Collection, Dict, List, Optional, Tuple, Type, Union + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.metrics import CONNECTOR_PIPELINE_TIMER, CONNECTOR_TIMERS, TIMERS +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.metrics.utils import to_snake_case +from ray.rllib.utils.typing import EpisodeType, StateDict +from ray.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="alpha") +class ConnectorPipelineV2(ConnectorV2): + """Utility class for quick manipulation of a connector pipeline.""" + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + self._fix_spaces(input_observation_space, input_action_space) + return self.observation_space + + @override(ConnectorV2) + def recompute_output_action_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + self._fix_spaces(input_observation_space, input_action_space) + return self.action_space + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + connectors: Optional[List[ConnectorV2]] = None, + **kwargs, + ): + """Initializes a ConnectorPipelineV2 instance. + + Args: + input_observation_space: The (optional) input observation space for this + connector piece. This is the space coming from a previous connector + piece in the (env-to-module or learner) pipeline or is directly + defined within the gym.Env. + input_action_space: The (optional) input action space for this connector + piece. This is the space coming from a previous connector piece in the + (module-to-env) pipeline or is directly defined within the gym.Env. + connectors: A list of individual ConnectorV2 pieces to be added to this + pipeline during construction. Note that you can always add (or remove) + more ConnectorV2 pieces later on the fly. + """ + self.connectors = [] + + for conn in connectors: + # If we have a `ConnectorV2` instance just append. + if isinstance(conn, ConnectorV2): + self.connectors.append(conn) + # If, we have a class with `args` and `kwargs`, build the instance. + # Note that this way of constructing a pipeline should only be + # used internally when restoring the pipeline state from a + # checkpoint. + elif isinstance(conn, tuple) and len(conn) == 3: + self.connectors.append(conn[0](*conn[1], **conn[2])) + + super().__init__(input_observation_space, input_action_space, **kwargs) + + def __len__(self): + return len(self.connectors) + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + metrics: Optional[MetricsLogger] = None, + **kwargs, + ) -> Any: + """In a pipeline, we simply call each of our connector pieces after each other. + + Each connector piece receives as input the output of the previous connector + piece in the pipeline. + """ + shared_data = shared_data if shared_data is not None else {} + full_stats = None + if metrics: + full_stats = metrics.log_time( + kwargs.get("metrics_prefix_key", ()) + (CONNECTOR_PIPELINE_TIMER,) + ) + full_stats.__enter__() + + # Loop through connector pieces and call each one with the output of the + # previous one. Thereby, time each connector piece's call. + for connector in self.connectors: + # TODO (sven): Add MetricsLogger to non-Learner components that have a + # LearnerConnector pipeline. + stats = None + if metrics: + stats = metrics.log_time( + kwargs.get("metrics_prefix_key", ()) + + ( + TIMERS, + CONNECTOR_TIMERS, + to_snake_case(connector.__class__.__name__), + ) + ) + stats.__enter__() + + batch = connector( + rl_module=rl_module, + batch=batch, + episodes=episodes, + explore=explore, + shared_data=shared_data, + metrics=metrics, + # Deprecated arg. + data=batch, + **kwargs, + ) + + if metrics: + stats.__exit__(None, None, None) + + if not isinstance(batch, dict): + raise ValueError( + f"`data` returned by ConnectorV2 {connector} must be a dict! " + f"You returned {batch}. Check your (custom) connectors' " + f"`__call__()` method's return value and make sure you return " + f"the `batch` arg passed in (either altered or unchanged)." + ) + + if metrics: + full_stats.__exit__(None, None, None) + + return batch + + def remove(self, name_or_class: Union[str, Type]): + """Remove a single connector piece in this pipeline by its name or class. + + Args: + name_or_class: The name of the connector piece to be removed from the + pipeline. + """ + idx = -1 + for i, c in enumerate(self.connectors): + if (isinstance(name_or_class, type) and c.__class__ is name_or_class) or ( + isinstance(name_or_class, str) and c.__class__.__name__ == name_or_class + ): + idx = i + break + if idx >= 0: + del self.connectors[idx] + self._fix_spaces(self.input_observation_space, self.input_action_space) + logger.info( + f"Removed connector {name_or_class} from {self.__class__.__name__}." + ) + else: + logger.warning( + f"Trying to remove a non-existent connector {name_or_class}." + ) + + def insert_before( + self, + name_or_class: Union[str, type], + connector: ConnectorV2, + ) -> ConnectorV2: + """Insert a new connector piece before an existing piece (by name or class). + + Args: + name_or_class: Name or class of the connector piece before which `connector` + will get inserted. + connector: The new connector piece to be inserted. + + Returns: + The ConnectorV2 before which `connector` has been inserted. + """ + idx = -1 + for idx, c in enumerate(self.connectors): + if ( + isinstance(name_or_class, str) and c.__class__.__name__ == name_or_class + ) or (isinstance(name_or_class, type) and c.__class__ is name_or_class): + break + if idx < 0: + raise ValueError( + f"Can not find connector with name or type '{name_or_class}'!" + ) + next_connector = self.connectors[idx] + + self.connectors.insert(idx, connector) + self._fix_spaces(self.input_observation_space, self.input_action_space) + + logger.info( + f"Inserted {connector.__class__.__name__} before {name_or_class} " + f"to {self.__class__.__name__}." + ) + return next_connector + + def insert_after( + self, + name_or_class: Union[str, Type], + connector: ConnectorV2, + ) -> ConnectorV2: + """Insert a new connector piece after an existing piece (by name or class). + + Args: + name_or_class: Name or class of the connector piece after which `connector` + will get inserted. + connector: The new connector piece to be inserted. + + Returns: + The ConnectorV2 after which `connector` has been inserted. + """ + idx = -1 + for idx, c in enumerate(self.connectors): + if ( + isinstance(name_or_class, str) and c.__class__.__name__ == name_or_class + ) or (isinstance(name_or_class, type) and c.__class__ is name_or_class): + break + if idx < 0: + raise ValueError( + f"Can not find connector with name or type '{name_or_class}'!" + ) + prev_connector = self.connectors[idx] + + self.connectors.insert(idx + 1, connector) + self._fix_spaces(self.input_observation_space, self.input_action_space) + + logger.info( + f"Inserted {connector.__class__.__name__} after {name_or_class} " + f"to {self.__class__.__name__}." + ) + + return prev_connector + + def prepend(self, connector: ConnectorV2) -> None: + """Prepend a new connector at the beginning of a connector pipeline. + + Args: + connector: The new connector piece to be prepended to this pipeline. + """ + self.connectors.insert(0, connector) + self._fix_spaces(self.input_observation_space, self.input_action_space) + + logger.info( + f"Added {connector.__class__.__name__} to the beginning of " + f"{self.__class__.__name__}." + ) + + def append(self, connector: ConnectorV2) -> None: + """Append a new connector at the end of a connector pipeline. + + Args: + connector: The new connector piece to be appended to this pipeline. + """ + self.connectors.append(connector) + self._fix_spaces(self.input_observation_space, self.input_action_space) + + logger.info( + f"Added {connector.__class__.__name__} to the end of " + f"{self.__class__.__name__}." + ) + + @override(ConnectorV2) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + state = {} + for conn in self.connectors: + conn_name = type(conn).__name__ + if self._check_component(conn_name, components, not_components): + sts = conn.get_state( + components=self._get_subcomponents(conn_name, components), + not_components=self._get_subcomponents(conn_name, not_components), + **kwargs, + ) + # Ignore empty dicts. + if sts: + state[conn_name] = sts + return state + + @override(ConnectorV2) + def set_state(self, state: Dict[str, Any]) -> None: + for conn in self.connectors: + conn_name = type(conn).__name__ + if conn_name in state: + conn.set_state(state[conn_name]) + + @override(Checkpointable) + def get_checkpointable_components(self) -> List[Tuple[str, "Checkpointable"]]: + return [(type(conn).__name__, conn) for conn in self.connectors] + + # Note that we don't have to override Checkpointable.get_ctor_args_and_kwargs and + # don't have to return the `connectors` c'tor kwarg from there. This is b/c all + # connector pieces in this pipeline are themselves Checkpointable components, + # so they will be properly written into this pipeline's checkpoint. + @override(Checkpointable) + def get_ctor_args_and_kwargs(self) -> Tuple[Tuple, Dict[str, Any]]: + return ( + (self.input_observation_space, self.input_action_space), # *args + { + "connectors": [ + (type(conn), *conn.get_ctor_args_and_kwargs()) + for conn in self.connectors + ] + }, + ) + + @override(ConnectorV2) + def reset_state(self) -> None: + for conn in self.connectors: + conn.reset_state() + + @override(ConnectorV2) + def merge_states(self, states: List[Dict[str, Any]]) -> Dict[str, Any]: + merged_states = {} + if not states: + return merged_states + for i, (key, item) in enumerate(states[0].items()): + state_list = [state[key] for state in states] + conn = self.connectors[i] + merged_states[key] = conn.merge_states(state_list) + return merged_states + + def __repr__(self, indentation: int = 0): + return "\n".join( + [" " * indentation + self.__class__.__name__] + + [c.__str__(indentation + 4) for c in self.connectors] + ) + + def __getitem__( + self, + key: Union[str, int, Type], + ) -> Union[ConnectorV2, List[ConnectorV2]]: + """Returns a single ConnectorV2 or list of ConnectorV2s that fit `key`. + + If key is an int, we return a single ConnectorV2 at that index in this pipeline. + If key is a ConnectorV2 type or a string matching the class name of a + ConnectorV2 in this pipeline, we return a list of all ConnectorV2s in this + pipeline matching the specified class. + + Args: + key: The key to find or to index by. + + Returns: + A single ConnectorV2 or a list of ConnectorV2s matching `key`. + """ + # Key is an int -> Index into pipeline and return. + if isinstance(key, int): + return self.connectors[key] + # Key is a class. + elif isinstance(key, type): + results = [] + for c in self.connectors: + if issubclass(c.__class__, key): + results.append(c) + return results + # Key is a string -> Find connector(s) by name. + elif isinstance(key, str): + results = [] + for c in self.connectors: + if c.name == key: + results.append(c) + return results + # Slicing not supported (yet). + elif isinstance(key, slice): + raise NotImplementedError( + "Slicing of ConnectorPipelineV2 is currently not supported!" + ) + else: + raise NotImplementedError( + f"Indexing ConnectorPipelineV2 by {type(key)} is currently not " + f"supported!" + ) + + @property + def observation_space(self): + if len(self) > 0: + return self.connectors[-1].observation_space + return self._observation_space + + @property + def action_space(self): + if len(self) > 0: + return self.connectors[-1].action_space + return self._action_space + + def _fix_spaces(self, input_observation_space, input_action_space): + if len(self) > 0: + # Fix each connector's input_observation- and input_action space in + # the pipeline. + obs_space = input_observation_space + act_space = input_action_space + for con in self.connectors: + con.input_action_space = act_space + con.input_observation_space = obs_space + obs_space = con.observation_space + act_space = con.action_space diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/connector_v2.py b/lib/python3.12/site-packages/ray/rllib/connectors/connector_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..9b1129eca8b88ff2963dd066afbcfa9b31cb90c4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/connector_v2.py @@ -0,0 +1,1017 @@ +import abc +import inspect +from collections import defaultdict +from typing import ( + Any, + Callable, + Collection, + Dict, + Iterator, + List, + Optional, + Tuple, + Union, +) + +import gymnasium as gym +import tree + +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import OverrideToImplementCustomLogic, override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.spaces.space_utils import BatchedNdArray +from ray.rllib.utils.typing import AgentID, EpisodeType, ModuleID, StateDict +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class ConnectorV2(Checkpointable, abc.ABC): + """Base class defining the API for an individual "connector piece". + + A ConnectorV2 ("connector piece") is usually part of a whole series of connector + pieces within a so-called connector pipeline, which in itself also abides to this + very API. + For example, you might have a connector pipeline consisting of two connector pieces, + A and B, both instances of subclasses of ConnectorV2 and each one performing a + particular transformation on their input data. The resulting connector pipeline + (A->B) itself also abides to this very ConnectorV2 API and could thus be part of yet + another, higher-level connector pipeline, e.g. (A->B)->C->D. + + Any ConnectorV2 instance (individual pieces or several connector pieces in a + pipeline) is a callable and users should override the `__call__()` method. + When called, they take the outputs of a previous connector piece (or an empty dict + if there are no previous pieces) and all the data collected thus far in the + ongoing episode(s) (only applies to connectors used in EnvRunners) or retrieved + from a replay buffer or from an environment sampling step (only applies to + connectors used in Learner pipelines). From this input data, a ConnectorV2 then + performs a transformation step. + + There are 3 types of pipelines any ConnectorV2 piece can belong to: + 1) EnvToModulePipeline: The connector transforms environment data before it gets to + the RLModule. This type of pipeline is used by an EnvRunner for transforming + env output data into RLModule readable data (for the next RLModule forward pass). + For example, such a pipeline would include observation postprocessors, -filters, + or any RNN preparation code related to time-sequences and zero-padding. + 2) ModuleToEnvPipeline: This type of pipeline is used by an + EnvRunner to transform RLModule output data to env readable actions (for the next + `env.step()` call). For example, in case the RLModule only outputs action + distribution parameters (but not actual actions), the ModuleToEnvPipeline would + take care of sampling the actions to be sent back to the end from the + resulting distribution (made deterministic if exploration is off). + 3) LearnerConnectorPipeline: This connector pipeline type transforms data coming + from an `EnvRunner.sample()` call or a replay buffer and will then be sent into the + RLModule's `forward_train()` method in order to compute loss function inputs. + This type of pipeline is used by a Learner worker to transform raw training data + (a batch or a list of episodes) to RLModule readable training data (for the next + RLModule `forward_train()` call). + + Some connectors might be stateful, for example for keeping track of observation + filtering stats (mean and stddev values). Any Algorithm, which uses connectors is + responsible for frequently synchronizing the states of all connectors and connector + pipelines between the EnvRunners (owning the env-to-module and module-to-env + pipelines) and the Learners (owning the Learner pipelines). + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + **kwargs, + ): + """Initializes a ConnectorV2 instance. + + Args: + input_observation_space: The (optional) input observation space for this + connector piece. This is the space coming from a previous connector + piece in the (env-to-module or learner) pipeline or is directly + defined within the gym.Env. + input_action_space: The (optional) input action space for this connector + piece. This is the space coming from a previous connector piece in the + (module-to-env) pipeline or is directly defined within the gym.Env. + **kwargs: Forward API-compatibility kwargs. + """ + self._observation_space = None + self._action_space = None + self._input_observation_space = None + self._input_action_space = None + + self.input_action_space = input_action_space + self.input_observation_space = input_observation_space + + # Store child's constructor args and kwargs for the default + # `get_ctor_args_and_kwargs` implementation (to be able to restore from a + # checkpoint). + if self.__class__.__dict__.get("__init__") is not None: + caller_frame = inspect.stack()[1].frame + arg_info = inspect.getargvalues(caller_frame) + # Separate positional arguments and keyword arguments. + caller_locals = ( + arg_info.locals + ) # Dictionary of all local variables in the caller + self._ctor_kwargs = { + arg: caller_locals[arg] for arg in arg_info.args if arg != "self" + } + else: + self._ctor_kwargs = { + "input_observation_space": self.input_observation_space, + "input_action_space": self.input_action_space, + } + + @OverrideToImplementCustomLogic + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + """Re-computes a new (output) observation space based on the input spaces. + + This method should be overridden by users to make sure a ConnectorPipelineV2 + knows how the input spaces through its individual ConnectorV2 pieces are being + transformed. + + .. testcode:: + + from gymnasium.spaces import Box, Discrete + import numpy as np + + from ray.rllib.connectors.connector_v2 import ConnectorV2 + from ray.rllib.utils.numpy import one_hot + from ray.rllib.utils.test_utils import check + + class OneHotConnector(ConnectorV2): + def recompute_output_observation_space( + self, + input_observation_space, + input_action_space, + ): + return Box(0.0, 1.0, (input_observation_space.n,), np.float32) + + def __call__( + self, + *, + rl_module, + batch, + episodes, + explore=None, + shared_data=None, + metrics=None, + **kwargs, + ): + assert "obs" in batch + batch["obs"] = one_hot(batch["obs"]) + return batch + + connector = OneHotConnector(input_observation_space=Discrete(2)) + batch = {"obs": np.array([1, 0, 0], np.int32)} + output = connector(rl_module=None, batch=batch, episodes=None) + + check(output, {"obs": np.array([[0.0, 1.0], [1.0, 0.0], [1.0, 0.0]])}) + + If this ConnectorV2 does not change the observation space in any way, leave + this parent method implementation untouched. + + Args: + input_observation_space: The input observation space (either coming from the + environment if `self` is the first connector piece in the pipeline or + from the previous connector piece in the pipeline). + input_action_space: The input action space (either coming from the + environment if `self` is the first connector piece in the pipeline or + from the previous connector piece in the pipeline). + + Returns: + The new observation space (after data has passed through this ConnectorV2 + piece). + """ + return self.input_observation_space + + @OverrideToImplementCustomLogic + def recompute_output_action_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + """Re-computes a new (output) action space based on the input space. + + This method should be overridden by users to make sure a ConnectorPipelineV2 + knows how the input spaces through its individual ConnectorV2 pieces are being + transformed. + + If this ConnectorV2 does not change the action space in any way, leave + this parent method implementation untouched. + + Args: + input_observation_space: The input observation space (either coming from the + environment if `self` is the first connector piece in the pipeline or + from the previous connector piece in the pipeline). + input_action_space: The input action space (either coming from the + environment if `self` is the first connector piece in the pipeline or + from the previous connector piece in the pipeline). + + Returns: + The new action space (after data has passed through this ConenctorV2 + piece). + """ + return self.input_action_space + + @abc.abstractmethod + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + metrics: Optional[MetricsLogger] = None, + **kwargs, + ) -> Any: + """Method for transforming an input `batch` into an output `batch`. + + Args: + rl_module: The RLModule object that the connector connects to or from. + batch: The input data to be transformed by this connector. Transformations + might either be done in-place or a new structure may be returned. + Note that the information in `batch` will eventually either become the + forward batch for the RLModule (env-to-module and learner connectors) + or the input to the `env.step()` call (module-to-env connectors). Note + that in the first case (`batch` is a forward batch for RLModule), the + information in `batch` will be discarded after that RLModule forward + pass. Any transformation of information (e.g. observation preprocessing) + that you have only done inside `batch` will be lost, unless you have + written it back into the corresponding `episodes` during the connector + pass. + episodes: The list of SingleAgentEpisode or MultiAgentEpisode objects, + each corresponding to one slot in the vector env. Note that episodes + can be read from (e.g. to place information into `batch`), but also + written to. You should only write back (changed, transformed) + information into the episodes, if you want these changes to be + "permanent". For example if you sample from an environment, pick up + observations from the episodes and place them into `batch`, then + transform these observations, and would like to make these + transformations permanent (note that `batch` gets discarded after the + RLModule forward pass), then you have to write the transformed + observations back into the episode to make sure you do not have to + perform the same transformation again on the learner (or replay buffer) + side. The Learner will hence work on the already changed episodes (and + compile the train batch using the Learner connector). + explore: Whether `explore` is currently on. Per convention, if True, the + RLModule's `forward_exploration` method should be called, if False, the + EnvRunner should call `forward_inference` instead. + shared_data: Optional additional context data that needs to be exchanged + between different ConnectorV2 pieces (in the same pipeline) or across + ConnectorV2 pipelines (meaning between env-to-module and module-to-env). + metrics: Optional MetricsLogger instance to log custom metrics to. + kwargs: Forward API-compatibility kwargs. + + Returns: + The transformed connector output. + """ + + @staticmethod + def single_agent_episode_iterator( + episodes: List[EpisodeType], + agents_that_stepped_only: bool = True, + zip_with_batch_column: Optional[Union[List[Any], Dict[Tuple, Any]]] = None, + ) -> Iterator[SingleAgentEpisode]: + """An iterator over a list of episodes yielding always SingleAgentEpisodes. + + In case items in the list are MultiAgentEpisodes, these are broken down + into their individual agents' SingleAgentEpisodes and those are then yielded + one after the other. + + Useful for connectors that operate on both single-agent and multi-agent + episodes. + + Args: + episodes: The list of SingleAgent- or MultiAgentEpisode objects. + agents_that_stepped_only: If True (and multi-agent setup), will only place + items of those agents into the batch that have just stepped in the + actual MultiAgentEpisode (this is checked via a + `MultiAgentEpside.episode.get_agents_to_act()`). Note that this setting + is ignored in a single-agent setups b/c the agent steps at each timestep + regardless. + zip_with_batch_column: If provided, must be a list of batch items + corresponding to the given `episodes` (single agent case) or a dict + mapping (AgentID, ModuleID) tuples to lists of individual batch items + corresponding to this agent/module combination. The iterator will then + yield tuples of SingleAgentEpisode objects (1st item) along with the + data item (2nd item) that this episode was responsible for generating + originally. + + Yields: + All SingleAgentEpisodes in the input list, whereby MultiAgentEpisodes will + be broken down into their individual SingleAgentEpisode components. + """ + list_indices = defaultdict(int) + + # Single-agent case. + if episodes and isinstance(episodes[0], SingleAgentEpisode): + if zip_with_batch_column is not None: + if len(zip_with_batch_column) != len(episodes): + raise ValueError( + "Invalid `zip_with_batch_column` data: Must have the same " + f"length as the list of episodes ({len(episodes)}), but has " + f"length {len(zip_with_batch_column)}!" + ) + # Simple case: Items are stored in lists directly under the column (str) + # key. + if isinstance(zip_with_batch_column, list): + for episode, data in zip(episodes, zip_with_batch_column): + yield episode, data + # Normal single-agent case: Items are stored in dicts under the column + # (str) key. These dicts map (eps_id,)-tuples to lists of individual + # items. + else: + for episode, (eps_id_tuple, data) in zip( + episodes, + zip_with_batch_column.items(), + ): + assert episode.id_ == eps_id_tuple[0] + d = data[list_indices[eps_id_tuple]] + list_indices[eps_id_tuple] += 1 + yield episode, d + else: + for episode in episodes: + yield episode + return + + # Multi-agent case. + for episode in episodes: + for agent_id in ( + episode.get_agents_that_stepped() + if agents_that_stepped_only + else episode.agent_ids + ): + sa_episode = episode.agent_episodes[agent_id] + # for sa_episode in episode.agent_episodes.values(): + if zip_with_batch_column is not None: + key = ( + sa_episode.multi_agent_episode_id, + sa_episode.agent_id, + sa_episode.module_id, + ) + if len(zip_with_batch_column[key]) <= list_indices[key]: + raise ValueError( + "Invalid `zip_with_batch_column` data: Must structurally " + "match the single-agent contents in the given list of " + "(multi-agent) episodes!" + ) + d = zip_with_batch_column[key][list_indices[key]] + list_indices[key] += 1 + yield sa_episode, d + else: + yield sa_episode + + @staticmethod + def add_batch_item( + batch: Dict[str, Any], + column: str, + item_to_add: Any, + single_agent_episode: Optional[SingleAgentEpisode] = None, + ) -> None: + """Adds a data item under `column` to the given `batch`. + + The `item_to_add` is stored in the `batch` in the following manner: + 1) If `single_agent_episode` is not provided (None), will store the item in a + list directly under `column`: + `column` -> [item, item, ...] + 2) If `single_agent_episode`'s `agent_id` and `module_id` properties are None + (`single_agent_episode` is not part of a multi-agent episode), will append + `item_to_add` to a list under a `(,)` key under `column`: + `column` -> `(,)` -> [item, item, ...] + 3) If `single_agent_episode`'s `agent_id` and `module_id` are NOT None + (`single_agent_episode` is part of a multi-agent episode), will append + `item_to_add` to a list under a `(,,)` key + under `column`: + `column` -> `(,,)` -> [item, item, ...] + + See the these examples here for clarification of these three cases: + + .. testcode:: + + from ray.rllib.connectors.connector_v2 import ConnectorV2 + from ray.rllib.env.multi_agent_episode import MultiAgentEpisode + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # 1) Simple case (no episodes provided) -> Store data in a list directly + # under `column`: + batch = {} + ConnectorV2.add_batch_item(batch, "test_col", item_to_add=5) + ConnectorV2.add_batch_item(batch, "test_col", item_to_add=6) + check(batch, {"test_col": [5, 6]}) + ConnectorV2.add_batch_item(batch, "test_col_2", item_to_add=-10) + check(batch, { + "test_col": [5, 6], + "test_col_2": [-10], + }) + + # 2) Single-agent case (SingleAgentEpisode provided) -> Store data in a list + # under the keys: `column` -> `(,)` -> [...]: + batch = {} + episode = SingleAgentEpisode( + id_="SA-EPS0", + observations=[0, 1, 2, 3], + actions=[1, 2, 3], + rewards=[1.0, 2.0, 3.0], + ) + ConnectorV2.add_batch_item(batch, "test_col", 5, episode) + ConnectorV2.add_batch_item(batch, "test_col", 6, episode) + ConnectorV2.add_batch_item(batch, "test_col_2", -10, episode) + check(batch, { + "test_col": {("SA-EPS0",): [5, 6]}, + "test_col_2": {("SA-EPS0",): [-10]}, + }) + + # 3) Multi-agent case (SingleAgentEpisode provided that has `agent_id` and + # `module_id` information) -> Store data in a list under the keys: + # `column` -> `(,,)` -> [...]: + batch = {} + ma_episode = MultiAgentEpisode( + id_="MA-EPS1", + observations=[ + {"ag0": 0, "ag1": 1}, {"ag0": 2, "ag1": 4} + ], + actions=[{"ag0": 0, "ag1": 1}], + rewards=[{"ag0": -0.1, "ag1": -0.2}], + # ag0 maps to mod0, ag1 maps to mod1, etc.. + agent_to_module_mapping_fn=lambda aid, eps: f"mod{aid[2:]}", + ) + ConnectorV2.add_batch_item( + batch, + "test_col", + item_to_add=5, + single_agent_episode=ma_episode.agent_episodes["ag0"], + ) + ConnectorV2.add_batch_item( + batch, + "test_col", + item_to_add=6, + single_agent_episode=ma_episode.agent_episodes["ag0"], + ) + ConnectorV2.add_batch_item( + batch, + "test_col_2", + item_to_add=10, + single_agent_episode=ma_episode.agent_episodes["ag1"], + ) + check( + batch, + { + "test_col": {("MA-EPS1", "ag0", "mod0"): [5, 6]}, + "test_col_2": {("MA-EPS1", "ag1", "mod1"): [10]}, + }, + ) + + Args: + batch: The batch to store `item_to_add` in. + column: The column name (str) within the `batch` to store `item_to_add` + under. + item_to_add: The data item to store in the batch. + single_agent_episode: An optional SingleAgentEpisode. + If provided and its `agent_id` and `module_id` properties are None, + creates a further sub dictionary under `column`, mapping from + `(,)` to a list of data items (to which `item_to_add` will + be appended in this call). + If provided and its `agent_id` and `module_id` properties are NOT None, + creates a further sub dictionary under `column`, mapping from + `(,,,)` to a list of data items (to which + `item_to_add` will be appended in this call). + If not provided, will append `item_to_add` to a list directly under + `column`. + """ + sub_key = None + # SAEpisode is provided ... + if single_agent_episode is not None: + module_id = single_agent_episode.module_id + # ... and has `module_id` AND that `module_id` is already a top-level key in + # `batch` (`batch` is already in module-major form, mapping ModuleID to + # columns mapping to data). + if module_id is not None and module_id in batch: + raise ValueError( + "Can't call `add_batch_item` on a `batch` that is already " + "module-major (meaning ModuleID is top-level with column names on " + "the level thereunder)! Make sure to only call `add_batch_items` " + "before the `AgentToModuleMapping` ConnectorV2 piece is applied." + ) + + # ... and has `agent_id` -> Use `single_agent_episode`'s agent ID and + # module ID. + elif single_agent_episode.agent_id is not None: + sub_key = ( + single_agent_episode.multi_agent_episode_id, + single_agent_episode.agent_id, + single_agent_episode.module_id, + ) + # Otherwise, just use episode's ID. + else: + sub_key = (single_agent_episode.id_,) + + if column not in batch: + batch[column] = [] if sub_key is None else {sub_key: []} + if sub_key is not None: + if sub_key not in batch[column]: + batch[column][sub_key] = [] + batch[column][sub_key].append(item_to_add) + else: + batch[column].append(item_to_add) + + @staticmethod + def add_n_batch_items( + batch: Dict[str, Any], + column: str, + items_to_add: Any, + num_items: int, + single_agent_episode: Optional[SingleAgentEpisode] = None, + ) -> None: + """Adds a list of items (or batched item) under `column` to the given `batch`. + + If `items_to_add` is not a list, but an already batched struct (of np.ndarray + leafs), the `items_to_add` will be appended to possibly existing data under the + same `column` as-is. A subsequent `BatchIndividualItems` ConnectorV2 piece will + recognize this and batch the data properly into a single (batched) item. + This is much faster than first splitting up `items_to_add` and then adding each + item individually. + + If `single_agent_episode` is provided and its `agent_id` and `module_id` + properties are None, creates a further sub dictionary under `column`, mapping + from `(,)` to a list of data items (to which `items_to_add` will + be appended in this call). + If `single_agent_episode` is provided and its `agent_id` and `module_id` + properties are NOT None, creates a further sub dictionary under `column`, + mapping from `(,,,)` to a list of data items (to + which `items_to_add` will be appended in this call). + If `single_agent_episode` is not provided, will append `items_to_add` to a list + directly under `column`. + + .. testcode:: + + import numpy as np + + from ray.rllib.connectors.connector_v2 import ConnectorV2 + from ray.rllib.env.multi_agent_episode import MultiAgentEpisode + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Simple case (no episodes provided) -> Store data in a list directly under + # `column`: + batch = {} + ConnectorV2.add_n_batch_items( + batch, + "test_col", + # List of (complex) structs. + [{"a": np.array(3), "b": 4}, {"a": np.array(5), "b": 6}], + num_items=2, + ) + check( + batch["test_col"], + [{"a": np.array(3), "b": 4}, {"a": np.array(5), "b": 6}], + ) + # In a new column (test_col_2), store some already batched items. + # This way, you may avoid having to disassemble an already batched item + # (e.g. a numpy array of shape (10, 2)) into its individual items (e.g. + # split the array into a list of len=10) and then adding these individually. + # The performance gains may be quite large when providing already batched + # items (such as numpy arrays with a batch dim): + ConnectorV2.add_n_batch_items( + batch, + "test_col_2", + # One (complex) already batched struct. + {"a": np.array([3, 5]), "b": np.array([4, 6])}, + num_items=2, + ) + # Add more already batched items (this time with a different batch size) + ConnectorV2.add_n_batch_items( + batch, + "test_col_2", + {"a": np.array([7, 7, 7]), "b": np.array([8, 8, 8])}, + num_items=3, # <- in this case, this must be the batch size + ) + check( + batch["test_col_2"], + [ + {"a": np.array([3, 5]), "b": np.array([4, 6])}, + {"a": np.array([7, 7, 7]), "b": np.array([8, 8, 8])}, + ], + ) + + # Single-agent case (SingleAgentEpisode provided) -> Store data in a list + # under the keys: `column` -> `(,)`: + batch = {} + episode = SingleAgentEpisode( + id_="SA-EPS0", + observations=[0, 1, 2, 3], + actions=[1, 2, 3], + rewards=[1.0, 2.0, 3.0], + ) + ConnectorV2.add_n_batch_items( + batch=batch, + column="test_col", + items_to_add=[5, 6, 7], + num_items=3, + single_agent_episode=episode, + ) + check(batch, { + "test_col": {("SA-EPS0",): [5, 6, 7]}, + }) + + # Multi-agent case (SingleAgentEpisode provided that has `agent_id` and + # `module_id` information) -> Store data in a list under the keys: + # `column` -> `(,,)`: + batch = {} + ma_episode = MultiAgentEpisode( + id_="MA-EPS1", + observations=[ + {"ag0": 0, "ag1": 1}, {"ag0": 2, "ag1": 4} + ], + actions=[{"ag0": 0, "ag1": 1}], + rewards=[{"ag0": -0.1, "ag1": -0.2}], + # ag0 maps to mod0, ag1 maps to mod1, etc.. + agent_to_module_mapping_fn=lambda aid, eps: f"mod{aid[2:]}", + ) + ConnectorV2.add_batch_item( + batch, + "test_col", + item_to_add=5, + single_agent_episode=ma_episode.agent_episodes["ag0"], + ) + ConnectorV2.add_batch_item( + batch, + "test_col", + item_to_add=6, + single_agent_episode=ma_episode.agent_episodes["ag0"], + ) + ConnectorV2.add_batch_item( + batch, + "test_col_2", + item_to_add=10, + single_agent_episode=ma_episode.agent_episodes["ag1"], + ) + check( + batch, + { + "test_col": {("MA-EPS1", "ag0", "mod0"): [5, 6]}, + "test_col_2": {("MA-EPS1", "ag1", "mod1"): [10]}, + }, + ) + + Args: + batch: The batch to store n `items_to_add` in. + column: The column name (str) within the `batch` to store `item_to_add` + under. + items_to_add: The list of data items to store in the batch OR an already + batched (possibly nested) struct. In the latter case, the `items_to_add` + will be appended to possibly existing data under the same `column` + as-is. A subsequent `BatchIndividualItems` ConnectorV2 piece will + recognize this and batch the data properly into a single (batched) item. + This is much faster than first splitting up `items_to_add` and then + adding each item individually. + num_items: The number of items in `items_to_add`. This arg is mostly for + asserting the correct usage of this method by checking, whether the + given data in `items_to_add` really has the right amount of individual + items. + single_agent_episode: An optional SingleAgentEpisode. + If provided and its `agent_id` and `module_id` properties are None, + creates a further sub dictionary under `column`, mapping from + `(,)` to a list of data items (to which `items_to_add` will + be appended in this call). + If provided and its `agent_id` and `module_id` properties are NOT None, + creates a further sub dictionary under `column`, mapping from + `(,,,)` to a list of data items (to which + `items_to_add` will be appended in this call). + If not provided, will append `items_to_add` to a list directly under + `column`. + """ + # Process n list items by calling `add_batch_item` on each of them individually. + if isinstance(items_to_add, list): + if len(items_to_add) != num_items: + raise ValueError( + f"Mismatch between `num_items` ({num_items}) and the length " + f"of the provided list ({len(items_to_add)}) in " + f"{ConnectorV2.__name__}.add_n_batch_items()!" + ) + for item in items_to_add: + ConnectorV2.add_batch_item( + batch=batch, + column=column, + item_to_add=item, + single_agent_episode=single_agent_episode, + ) + return + + # Process a batched (possibly complex) struct. + # We could just unbatch the item (split it into a list) and then add each + # individual item to our `batch`. However, this comes with a heavy performance + # penalty. Instead, we tag the thus added array(s) here as "_has_batch_dim=True" + # and then know that when batching the entire list under the respective + # (eps_id, agent_id, module_id)-tuple key, we need to concatenate, not stack + # the items in there. + def _tag(s): + return BatchedNdArray(s) + + ConnectorV2.add_batch_item( + batch=batch, + column=column, + # Convert given input into BatchedNdArray(s) such that the `batch` utility + # knows that it'll have to concat, not stack. + item_to_add=tree.map_structure(_tag, items_to_add), + single_agent_episode=single_agent_episode, + ) + + @staticmethod + def foreach_batch_item_change_in_place( + batch: Dict[str, Any], + column: Union[str, List[str], Tuple[str]], + func: Callable[ + [Any, Optional[int], Optional[AgentID], Optional[ModuleID]], Any + ], + ) -> None: + """Runs the provided `func` on all items under one or more columns in the batch. + + Use this method to conveniently loop through all items in a batch + and transform them in place. + + `func` takes the following as arguments: + - The item itself. If column is a list of column names, this argument is a tuple + of items. + - The EpisodeID. This value might be None. + - The AgentID. This value might be None in the single-agent case. + - The ModuleID. This value might be None in the single-agent case. + + The return value(s) of `func` are used to directly override the values in the + given `batch`. + + Args: + batch: The batch to process in-place. + column: A single column name (str) or a list thereof. If a list is provided, + the first argument to `func` is a tuple of items. If a single + str is provided, the first argument to `func` is an individual + item. + func: The function to call on each item or tuple of item(s). + + .. testcode:: + + from ray.rllib.connectors.connector_v2 import ConnectorV2 + from ray.rllib.utils.test_utils import check + + # Simple case: Batch items are in lists directly under their column names. + batch = { + "col1": [0, 1, 2, 3], + "col2": [0, -1, -2, -3], + } + # Increase all ints by 1. + ConnectorV2.foreach_batch_item_change_in_place( + batch=batch, + column="col1", + func=lambda item, *args: item + 1, + ) + check(batch["col1"], [1, 2, 3, 4]) + + # Further increase all ints by 1 in col1 and flip sign in col2. + ConnectorV2.foreach_batch_item_change_in_place( + batch=batch, + column=["col1", "col2"], + func=(lambda items, *args: (items[0] + 1, -items[1])), + ) + check(batch["col1"], [2, 3, 4, 5]) + check(batch["col2"], [0, 1, 2, 3]) + + # Single-agent case: Batch items are in lists under (eps_id,)-keys in a dict + # under their column names. + batch = { + "col1": { + ("eps1",): [0, 1, 2, 3], + ("eps2",): [400, 500, 600], + }, + } + # Increase all ints of eps1 by 1 and divide all ints of eps2 by 100. + ConnectorV2.foreach_batch_item_change_in_place( + batch=batch, + column="col1", + func=lambda item, eps_id, *args: ( + item + 1 if eps_id == "eps1" else item / 100 + ), + ) + check(batch["col1"], { + ("eps1",): [1, 2, 3, 4], + ("eps2",): [4, 5, 6], + }) + + # Multi-agent case: Batch items are in lists under + # (eps_id, agent_id, module_id)-keys in a dict + # under their column names. + batch = { + "col1": { + ("eps1", "ag1", "mod1"): [1, 2, 3, 4], + ("eps2", "ag1", "mod2"): [400, 500, 600], + ("eps2", "ag2", "mod3"): [-1, -2, -3, -4, -5], + }, + } + # Decrease all ints of "eps1" by 1, divide all ints of "mod2" by 100, and + # flip sign of all ints of "ag2". + ConnectorV2.foreach_batch_item_change_in_place( + batch=batch, + column="col1", + func=lambda item, eps_id, ag_id, mod_id: ( + item - 1 + if eps_id == "eps1" + else item / 100 + if mod_id == "mod2" + else -item + ), + ) + check(batch["col1"], { + ("eps1", "ag1", "mod1"): [0, 1, 2, 3], + ("eps2", "ag1", "mod2"): [4, 5, 6], + ("eps2", "ag2", "mod3"): [1, 2, 3, 4, 5], + }) + """ + data_to_process = [batch.get(c) for c in force_list(column)] + single_col = isinstance(column, str) + if any(d is None for d in data_to_process): + raise ValueError( + f"Invalid column name(s) ({column})! One or more not found in " + f"given batch. Found columns {list(batch.keys())}." + ) + + # Simple case: Data items are stored in a list directly under the column + # name(s). + if isinstance(data_to_process[0], list): + for list_pos, data_tuple in enumerate(zip(*data_to_process)): + results = func( + data_tuple[0] if single_col else data_tuple, + None, # episode_id + None, # agent_id + None, # module_id + ) + # Tuple'ize results if single_col. + results = (results,) if single_col else results + for col_slot, result in enumerate(force_list(results)): + data_to_process[col_slot][list_pos] = result + # Single-agent/multi-agent cases. + else: + for key, d0_list in data_to_process[0].items(): + # Multi-agent case: There is a dict mapping from a + # (eps id, AgentID, ModuleID)-tuples to lists of individual data items. + if len(key) == 3: + eps_id, agent_id, module_id = key + # Single-agent case: There is a dict mapping from a (eps_id,)-tuple + # to lists of individual data items. + # AgentID and ModuleID are both None. + else: + eps_id = key[0] + agent_id = module_id = None + other_lists = [d[key] for d in data_to_process[1:]] + for list_pos, data_tuple in enumerate(zip(d0_list, *other_lists)): + results = func( + data_tuple[0] if single_col else data_tuple, + eps_id, + agent_id, + module_id, + ) + # Tuple'ize results if single_col. + results = (results,) if single_col else results + for col_slot, result in enumerate(results): + data_to_process[col_slot][key][list_pos] = result + + @staticmethod + def switch_batch_from_column_to_module_ids( + batch: Dict[str, Dict[ModuleID, Any]] + ) -> Dict[ModuleID, Dict[str, Any]]: + """Switches the first two levels of a `col_name -> ModuleID -> data` type batch. + + Assuming that the top level consists of column names as keys and the second + level (under these columns) consists of ModuleID keys, the resulting batch + will have these two reversed and thus map ModuleIDs to dicts mapping column + names to data items. + + .. testcode:: + + from ray.rllib.utils.test_utils import check + + batch = { + "obs": {"module_0": [1, 2, 3]}, + "actions": {"module_0": [4, 5, 6], "module_1": [7]}, + } + switched_batch = ConnectorV2.switch_batch_from_column_to_module_ids(batch) + check( + switched_batch, + { + "module_0": {"obs": [1, 2, 3], "actions": [4, 5, 6]}, + "module_1": {"actions": [7]}, + }, + ) + + Args: + batch: The batch to switch from being column name based (then ModuleIDs) + to being ModuleID based (then column names). + + Returns: + A new batch dict mapping ModuleIDs to dicts mapping column names (e.g. + "obs") to data. + """ + module_data = defaultdict(dict) + for column, column_data in batch.items(): + for module_id, data in column_data.items(): + module_data[module_id][column] = data + return dict(module_data) + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + return {} + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + pass + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self) -> Tuple[Tuple, Dict[str, Any]]: + return ( + (), # *args + self._ctor_kwargs, # **kwargs + ) + + def reset_state(self) -> None: + """Resets the state of this ConnectorV2 to some initial value. + + Note that this may NOT be the exact state that this ConnectorV2 was originally + constructed with. + """ + return + + def merge_states(self, states: List[Dict[str, Any]]) -> Dict[str, Any]: + """Computes a resulting state given self's state and a list of other states. + + Algorithms should use this method for merging states between connectors + running on parallel EnvRunner workers. For example, to synchronize the connector + states of n remote workers and a local worker, one could: + - Gather all remote worker connector states in a list. + - Call `self.merge_states()` on the local worker passing it the states list. + - Broadcast the resulting local worker's connector state back to all remote + workers. After this, all workers (including the local one) hold a + merged/synchronized new connecto state. + + Args: + states: The list of n other ConnectorV2 states to merge with self's state + into a single resulting state. + + Returns: + The resulting state dict. + """ + return {} + + @property + def observation_space(self): + """Getter for our (output) observation space. + + Logic: Use user provided space (if set via `observation_space` setter) + otherwise, use the same as the input space, assuming this connector piece + does not alter the space. + """ + return self._observation_space + + @property + def action_space(self): + """Getter for our (output) action space. + + Logic: Use user provided space (if set via `action_space` setter) + otherwise, use the same as the input space, assuming this connector piece + does not alter the space. + """ + return self._action_space + + @property + def input_observation_space(self): + return self._input_observation_space + + @input_observation_space.setter + def input_observation_space(self, value): + self._input_observation_space = value + if value is not None: + self._observation_space = self.recompute_output_observation_space( + value, self.input_action_space + ) + + @property + def input_action_space(self): + return self._input_action_space + + @input_action_space.setter + def input_action_space(self, value): + self._input_action_space = value + if value is not None: + self._action_space = self.recompute_output_action_space( + self.input_observation_space, value + ) + + def __str__(self, indentation: int = 0): + return " " * indentation + self.__class__.__name__ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__init__.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9ae184ee597b40a1c89fc8f7f6be149f0ab8d2bb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__init__.py @@ -0,0 +1,39 @@ +from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import ( + AddObservationsFromEpisodesToBatch, +) +from ray.rllib.connectors.common.add_states_from_episodes_to_batch import ( + AddStatesFromEpisodesToBatch, +) +from ray.rllib.connectors.common.add_time_dim_to_batch_and_zero_pad import ( + AddTimeDimToBatchAndZeroPad, +) +from ray.rllib.connectors.common.agent_to_module_mapping import AgentToModuleMapping +from ray.rllib.connectors.common.batch_individual_items import BatchIndividualItems +from ray.rllib.connectors.common.numpy_to_tensor import NumpyToTensor +from ray.rllib.connectors.env_to_module.env_to_module_pipeline import ( + EnvToModulePipeline, +) +from ray.rllib.connectors.env_to_module.flatten_observations import ( + FlattenObservations, +) +from ray.rllib.connectors.env_to_module.mean_std_filter import MeanStdFilter +from ray.rllib.connectors.env_to_module.prev_actions_prev_rewards import ( + PrevActionsPrevRewards, +) +from ray.rllib.connectors.env_to_module.write_observations_to_episodes import ( + WriteObservationsToEpisodes, +) + +__all__ = [ + "AddObservationsFromEpisodesToBatch", + "AddStatesFromEpisodesToBatch", + "AddTimeDimToBatchAndZeroPad", + "AgentToModuleMapping", + "BatchIndividualItems", + "EnvToModulePipeline", + "FlattenObservations", + "MeanStdFilter", + "NumpyToTensor", + "PrevActionsPrevRewards", + "WriteObservationsToEpisodes", +] diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd894912971f5d2ba7b46bad911196fda1da5b7a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/env_to_module_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/env_to_module_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc3ed1928fe4bed4c26870bc5809ad2032cb1cdc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/env_to_module_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/flatten_observations.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/flatten_observations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f056bd37da528d9ab760c5e06ad9cab77d003038 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/flatten_observations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/frame_stacking.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/frame_stacking.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7cbaa858bc302f13327130946dae88762b7b789d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/frame_stacking.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/mean_std_filter.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/mean_std_filter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65cdb907220d6bedcc597942f95084954d0acb0b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/mean_std_filter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/observation_preprocessor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/observation_preprocessor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc3ca6043c3e3b7d31dad5fca82c126a3c6ade54 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/observation_preprocessor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/prev_actions_prev_rewards.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/prev_actions_prev_rewards.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33d1aba3f39e8cca2cfd6b7243ba355d4d6dff71 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/prev_actions_prev_rewards.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/write_observations_to_episodes.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/write_observations_to_episodes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6fe3c638de9025ac6761b379f05031f9aa2089bc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/__pycache__/write_observations_to_episodes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/env_to_module_pipeline.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/env_to_module_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..04a5389737e1ffded64793b77b48207687ad6af0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/env_to_module_pipeline.py @@ -0,0 +1,55 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_pipeline_v2 import ConnectorPipelineV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics import ( + ENV_TO_MODULE_SUM_EPISODES_LENGTH_IN, + ENV_TO_MODULE_SUM_EPISODES_LENGTH_OUT, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class EnvToModulePipeline(ConnectorPipelineV2): + @override(ConnectorPipelineV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]] = None, + episodes: List[EpisodeType], + explore: bool, + shared_data: Optional[dict] = None, + metrics: Optional[MetricsLogger] = None, + **kwargs, + ): + # Log the sum of lengths of all episodes incoming. + if metrics: + metrics.log_value( + ENV_TO_MODULE_SUM_EPISODES_LENGTH_IN, + sum(map(len, episodes)), + ) + + # Make sure user does not necessarily send initial input into this pipeline. + # Might just be empty and to be populated from `episodes`. + ret = super().__call__( + rl_module=rl_module, + batch=batch if batch is not None else {}, + episodes=episodes, + explore=explore, + shared_data=shared_data if shared_data is not None else {}, + metrics=metrics, + **kwargs, + ) + + # Log the sum of lengths of all episodes outgoing. + if metrics: + metrics.log_value( + ENV_TO_MODULE_SUM_EPISODES_LENGTH_OUT, + sum(map(len, episodes)), + ) + + return ret diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/flatten_observations.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/flatten_observations.py new file mode 100644 index 0000000000000000000000000000000000000000..4435ea861bf315e5e8e2b7fd7359a4a0b3872e79 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/flatten_observations.py @@ -0,0 +1,5 @@ +from functools import partial + +from ray.rllib.connectors.common.flatten_observations import FlattenObservations + +FlattenObservations = partial(FlattenObservations, as_learner_connector=False) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/frame_stacking.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/frame_stacking.py new file mode 100644 index 0000000000000000000000000000000000000000..8269d859059fce9d40e3219a0618a720f43e41dd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/frame_stacking.py @@ -0,0 +1,5 @@ +from functools import partial + +from ray.rllib.connectors.common.frame_stacking import FrameStacking + +FrameStackingEnvToModule = partial(FrameStacking, as_learner_connector=False) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/mean_std_filter.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/mean_std_filter.py new file mode 100644 index 0000000000000000000000000000000000000000..c668a0492d08eb27997a65afd5bef0b07456981d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/mean_std_filter.py @@ -0,0 +1,253 @@ +from typing import Any, Collection, Dict, List, Optional, Union + +import gymnasium as gym +import numpy as np +import tree +from gymnasium.spaces import Discrete, MultiDiscrete + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.filter import MeanStdFilter as _MeanStdFilter, RunningStat +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.typing import AgentID, EpisodeType, StateDict +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class MeanStdFilter(ConnectorV2): + """A connector used to mean-std-filter observations. + + Incoming observations are filtered such that the output of this filter is on + average 0.0 and has a standard deviation of 1.0. If the observation space is + a (possibly nested) dict, this filtering is applied separately per element of + the observation space (except for discrete- and multi-discrete elements, which + are left as-is). + + This connector is stateful as it continues to update its internal stats on mean + and std values as new data is pushed through it (unless `update_stats` is False). + """ + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + _input_observation_space_struct = get_base_struct_from_space( + input_observation_space + ) + + # Adjust our observation space's Boxes (only if clipping is active). + _observation_space_struct = tree.map_structure( + lambda s: ( + s + if not isinstance(s, gym.spaces.Box) + else gym.spaces.Box( + low=-self.clip_by_value, + high=self.clip_by_value, + shape=s.shape, + dtype=s.dtype, + ) + ), + _input_observation_space_struct, + ) + if isinstance(input_observation_space, (gym.spaces.Dict, gym.spaces.Tuple)): + return type(input_observation_space)(_observation_space_struct) + else: + return _observation_space_struct + + def __init__( + self, + *, + multi_agent: bool = False, + de_mean_to_zero: bool = True, + de_std_to_one: bool = True, + clip_by_value: Optional[float] = 10.0, + update_stats: bool = True, + **kwargs, + ): + """Initializes a MeanStdFilter instance. + + Args: + multi_agent: Whether this is a connector operating on a multi-agent + observation space mapping AgentIDs to individual agents' observations. + de_mean_to_zero: Whether to transform the mean values of the output data to + 0.0. This is done by subtracting the incoming data by the currently + stored mean value. + de_std_to_one: Whether to transform the standard deviation values of the + output data to 1.0. This is done by dividing the incoming data by the + currently stored std value. + clip_by_value: If not None, clip the incoming data within the interval: + [-clip_by_value, +clip_by_value]. + update_stats: Whether to update the internal mean and std stats with each + incoming sample (with each `__call__()`) or not. You should set this to + False if you would like to perform inference in a production + environment, without continuing to "learn" stats from new data. + """ + super().__init__(**kwargs) + + self._multi_agent = multi_agent + + # We simply use the old MeanStdFilter until non-connector env_runner is fully + # deprecated to avoid duplicate code + self.de_mean_to_zero = de_mean_to_zero + self.de_std_to_one = de_std_to_one + self.clip_by_value = clip_by_value + self._update_stats = update_stats + + self._filters: Optional[Dict[AgentID, _MeanStdFilter]] = None + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + persistent_data: Optional[dict] = None, + **kwargs, + ) -> Any: + if self._filters is None: + self._init_new_filters() + + # This connector acts as a classic preprocessor. We process and then replace + # observations inside the episodes directly. Thus, all following connectors + # will only see and operate on the already normalized data (w/o having access + # anymore to the original observations). + for sa_episode in self.single_agent_episode_iterator(episodes): + sa_obs = sa_episode.get_observations(indices=-1) + try: + normalized_sa_obs = self._filters[sa_episode.agent_id]( + sa_obs, update=self._update_stats + ) + except KeyError: + raise KeyError( + "KeyError trying to access a filter by agent ID " + f"`{sa_episode.agent_id}`! You probably did NOT pass the " + f"`multi_agent=True` flag into the `MeanStdFilter()` constructor. " + ) + sa_episode.set_observations(at_indices=-1, new_data=normalized_sa_obs) + # We set the Episode's observation space to ours so that we can safely + # set the last obs to the new value (without causing a space mismatch + # error). + sa_episode.observation_space = self.observation_space + + # Leave `batch` as is. RLlib's default connector will automatically + # populate the OBS column therein from the episodes' now transformed + # observations. + return batch + + @override(ConnectorV2) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + if self._filters is None: + self._init_new_filters() + return self._get_state_from_filters(self._filters) + + @override(ConnectorV2) + def set_state(self, state: StateDict) -> None: + if self._filters is None: + self._init_new_filters() + for agent_id, agent_state in state.items(): + filter = self._filters[agent_id] + filter.shape = agent_state["shape"] + filter.demean = agent_state["de_mean_to_zero"] + filter.destd = agent_state["de_std_to_one"] + filter.clip = agent_state["clip_by_value"] + filter.running_stats = tree.unflatten_as( + filter.shape, + [RunningStat.from_state(s) for s in agent_state["running_stats"]], + ) + # Do not update the buffer. + + @override(ConnectorV2) + def reset_state(self) -> None: + """Creates copy of current state and resets accumulated state""" + if not self._update_stats: + raise ValueError( + f"State of {type(self).__name__} can only be changed when " + f"`update_stats` was set to False." + ) + self._init_new_filters() + + @override(ConnectorV2) + def merge_states(self, states: List[Dict[str, Any]]) -> Dict[str, Any]: + if self._filters is None: + self._init_new_filters() + + # Make sure data is uniform across given states. + ref = next(iter(states[0].values())) + + for state in states: + for agent_id, agent_state in state.items(): + assert ( + agent_state["shape"] == ref["shape"] + and agent_state["de_mean_to_zero"] == ref["de_mean_to_zero"] + and agent_state["de_std_to_one"] == ref["de_std_to_one"] + and agent_state["clip_by_value"] == ref["clip_by_value"] + ) + + _filter = _MeanStdFilter( + ref["shape"], + demean=ref["de_mean_to_zero"], + destd=ref["de_std_to_one"], + clip=ref["clip_by_value"], + ) + # Override running stats of the filter with the ones stored in + # `agent_state`. + _filter.buffer = tree.unflatten_as( + agent_state["shape"], + [ + RunningStat.from_state(stats) + for stats in agent_state["running_stats"] + ], + ) + + # Leave the buffers as-is, since they should always only reflect + # what has happened on the particular env runner. + self._filters[agent_id].apply_changes(_filter, with_buffer=False) + + return MeanStdFilter._get_state_from_filters(self._filters) + + def _init_new_filters(self): + filter_shape = tree.map_structure( + lambda s: ( + None if isinstance(s, (Discrete, MultiDiscrete)) else np.array(s.shape) + ), + get_base_struct_from_space(self.input_observation_space), + ) + if not self._multi_agent: + filter_shape = {None: filter_shape} + + del self._filters + self._filters = { + agent_id: _MeanStdFilter( + agent_filter_shape, + demean=self.de_mean_to_zero, + destd=self.de_std_to_one, + clip=self.clip_by_value, + ) + for agent_id, agent_filter_shape in filter_shape.items() + } + + @staticmethod + def _get_state_from_filters(filters: Dict[AgentID, Dict[str, Any]]): + ret = {} + for agent_id, agent_filter in filters.items(): + ret[agent_id] = { + "shape": agent_filter.shape, + "de_mean_to_zero": agent_filter.demean, + "de_std_to_one": agent_filter.destd, + "clip_by_value": agent_filter.clip, + "running_stats": [ + s.to_state() for s in tree.flatten(agent_filter.running_stats) + ], + } + return ret diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/observation_preprocessor.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/observation_preprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..2910173500ad8a8adbe1e728d3be6df5ff82d6e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/observation_preprocessor.py @@ -0,0 +1,181 @@ +import abc +from typing import Any, Dict, List, Optional + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class SingleAgentObservationPreprocessor(ConnectorV2, abc.ABC): + """Env-to-module connector preprocessing the most recent single-agent observation. + + This is a convenience class that simplifies the writing of few-step preprocessor + connectors. + + Note that this class also works in a multi-agent setup, in which case RLlib + separately calls this connector piece with each agents' observation and + `SingleAgentEpisode` object. + + Users must implement the `preprocess()` method, which simplifies the usual procedure + of extracting some data from a list of episodes and adding it to the batch to a mere + "old-observation --transform--> return new-observation" step. + """ + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + # Users should override this method only in case the + # `SingleAgentObservationPreprocessor` changes the observation space of the + # pipeline. In this case, return the new observation space based on the + # incoming one (`input_observation_space`). + return super().recompute_output_observation_space( + input_observation_space, input_action_space + ) + + @abc.abstractmethod + def preprocess(self, observation, episode: SingleAgentEpisode): + """Override to implement the preprocessing logic. + + Args: + observation: A single (non-batched) observation item for a single agent to + be preprocessed by this connector. + episode: The `SingleAgentEpisode` instance, from which `observation` was + taken. You can extract information on the particular AgentID and the + ModuleID through `episode.agent_id` and `episode.module_id`. + + Returns: + The new observation for the agent after `observation` has been preprocessed. + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + persistent_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # We process and then replace observations inside the episodes directly. + # Thus, all following connectors will only see and operate on the already + # processed observation (w/o having access anymore to the original + # observations). + for sa_episode in self.single_agent_episode_iterator(episodes): + observation = sa_episode.get_observations(-1) + + # Process the observation and write the new observation back into the + # episode. + new_observation = self.preprocess( + observation=observation, + episode=sa_episode, + ) + sa_episode.set_observations(at_indices=-1, new_data=new_observation) + # We set the Episode's observation space to ours so that we can safely + # set the last obs to the new value (without causing a space mismatch + # error). + sa_episode.observation_space = self.observation_space + + # Leave `batch` as is. RLlib's default connector automatically populates + # the OBS column therein from the episodes' now transformed observations. + return batch + + +@PublicAPI(stability="alpha") +class MultiAgentObservationPreprocessor(ConnectorV2, abc.ABC): + """Env-to-module connector preprocessing the most recent multi-agent observation. + + The observation is always a dict of individual agents' observations. + + This is a convenience class that simplifies the writing of few-step preprocessor + connectors. + + Users must implement the `preprocess()` method, which simplifies the usual procedure + of extracting some data from a list of episodes and adding it to the batch to a mere + "old-observation --transform--> return new-observation" step. + """ + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + # Users should override this method only in case the + # `MultiAgentObservationPreprocessor` changes the observation space of the + # pipeline. In this case, return the new observation space based on the + # incoming one (`input_observation_space`). + return super().recompute_output_observation_space( + input_observation_space, input_action_space + ) + + @abc.abstractmethod + def preprocess(self, observations, episode: MultiAgentEpisode): + """Override to implement the preprocessing logic. + + Args: + observations: An observation dict containing each stepping agents' + (non-batched) observation to be preprocessed by this connector. + episode: The MultiAgentEpisode instance, where the `observation` dict + originated from. + + Returns: + The new multi-agent observation dict after `observations` has been + preprocessed. + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + persistent_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # We process and then replace observations inside the episodes directly. + # Thus, all following connectors will only see and operate on the already + # processed observation (w/o having access anymore to the original + # observations). + for ma_episode in episodes: + observations = ma_episode.get_observations(-1) + + # Process the observation and write the new observation back into the + # episode. + new_observation = self.preprocess( + observations=observations, + episode=ma_episode, + ) + # TODO (sven): Implement set_observations API for multi-agent episodes. + # For now, we'll hack it through the single agent APIs. + # ma_episode.set_observations(at_indices=-1, new_data=new_observation) + for agent_id, obs in new_observation.items(): + ma_episode.agent_episodes[agent_id].set_observations( + at_indices=-1, + new_data=obs, + ) + # We set the Episode's observation space to ours so that we can safely + # set the last obs to the new value (without causing a space mismatch + # error). + ma_episode.observation_space = self.observation_space + + # Leave `batch` as is. RLlib's default connector automatically populates + # the OBS column therein from the episodes' now transformed observations. + return batch + + +# Backward compatibility +ObservationPreprocessor = SingleAgentObservationPreprocessor diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/prev_actions_prev_rewards.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/prev_actions_prev_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..14016971ceb7fd6b17a96abb03accecd51dedb35 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/prev_actions_prev_rewards.py @@ -0,0 +1,168 @@ +from typing import Any, Dict, List, Optional + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.spaces.space_utils import ( + batch as batch_fn, + flatten_to_single_ndarray, +) +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class PrevActionsPrevRewards(ConnectorV2): + """A connector piece that adds previous rewards and actions to the input obs. + + - Requires Columns.OBS to be already a part of the batch. + - This connector makes the assumption that under the Columns.OBS key in batch, + there is either a list of individual env observations to be flattened (single-agent + case) or a dict mapping (AgentID, ModuleID)-tuples to lists of data items to be + flattened (multi-agent case). + - Converts Columns.OBS data into a dict (or creates a sub-dict if obs are + already a dict), and adds "prev_rewards" and "prev_actions" + to this dict. The original observations are stored under the self.ORIG_OBS_KEY in + that dict. + - If your RLModule does not handle dict inputs, you will have to plug in an + `FlattenObservations` connector piece after this one. + - Does NOT work in a Learner pipeline as it operates on individual observation + items (as opposed to batched/time-ranked data). + - Therefore, assumes that the altered (flattened) observations will be written + back into the episode by a later connector piece in the env-to-module pipeline + (which this piece is part of as well). + - Only reads reward- and action information from the given list of Episode objects. + - Does NOT write any observations (or other data) to the given Episode objects. + """ + + ORIG_OBS_KEY = "_orig_obs" + PREV_ACTIONS_KEY = "prev_n_actions" + PREV_REWARDS_KEY = "prev_n_rewards" + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + if self._multi_agent: + ret = {} + for agent_id, obs_space in input_observation_space.spaces.items(): + act_space = input_action_space[agent_id] + ret[agent_id] = self._convert_individual_space(obs_space, act_space) + return gym.spaces.Dict(ret) + else: + return self._convert_individual_space( + input_observation_space, input_action_space + ) + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + multi_agent: bool = False, + n_prev_actions: int = 1, + n_prev_rewards: int = 1, + **kwargs, + ): + """Initializes a PrevActionsPrevRewards instance. + + Args: + multi_agent: Whether this is a connector operating on a multi-agent + observation space mapping AgentIDs to individual agents' observations. + n_prev_actions: The number of previous actions to include in the output + data. Discrete actions are ont-hot'd. If > 1, will concatenate the + individual action tensors. + n_prev_rewards: The number of previous rewards to include in the output + data. + """ + super().__init__( + input_observation_space=input_observation_space, + input_action_space=input_action_space, + **kwargs, + ) + + self._multi_agent = multi_agent + self.n_prev_actions = n_prev_actions + self.n_prev_rewards = n_prev_rewards + + # TODO: Move into input_observation_space setter + # Thus far, this connector piece only operates on discrete action spaces. + # act_spaces = [self.input_action_space] + # if self._multi_agent: + # act_spaces = self.input_action_space.spaces.values() + # if not all(isinstance(s, gym.spaces.Discrete) for s in act_spaces): + # raise ValueError( + # f"{type(self).__name__} only works on Discrete action spaces " + # f"thus far (or, for multi-agent, on Dict spaces mapping AgentIDs to " + # f"the individual agents' Discrete action spaces)!" + # ) + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + for sa_episode in self.single_agent_episode_iterator( + episodes, agents_that_stepped_only=True + ): + # Episode is not numpy'ized yet and thus still operates on lists of items. + assert not sa_episode.is_numpy + + augmented_obs = {self.ORIG_OBS_KEY: sa_episode.get_observations(-1)} + + if self.n_prev_actions: + augmented_obs[self.PREV_ACTIONS_KEY] = flatten_to_single_ndarray( + batch_fn( + sa_episode.get_actions( + indices=slice(-self.n_prev_actions, None), + fill=0.0, + one_hot_discrete=True, + ) + ) + ) + + if self.n_prev_rewards: + augmented_obs[self.PREV_REWARDS_KEY] = np.array( + sa_episode.get_rewards( + indices=slice(-self.n_prev_rewards, None), + fill=0.0, + ) + ) + + # Write new observation directly back into the episode. + sa_episode.set_observations(at_indices=-1, new_data=augmented_obs) + # We set the Episode's observation space to ours so that we can safely + # set the last obs to the new value (without causing a space mismatch + # error). + sa_episode.observation_space = self.observation_space + + return batch + + def _convert_individual_space(self, obs_space, act_space): + return gym.spaces.Dict( + { + self.ORIG_OBS_KEY: obs_space, + # Currently only works for Discrete action spaces. + self.PREV_ACTIONS_KEY: Box( + 0.0, 1.0, (act_space.n * self.n_prev_actions,), np.float32 + ), + self.PREV_REWARDS_KEY: Box( + float("-inf"), + float("inf"), + (self.n_prev_rewards,), + np.float32, + ), + } + ) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/write_observations_to_episodes.py b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/write_observations_to_episodes.py new file mode 100644 index 0000000000000000000000000000000000000000..9b92da4984fd20cb8fe704692e1473eab4b925a4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/env_to_module/write_observations_to_episodes.py @@ -0,0 +1,131 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class WriteObservationsToEpisodes(ConnectorV2): + """Writes the observations from the batch into the running episodes. + + Note: This is one of the default env-to-module ConnectorV2 pieces that are added + automatically by RLlib into every env-to-module connector pipelines, unless + `config.add_default_connectors_to_env_to_module_pipeline` is set to False. + + The default env-to-module connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + This ConnectorV2: + - Operates on a batch that already has observations in it and a list of Episode + objects. + - Writes the observation(s) from the batch to all the given episodes. Thereby + the number of observations in the batch must match the length of the list of + episodes given. + - Does NOT alter any observations (or other data) in the batch. + - Can only be used in an EnvToModule pipeline (writing into Episode objects in a + Learner pipeline does not make a lot of sense as - after the learner update - the + list of episodes is discarded). + + .. testcode:: + + import gymnasium as gym + import numpy as np + + from ray.rllib.connectors.env_to_module import WriteObservationsToEpisodes + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Assume we have two episodes (vectorized), then our forward batch will carry + # two observation records (batch size = 2). + # The connector in this example will write these two (possibly transformed) + # observations back into the two respective SingleAgentEpisode objects. + batch = { + "obs": [np.array([0.0, 1.0], np.float32), np.array([2.0, 3.0], np.float32)], + } + + # Our two episodes have one observation each (i.e. the reset one). This is the + # one that will be overwritten by the connector in this example. + obs_space = gym.spaces.Box(-10.0, 10.0, (2,), np.float32) + act_space = gym.spaces.Discrete(2) + episodes = [ + SingleAgentEpisode( + observation_space=obs_space, + observations=[np.array([-10, -20], np.float32)], + len_lookback_buffer=0, + ) for _ in range(2) + ] + # Make sure everything is setup correctly. + check(episodes[0].get_observations(0), [-10.0, -20.0]) + check(episodes[1].get_observations(-1), [-10.0, -20.0]) + + # Create our connector piece. + connector = WriteObservationsToEpisodes(obs_space, act_space) + + # Call the connector (and thereby write the transformed observations back + # into the episodes). + output_batch = connector( + rl_module=None, # This particular connector works without an RLModule. + batch=batch, + episodes=episodes, + explore=True, + shared_data={}, + ) + + # The connector does NOT change the data batch being passed through. + check(output_batch, batch) + + # However, the connector has overwritten the last observations in the episodes. + check(episodes[0].get_observations(-1), [0.0, 1.0]) + check(episodes[1].get_observations(0), [2.0, 3.0]) + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + observations = batch.get(Columns.OBS) + + if observations is None: + raise ValueError( + f"`batch` must already have a column named {Columns.OBS} in it " + f"for this connector to work!" + ) + + # Note that the following loop works with multi-agent as well as with + # single-agent episode, as long as the following conditions are met (these + # will be validated by `self.single_agent_episode_iterator()`): + # - Per single agent episode, one observation item is expected to exist in + # `data`, either in a list directly under the "obs" key OR for multi-agent: + # in a list sitting under a key `(agent_id, module_id)` of a dict sitting + # under the "obs" key. + for sa_episode, obs in self.single_agent_episode_iterator( + episodes=episodes, zip_with_batch_column=observations + ): + # Make sure episodes are NOT numpy'ized yet (we are expecting to run in an + # env-to-module pipeline). + assert not sa_episode.is_numpy + # Write new information into the episode. + sa_episode.set_observations(at_indices=-1, new_data=obs) + # Change the observation space of the sa_episode. + sa_episode.observation_space = self.observation_space + + # Return the unchanged `batch`. + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__init__.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1bf09d5c9d67493ebfa30dbe3b2f7494efe5ae67 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__init__.py @@ -0,0 +1,47 @@ +from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import ( + AddObservationsFromEpisodesToBatch, +) +from ray.rllib.connectors.common.add_states_from_episodes_to_batch import ( + AddStatesFromEpisodesToBatch, +) +from ray.rllib.connectors.common.add_time_dim_to_batch_and_zero_pad import ( + AddTimeDimToBatchAndZeroPad, +) +from ray.rllib.connectors.common.agent_to_module_mapping import AgentToModuleMapping +from ray.rllib.connectors.common.batch_individual_items import BatchIndividualItems +from ray.rllib.connectors.common.numpy_to_tensor import NumpyToTensor +from ray.rllib.connectors.learner.add_columns_from_episodes_to_train_batch import ( + AddColumnsFromEpisodesToTrainBatch, +) +from ray.rllib.connectors.learner.add_infos_from_episodes_to_train_batch import ( + AddInfosFromEpisodesToTrainBatch, +) +from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import ( # noqa + AddNextObservationsFromEpisodesToTrainBatch, +) +from ray.rllib.connectors.learner.add_one_ts_to_episodes_and_truncate import ( + AddOneTsToEpisodesAndTruncate, +) +from ray.rllib.connectors.learner.compute_returns_to_go import ComputeReturnsToGo +from ray.rllib.connectors.learner.general_advantage_estimation import ( + GeneralAdvantageEstimation, +) +from ray.rllib.connectors.learner.learner_connector_pipeline import ( + LearnerConnectorPipeline, +) + +__all__ = [ + "AddColumnsFromEpisodesToTrainBatch", + "AddInfosFromEpisodesToTrainBatch", + "AddNextObservationsFromEpisodesToTrainBatch", + "AddObservationsFromEpisodesToBatch", + "AddOneTsToEpisodesAndTruncate", + "AddStatesFromEpisodesToBatch", + "AddTimeDimToBatchAndZeroPad", + "AgentToModuleMapping", + "BatchIndividualItems", + "ComputeReturnsToGo", + "GeneralAdvantageEstimation", + "LearnerConnectorPipeline", + "NumpyToTensor", +] diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62b68ca7df7bc397081f585620b07b3bb2d82eb3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_columns_from_episodes_to_train_batch.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_columns_from_episodes_to_train_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7dbb0dc2189fb7e0c1c01ecd37712d57b7f743dc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_columns_from_episodes_to_train_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_infos_from_episodes_to_train_batch.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_infos_from_episodes_to_train_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84c1197386b72857dc21ec0b7000e37341402c82 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_infos_from_episodes_to_train_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_next_observations_from_episodes_to_train_batch.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_next_observations_from_episodes_to_train_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c49150ffffac26d02826f591b995db954c7eaae3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_next_observations_from_episodes_to_train_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_one_ts_to_episodes_and_truncate.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_one_ts_to_episodes_and_truncate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f987a9c46bdc0aa0ae9ce1ea2d22ebc7865138b6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/add_one_ts_to_episodes_and_truncate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/compute_returns_to_go.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/compute_returns_to_go.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7517c7d506fdaaff515a5cde6ed53e84da3576de Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/compute_returns_to_go.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/frame_stacking.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/frame_stacking.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7f730c8b11286ca2a640fbe2a208b76139827ea Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/frame_stacking.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/general_advantage_estimation.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/general_advantage_estimation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf96a822163598ae34c455f89503490429df3c80 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/general_advantage_estimation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/learner_connector_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/learner_connector_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69514884e52c95f0738665046f4124819eef4299 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/learner/__pycache__/learner_connector_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_columns_from_episodes_to_train_batch.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_columns_from_episodes_to_train_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..b55729c8231406b669bee077fa5d5b3862732db1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_columns_from_episodes_to_train_batch.py @@ -0,0 +1,167 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AddColumnsFromEpisodesToTrainBatch(ConnectorV2): + """Adds actions/rewards/terminateds/... to train batch. Excluding the infos column. + + Note: This is one of the default Learner ConnectorV2 pieces that are added + automatically by RLlib into every Learner connector pipeline, unless + `config.add_default_connectors_to_learner_pipeline` is set to False. + + The default Learner connector pipeline is: + [ + [0 or more user defined ConnectorV2 pieces], + AddObservationsFromEpisodesToBatch, + AddColumnsFromEpisodesToTrainBatch, + AddTimeDimToBatchAndZeroPad, + AddStatesFromEpisodesToBatch, + AgentToModuleMapping, # only in multi-agent setups! + BatchIndividualItems, + NumpyToTensor, + ] + + Does NOT add observations or infos to train batch. + Observations should have already been added by another ConnectorV2 piece: + `AddObservationsToTrainBatch` in the same pipeline. + Infos can be added manually by the user through setting this in the config: + + .. testcode:: + :skipif: True + + from ray.rllib.connectors.learner import AddInfosFromEpisodesToTrainBatch + config.training( + learner_connector=lambda obs_sp, act_sp: AddInfosFromEpisodesToTrainBatch() + )` + + If provided with `episodes` data, this connector piece makes sure that the final + train batch going into the RLModule for updating (`forward_train()` call) contains + at the minimum: + - Observations: From all episodes under the Columns.OBS key. + - Actions, rewards, terminal/truncation flags: From all episodes under the + respective keys. + - All data inside the episodes' `extra_model_outs` property, e.g. action logp and + action probs under the respective keys. + - Internal states: These will NOT be added to the batch by this connector piece + as this functionality is handled by a different default connector piece: + `AddStatesFromEpisodesToBatch`. + + If the user wants to customize their own data under the given keys (e.g. obs, + actions, ...), they can extract from the episodes or recompute from `data` + their own data and store it in `data` under those keys. In this case, the default + connector will not change the data under these keys and simply act as a + pass-through. + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + + # Actions. + if Columns.ACTIONS not in batch: + for sa_episode in self.single_agent_episode_iterator( + episodes, + agents_that_stepped_only=False, + ): + self.add_n_batch_items( + batch, + Columns.ACTIONS, + items_to_add=[ + sa_episode.get_actions(indices=ts) + for ts in range(len(sa_episode)) + ], + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + # Rewards. + if Columns.REWARDS not in batch: + for sa_episode in self.single_agent_episode_iterator( + episodes, + agents_that_stepped_only=False, + ): + self.add_n_batch_items( + batch, + Columns.REWARDS, + items_to_add=[ + sa_episode.get_rewards(indices=ts) + for ts in range(len(sa_episode)) + ], + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + # Terminateds. + if Columns.TERMINATEDS not in batch: + for sa_episode in self.single_agent_episode_iterator( + episodes, + agents_that_stepped_only=False, + ): + self.add_n_batch_items( + batch, + Columns.TERMINATEDS, + items_to_add=( + [False] * (len(sa_episode) - 1) + [sa_episode.is_terminated] + if len(sa_episode) > 0 + else [] + ), + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + # Truncateds. + if Columns.TRUNCATEDS not in batch: + for sa_episode in self.single_agent_episode_iterator( + episodes, + agents_that_stepped_only=False, + ): + self.add_n_batch_items( + batch, + Columns.TRUNCATEDS, + items_to_add=( + [False] * (len(sa_episode) - 1) + [sa_episode.is_truncated] + if len(sa_episode) > 0 + else [] + ), + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + # Extra model outputs (except for STATE_OUT, which will be handled by another + # default connector piece). Also, like with all the fields above, skip + # those that the user already seemed to have populated via custom connector + # pieces. + skip_columns = set(batch.keys()) | {Columns.STATE_IN, Columns.STATE_OUT} + for sa_episode in self.single_agent_episode_iterator( + episodes, + agents_that_stepped_only=False, + ): + for column in sa_episode.extra_model_outputs.keys(): + if column not in skip_columns: + self.add_n_batch_items( + batch, + column, + items_to_add=[ + sa_episode.get_extra_model_outputs(key=column, indices=ts) + for ts in range(len(sa_episode)) + ], + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_infos_from_episodes_to_train_batch.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_infos_from_episodes_to_train_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..145d6d86acb061d7e8a6ffe4808a67b81b216340 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_infos_from_episodes_to_train_batch.py @@ -0,0 +1,51 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AddInfosFromEpisodesToTrainBatch(ConnectorV2): + """Adds the infos column to th train batch. + + If provided with `episodes` data, this connector piece makes sure that the final + train batch going into the RLModule for updating (`forward_train()` call) contains + an `infos` column. + + If the user wants to customize their own data under the given keys (e.g. obs, + actions, ...), they can extract from the episodes or recompute from `data` + their own data and store it in `data` under those keys. In this case, the default + connector will not change the data under these keys and simply act as a + pass-through. + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # Infos. + if Columns.INFOS not in batch: + for sa_episode in self.single_agent_episode_iterator( + episodes, + agents_that_stepped_only=False, + ): + self.add_n_batch_items( + batch, + Columns.INFOS, + items_to_add=sa_episode.get_infos(slice(0, len(sa_episode))), + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_next_observations_from_episodes_to_train_batch.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_next_observations_from_episodes_to_train_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..e0ea6abfb234fb19c6cf321a8955e3b886498fbc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_next_observations_from_episodes_to_train_batch.py @@ -0,0 +1,103 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AddNextObservationsFromEpisodesToTrainBatch(ConnectorV2): + """Adds the NEXT_OBS column with the correct episode observations to train batch. + + - Operates on a list of Episode objects. + - Gets all observation(s) from all the given episodes (except the very first ones) + and adds them to the batch under construction in the NEXT_OBS column (as a list of + individual observations). + - Does NOT alter any observations (or other data) in the given episodes. + - Can be used in Learner connector pipelines. + + .. testcode:: + + import gymnasium as gym + import numpy as np + + from ray.rllib.connectors.learner import ( + AddNextObservationsFromEpisodesToTrainBatch + ) + from ray.rllib.core.columns import Columns + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Create two dummy SingleAgentEpisodes, each containing 3 observations, + # 2 actions and 2 rewards (both episodes are length=2). + obs_space = gym.spaces.Box(-1.0, 1.0, (2,), np.float32) + act_space = gym.spaces.Discrete(2) + + episodes = [SingleAgentEpisode( + observations=[obs_space.sample(), obs_space.sample(), obs_space.sample()], + actions=[act_space.sample(), act_space.sample()], + rewards=[1.0, 2.0], + len_lookback_buffer=0, + ) for _ in range(2)] + eps_1_next_obses = episodes[0].get_observations([1, 2]) + eps_2_next_obses = episodes[1].get_observations([1, 2]) + print(f"1st Episode's next obses are {eps_1_next_obses}") + print(f"2nd Episode's next obses are {eps_2_next_obses}") + + # Create an instance of this class. + connector = AddNextObservationsFromEpisodesToTrainBatch() + + # Call the connector with the two created episodes. + # Note that this particular connector works without an RLModule, so we + # simplify here for the sake of this example. + output_data = connector( + rl_module=None, + batch={}, + episodes=episodes, + explore=True, + shared_data={}, + ) + # The output data should now contain the last observations of both episodes, + # in a "per-episode organized" fashion. + check( + output_data, + { + Columns.NEXT_OBS: { + (episodes[0].id_,): eps_1_next_obses, + (episodes[1].id_,): eps_2_next_obses, + }, + }, + ) + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # If "obs" already in `batch`, early out. + if Columns.NEXT_OBS in batch: + return batch + + for sa_episode in self.single_agent_episode_iterator( + # This is a Learner-only connector -> Get all episodes (for train batch). + episodes, + agents_that_stepped_only=False, + ): + self.add_n_batch_items( + batch, + Columns.NEXT_OBS, + items_to_add=sa_episode.get_observations(slice(1, len(sa_episode) + 1)), + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_one_ts_to_episodes_and_truncate.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_one_ts_to_episodes_and_truncate.py new file mode 100644 index 0000000000000000000000000000000000000000..fcd3703eeb855e1ab1fad584c1aa01b5f922f5cb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/add_one_ts_to_episodes_and_truncate.py @@ -0,0 +1,168 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.utils.annotations import override +from ray.rllib.utils.postprocessing.episodes import add_one_ts_to_episodes_and_truncate +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class AddOneTsToEpisodesAndTruncate(ConnectorV2): + """Adds an artificial timestep to all incoming episodes at the end. + + In detail: The last observations, infos, actions, and all `extra_model_outputs` + will be duplicated and appended to each episode's data. An extra 0.0 reward + will be appended to the episode's rewards. The episode's timestep will be + increased by 1. Also, adds the truncated=True flag to each episode if the + episode is not already done (terminated or truncated). + + Useful for value function bootstrapping, where it is required to compute a + forward pass for the very last timestep within the episode, + i.e. using the following input dict: { + obs=[final obs], + state=[final state output], + prev. reward=[final reward], + etc.. + } + + .. testcode:: + + from ray.rllib.connectors.learner import AddOneTsToEpisodesAndTruncate + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Create 2 episodes (both to be extended by one timestep). + episode1 = SingleAgentEpisode( + observations=[0, 1, 2], + actions=[0, 1], + rewards=[0.0, 1.0], + terminated=False, + truncated=False, + len_lookback_buffer=0, + ).to_numpy() + check(len(episode1), 2) + check(episode1.is_truncated, False) + + episode2 = SingleAgentEpisode( + observations=[0, 1, 2, 3, 4, 5], + actions=[0, 1, 2, 3, 4], + rewards=[0.0, 1.0, 2.0, 3.0, 4.0], + terminated=True, # a terminated episode + truncated=False, + len_lookback_buffer=0, + ).to_numpy() + check(len(episode2), 5) + check(episode2.is_truncated, False) + check(episode2.is_terminated, True) + + # Create an instance of this class. + connector = AddOneTsToEpisodesAndTruncate() + + # Call the connector. + shared_data = {} + _ = connector( + rl_module=None, # Connector used here does not require RLModule. + batch={}, + episodes=[episode1, episode2], + shared_data=shared_data, + ) + # Check on the episodes. Both of them should now be 1 timestep longer. + check(len(episode1), 3) + check(episode1.is_truncated, True) + check(len(episode2), 6) + check(episode2.is_truncated, False) + check(episode2.is_terminated, True) + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # Build the loss mask to make sure the extra added timesteps do not influence + # the final loss and fix the terminateds and truncateds in the batch. + + # For proper v-trace execution, the rules must be as follows: + # Legend: + # T: terminal=True + # R: truncated=True + # B0: bootstrap with value 0 (also: terminal=True) + # Bx: bootstrap with some vf-computed value (also: terminal=True) + + # batch: - - - - - - - T B0- - - - - R Bx- - - - R Bx + # mask : t t t t t t t t f t t t t t t f t t t t t f + + # TODO (sven): Same situation as in TODO below, but for multi-agent episode. + # Maybe add a dedicated connector piece for this task? + # We extend the MultiAgentEpisode's ID by a running number here to make sure + # we treat each MAEpisode chunk as separate (for potentially upcoming v-trace + # and LSTM zero-padding) and don't mix data from different chunks. + if isinstance(episodes[0], MultiAgentEpisode): + for i, ma_episode in enumerate(episodes): + ma_episode.id_ += "_" + str(i) + # Also change the underlying single-agent episode's + # `multi_agent_episode_id` properties. + for sa_episode in ma_episode.agent_episodes.values(): + sa_episode.multi_agent_episode_id = ma_episode.id_ + + for i, sa_episode in enumerate( + self.single_agent_episode_iterator(episodes, agents_that_stepped_only=False) + ): + # TODO (sven): This is a little bit of a hack: By extending the Episode's + # ID, we make sure that each episode chunk in `episodes` is treated as a + # separate episode in the `self.add_n_batch_items` below. Some algos (e.g. + # APPO) may have >1 episode chunks from the same episode (same ID) in the + # training data, thus leading to a malformatted batch in case of + # RNN-triggered zero-padding of the train batch. + # For example, if e1 (id=a len=4) and e2 (id=a len=5) are two chunks of the + # same episode in `episodes`, the resulting batch would have an additional + # timestep in the middle of the episode's "row": + # { "obs": { + # ("a", <- eps ID): [0, 1, 2, 3 <- len=4, [additional 1 ts (bad)], + # 0, 1, 2, 3, 4 <- len=5, [additional 1 ts]] + # }} + sa_episode.id_ += "_" + str(i) + + len_ = len(sa_episode) + + # Extend all episodes by one ts. + add_one_ts_to_episodes_and_truncate([sa_episode]) + + loss_mask = [True for _ in range(len_)] + [False] + self.add_n_batch_items( + batch, + Columns.LOSS_MASK, + loss_mask, + len_ + 1, + sa_episode, + ) + + terminateds = ( + [False for _ in range(len_ - 1)] + + [bool(sa_episode.is_terminated)] + + [True] # extra timestep + ) + self.add_n_batch_items( + batch, + Columns.TERMINATEDS, + terminateds, + len_ + 1, + sa_episode, + ) + + # Signal to following connector pieces that the loss-mask which masks out + # invalid episode ts (for the extra added ts at the end) has already been + # added to `data`. + shared_data["_added_loss_mask_for_valid_episode_ts"] = True + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/compute_returns_to_go.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/compute_returns_to_go.py new file mode 100644 index 0000000000000000000000000000000000000000..6a2c860559bc9ea5bb610ef63353fac32873c4dc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/compute_returns_to_go.py @@ -0,0 +1,68 @@ +from typing import Any, Dict, List + +import scipy + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.utils.typing import EpisodeType + + +class ComputeReturnsToGo(ConnectorV2): + """Learner ConnectorV2 piece computing discounted returns to go till end of episode. + + This ConnectorV2: + - Operates on a list of Episode objects (single- or multi-agent). + - Should be used only in the Learner pipeline as a preparation for an upcoming loss + computation that requires the discounted returns to go (until the end of the + episode). + - For each agent, for each episode and at each timestep, sums up the rewards + (discounted) until the end of the episode and assigns the results to a new + column: RETURNS_TO_GO in the batch. + """ + + def __init__( + self, + input_observation_space=None, + input_action_space=None, + *, + gamma, + ): + """Initializes a ComputeReturnsToGo instance. + + Args: + gamma: The discount factor gamma. + """ + super().__init__(input_observation_space, input_action_space) + self.gamma = gamma + + def __call__( + self, + *, + rl_module: MultiRLModule, + episodes: List[EpisodeType], + batch: Dict[str, Any], + **kwargs, + ): + for sa_episode in self.single_agent_episode_iterator( + episodes, agents_that_stepped_only=False + ): + # Reverse the rewards sequence. + rewards_reversed = sa_episode.get_rewards()[::-1] + # Use lfilter to compute the discounted cumulative sums. + discounted_cumsum_reversed = scipy.signal.lfilter( + [1], [1, -self.gamma], rewards_reversed + ) + # Reverse the result to get the correct order. + discounted_returns = discounted_cumsum_reversed[::-1] + + # Add the results to the batch under a new column: RETURNS_TO_GO. + self.add_n_batch_items( + batch=batch, + column=Columns.RETURNS_TO_GO, + items_to_add=discounted_returns, + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/frame_stacking.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/frame_stacking.py new file mode 100644 index 0000000000000000000000000000000000000000..012a630f7d4a099231f381e3d1999ed6e30a46f2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/frame_stacking.py @@ -0,0 +1,5 @@ +from functools import partial + +from ray.rllib.connectors.common.frame_stacking import FrameStacking + +FrameStackingLearner = partial(FrameStacking, as_learner_connector=True) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/general_advantage_estimation.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/general_advantage_estimation.py new file mode 100644 index 0000000000000000000000000000000000000000..5d7f278e112b32b1249cc419e154947306f45b84 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/general_advantage_estimation.py @@ -0,0 +1,196 @@ +from typing import Any, Dict, List + +import numpy as np + +from ray.rllib.connectors.common.numpy_to_tensor import NumpyToTensor +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.evaluation.postprocessing import Postprocessing +from ray.rllib.utils.annotations import override +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.postprocessing.value_predictions import compute_value_targets +from ray.rllib.utils.postprocessing.zero_padding import ( + split_and_zero_pad_n_episodes, + unpad_data_if_necessary, +) +from ray.rllib.utils.typing import EpisodeType + + +class GeneralAdvantageEstimation(ConnectorV2): + """Learner ConnectorV2 piece computing GAE advantages and value targets on episodes. + + This ConnectorV2: + - Operates on a list of Episode objects (single- or multi-agent). + - Should be used only in the Learner pipeline and as one of the last pieces (due + to the fact that it requires the batch for the value functions to be already + complete). + - Requires the incoming episodes to already be elongated by one artificial timestep + at the end (last obs, actions, states, etc.. repeated, last reward=0.0, etc..), + making it possible to combine the per-timestep value computations with the + necessary "bootstrap" value computations at the episode (chunk) truncation points. + The extra timestep should be added using the `ray.rllib.connectors.learner. + add_one_ts_to_episodes_and_truncate.AddOneTsToEpisodesAndTruncate` connector piece. + + The GAE computation is performed in an efficient way through using the arriving + `batch` as forward batch for the value function, extracting the bootstrap values + (at the artificially added time steos) and all other value predictions (all other + timesteps), performing GAE, and adding the results back into `batch` (under + Postprocessing.ADVANTAGES and Postprocessing.VALUE_TARGETS. + """ + + def __init__( + self, + input_observation_space=None, + input_action_space=None, + *, + gamma, + lambda_, + ): + """Initializes a GeneralAdvantageEstimation instance. + + Args: + gamma: The discount factor gamma. + lambda_: The lambda parameter for General Advantage Estimation (GAE). + Defines the exponential weight used between actually measured rewards + vs value function estimates over multiple time steps. Specifically, + `lambda_` balances short-term, low-variance estimates with longer-term, + high-variance returns. A `lambda_` or 0.0 makes the GAE rely only on + immediate rewards (and vf predictions from there on, reducing variance, + but increasing bias), while a `lambda_` of 1.0 only incorporates vf + predictions at the truncation points of the given episodes or episode + chunks (reducing bias but increasing variance). + """ + super().__init__(input_observation_space, input_action_space) + self.gamma = gamma + self.lambda_ = lambda_ + + # Internal numpy-to-tensor connector to translate GAE results (advantages and + # vf targets) into tensors. + self._numpy_to_tensor_connector = None + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: MultiRLModule, + episodes: List[EpisodeType], + batch: Dict[str, Any], + **kwargs, + ): + # Device to place all GAE result tensors (advantages and value targets) on. + device = None + + # Extract all single-agent episodes. + sa_episodes_list = list( + self.single_agent_episode_iterator(episodes, agents_that_stepped_only=False) + ) + # Perform the value nets' forward passes. + # TODO (sven): We need to check here in the pipeline already, whether a module + # should even be updated or not (which we usually do after(!) the Learner + # pipeline). This is an open TODO to move this filter into a connector as well. + # For now, we'll just check, whether `mid` is in batch and skip if it isn't. + vf_preds = rl_module.foreach_module( + func=lambda mid, module: ( + module.compute_values(batch[mid]) + if mid in batch and isinstance(module, ValueFunctionAPI) + else None + ), + return_dict=True, + ) + # Loop through all modules and perform each one's GAE computation. + for module_id, module_vf_preds in vf_preds.items(): + # Skip those outputs of RLModules that are not implementers of + # `ValueFunctionAPI`. + if module_vf_preds is None: + continue + + module = rl_module[module_id] + device = module_vf_preds.device + # Convert to numpy for the upcoming GAE computations. + module_vf_preds = convert_to_numpy(module_vf_preds) + + # Collect (single-agent) episode lengths for this particular module. + episode_lens = [ + len(e) for e in sa_episodes_list if e.module_id in [None, module_id] + ] + + # Remove all zero-padding again, if applicable, for the upcoming + # GAE computations. + module_vf_preds = unpad_data_if_necessary(episode_lens, module_vf_preds) + # Compute value targets. + module_value_targets = compute_value_targets( + values=module_vf_preds, + rewards=unpad_data_if_necessary( + episode_lens, + convert_to_numpy(batch[module_id][Columns.REWARDS]), + ), + terminateds=unpad_data_if_necessary( + episode_lens, + convert_to_numpy(batch[module_id][Columns.TERMINATEDS]), + ), + truncateds=unpad_data_if_necessary( + episode_lens, + convert_to_numpy(batch[module_id][Columns.TRUNCATEDS]), + ), + gamma=self.gamma, + lambda_=self.lambda_, + ) + assert module_value_targets.shape[0] == sum(episode_lens) + + module_advantages = module_value_targets - module_vf_preds + # Drop vf-preds, not needed in loss. Note that in the DefaultPPORLModule, + # vf-preds are recomputed with each `forward_train` call anyway to compute + # the vf loss. + # Standardize advantages (used for more stable and better weighted + # policy gradient computations). + module_advantages = (module_advantages - module_advantages.mean()) / max( + 1e-4, module_advantages.std() + ) + + # Zero-pad the new computations, if necessary. + if module.is_stateful(): + module_advantages = np.stack( + split_and_zero_pad_n_episodes( + module_advantages, + episode_lens=episode_lens, + max_seq_len=module.model_config["max_seq_len"], + ), + axis=0, + ) + module_value_targets = np.stack( + split_and_zero_pad_n_episodes( + module_value_targets, + episode_lens=episode_lens, + max_seq_len=module.model_config["max_seq_len"], + ), + axis=0, + ) + batch[module_id][Postprocessing.ADVANTAGES] = module_advantages + batch[module_id][Postprocessing.VALUE_TARGETS] = module_value_targets + + # Convert all GAE results to tensors. + if self._numpy_to_tensor_connector is None: + self._numpy_to_tensor_connector = NumpyToTensor( + as_learner_connector=True, device=device + ) + tensor_results = self._numpy_to_tensor_connector( + rl_module=rl_module, + batch={ + mid: { + Postprocessing.ADVANTAGES: module_batch[Postprocessing.ADVANTAGES], + Postprocessing.VALUE_TARGETS: ( + module_batch[Postprocessing.VALUE_TARGETS] + ), + } + for mid, module_batch in batch.items() + if vf_preds[mid] is not None + }, + episodes=episodes, + ) + # Move converted tensors back to `batch`. + for mid, module_batch in tensor_results.items(): + batch[mid].update(module_batch) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/learner/learner_connector_pipeline.py b/lib/python3.12/site-packages/ray/rllib/connectors/learner/learner_connector_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..98f9e097044966c1ab434569f4a7e8fa707e3bc6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/learner/learner_connector_pipeline.py @@ -0,0 +1,61 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_pipeline_v2 import ConnectorPipelineV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics import ( + ALL_MODULES, + LEARNER_CONNECTOR, + LEARNER_CONNECTOR_SUM_EPISODES_LENGTH_IN, + LEARNER_CONNECTOR_SUM_EPISODES_LENGTH_OUT, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class LearnerConnectorPipeline(ConnectorPipelineV2): + @override(ConnectorPipelineV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]] = None, + episodes: List[EpisodeType], + explore: bool = False, + shared_data: Optional[dict] = None, + metrics: Optional[MetricsLogger] = None, + **kwargs, + ): + # Log the sum of lengths of all episodes incoming. + if metrics: + metrics.log_value( + (ALL_MODULES, LEARNER_CONNECTOR_SUM_EPISODES_LENGTH_IN), + sum(map(len, episodes)), + ) + + # Make sure user does not necessarily send initial input into this pipeline. + # Might just be empty and to be populated from `episodes`. + ret = super().__call__( + rl_module=rl_module, + batch=batch if batch is not None else {}, + episodes=episodes, + shared_data=shared_data if shared_data is not None else {}, + explore=explore, + metrics=metrics, + metrics_prefix_key=( + ALL_MODULES, + LEARNER_CONNECTOR, + ), + **kwargs, + ) + + # Log the sum of lengths of all episodes outgoing. + if metrics: + metrics.log_value( + (ALL_MODULES, LEARNER_CONNECTOR_SUM_EPISODES_LENGTH_OUT), + sum(map(len, episodes)), + ) + + return ret diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__init__.py b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6f6d55364221cbbdf8deb16148aa3d7dee97814a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__init__.py @@ -0,0 +1,29 @@ +from ray.rllib.connectors.common.module_to_agent_unmapping import ModuleToAgentUnmapping +from ray.rllib.connectors.common.tensor_to_numpy import TensorToNumpy +from ray.rllib.connectors.module_to_env.get_actions import GetActions +from ray.rllib.connectors.module_to_env.listify_data_for_vector_env import ( + ListifyDataForVectorEnv, +) +from ray.rllib.connectors.module_to_env.module_to_env_pipeline import ( + ModuleToEnvPipeline, +) +from ray.rllib.connectors.module_to_env.normalize_and_clip_actions import ( + NormalizeAndClipActions, +) +from ray.rllib.connectors.module_to_env.remove_single_ts_time_rank_from_batch import ( + RemoveSingleTsTimeRankFromBatch, +) +from ray.rllib.connectors.module_to_env.unbatch_to_individual_items import ( + UnBatchToIndividualItems, +) + +__all__ = [ + "GetActions", + "ListifyDataForVectorEnv", + "ModuleToAgentUnmapping", + "ModuleToEnvPipeline", + "NormalizeAndClipActions", + "RemoveSingleTsTimeRankFromBatch", + "TensorToNumpy", + "UnBatchToIndividualItems", +] diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..331f649216618d45ca36574fee97ecba21fda0e3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/get_actions.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/get_actions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..827ea64c20b2049b73fe16bf2c2395beb8da13a1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/get_actions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/listify_data_for_vector_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/listify_data_for_vector_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c24f5e0458ed5b9a3963a0ba1076e3a412459d03 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/listify_data_for_vector_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/module_to_env_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/module_to_env_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a5a69870e93a82dc08d9a01b6e504008267d46b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/module_to_env_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/normalize_and_clip_actions.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/normalize_and_clip_actions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b24d0953a5ebdff0258b25c38ae2e54fee74379f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/normalize_and_clip_actions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/remove_single_ts_time_rank_from_batch.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/remove_single_ts_time_rank_from_batch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e29cfd9fab8a0e4bd212ff71acdb63247c7e6c2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/remove_single_ts_time_rank_from_batch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/unbatch_to_individual_items.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/unbatch_to_individual_items.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57fb27bf4f08f1116ccc3441bd84f9086b4fb0d8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/__pycache__/unbatch_to_individual_items.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/get_actions.py b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/get_actions.py new file mode 100644 index 0000000000000000000000000000000000000000..1e862231b4fbd10500ffee1f4b03f604a8af2472 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/get_actions.py @@ -0,0 +1,91 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class GetActions(ConnectorV2): + """Connector piece sampling actions from ACTION_DIST_INPUTS from an RLModule. + + Note: This is one of the default module-to-env ConnectorV2 pieces that + are added automatically by RLlib into every module-to-env connector pipeline, + unless `config.add_default_connectors_to_module_to_env_pipeline` is set to + False. + + The default module-to-env connector pipeline is: + [ + GetActions, + TensorToNumpy, + UnBatchToIndividualItems, + ModuleToAgentUnmapping, # only in multi-agent setups! + RemoveSingleTsTimeRankFromBatch, + + [0 or more user defined ConnectorV2 pieces], + + NormalizeAndClipActions, + ListifyDataForVectorEnv, + ] + + If necessary, this connector samples actions, given action dist. inputs and a + dist. class. + The connector will only sample from the action distribution, if the + Columns.ACTIONS key cannot be found in `data`. Otherwise, it'll behave + as pass-through. If Columns.ACTIONS is NOT present in `data`, but + Columns.ACTION_DIST_INPUTS is, this connector will create a new action + distribution using the given RLModule and sample from its distribution class + (deterministically, if we are not exploring, stochastically, if we are). + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + is_multi_agent = isinstance(episodes[0], MultiAgentEpisode) + + if is_multi_agent: + for module_id, module_data in batch.copy().items(): + self._get_actions(module_data, rl_module[module_id], explore) + else: + self._get_actions(batch, rl_module, explore) + + return batch + + def _get_actions(self, batch, sa_rl_module, explore): + # Action have already been sampled -> Early out. + if Columns.ACTIONS in batch: + return + + # ACTION_DIST_INPUTS field returned by `forward_exploration|inference()` -> + # Create a new action distribution object. + if Columns.ACTION_DIST_INPUTS in batch: + if explore: + action_dist_class = sa_rl_module.get_exploration_action_dist_cls() + else: + action_dist_class = sa_rl_module.get_inference_action_dist_cls() + action_dist = action_dist_class.from_logits( + batch[Columns.ACTION_DIST_INPUTS], + ) + if not explore: + action_dist = action_dist.to_deterministic() + + # Sample actions from the distribution. + actions = action_dist.sample() + batch[Columns.ACTIONS] = actions + + # For convenience and if possible, compute action logp from distribution + # and add to output. + if Columns.ACTION_LOGP not in batch: + batch[Columns.ACTION_LOGP] = action_dist.logp(actions) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/listify_data_for_vector_env.py b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/listify_data_for_vector_env.py new file mode 100644 index 0000000000000000000000000000000000000000..c52489a8b1e5531f46287a554ab551f39e5e891a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/listify_data_for_vector_env.py @@ -0,0 +1,89 @@ +from typing import Any, Dict, List, Optional + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.utils.annotations import override +from ray.rllib.utils.spaces.space_utils import batch as batch_fn +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class ListifyDataForVectorEnv(ConnectorV2): + """Performs conversion from ConnectorV2-style format to env/episode insertion. + + Note: This is one of the default module-to-env ConnectorV2 pieces that + are added automatically by RLlib into every module-to-env connector pipeline, + unless `config.add_default_connectors_to_module_to_env_pipeline` is set to + False. + + The default module-to-env connector pipeline is: + [ + GetActions, + TensorToNumpy, + UnBatchToIndividualItems, + ModuleToAgentUnmapping, # only in multi-agent setups! + RemoveSingleTsTimeRankFromBatch, + + [0 or more user defined ConnectorV2 pieces], + + NormalizeAndClipActions, + ListifyDataForVectorEnv, + ] + + Single agent case: + Convert from: + [col] -> [(episode_id,)] -> [list of items]. + To: + [col] -> [list of items]. + + Multi-agent case: + Convert from: + [col] -> [(episode_id, agent_id, module_id)] -> list of items. + To: + [col] -> [list of multi-agent dicts]. + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + for column, column_data in batch.copy().items(): + # Multi-agent case: Create lists of multi-agent dicts under each column. + if isinstance(episodes[0], MultiAgentEpisode): + # For each environment/episode one entry. + new_column_data = [{} for _ in episodes] + episode_map_structure = shared_data.get("old_episode_ids", {}) + for key, value in batch[column].items(): + assert len(value) == 1 + eps_id, agent_id, module_id = key + # Get the episode index that corresponds to the ID in the batch. + # Note, the order in the `EnvRunner`s episode list needs to be + # kept for each batch column. + eps_id = episode_map_structure.get(eps_id, eps_id) + eps_index = next( + (i for i, eps in enumerate(episodes) if eps.id_ == eps_id), 0 + ) + + new_column_data[eps_index][agent_id] = value[0] + batch[column] = new_column_data + # Single-agent case: Create simple lists under each column. + else: + batch[column] = [ + d for key in batch[column].keys() for d in batch[column][key] + ] + # Batch actions for (single-agent) gym.vector.Env. + # All other columns, leave listify'ed. + if column in [Columns.ACTIONS_FOR_ENV, Columns.ACTIONS]: + batch[column] = batch_fn(batch[column]) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/module_to_env_pipeline.py b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/module_to_env_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..63d5f00bfa4cda7a59e1cd22b3476e2091241090 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/module_to_env_pipeline.py @@ -0,0 +1,7 @@ +from ray.rllib.connectors.connector_pipeline_v2 import ConnectorPipelineV2 +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class ModuleToEnvPipeline(ConnectorPipelineV2): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/normalize_and_clip_actions.py b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/normalize_and_clip_actions.py new file mode 100644 index 0000000000000000000000000000000000000000..1698b1d0fe1e1aa5a0a3b4bfd32e65b6efc3b60b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/normalize_and_clip_actions.py @@ -0,0 +1,146 @@ +import copy +from typing import Any, Dict, List, Optional + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.spaces.space_utils import ( + clip_action, + get_base_struct_from_space, + unsquash_action, +) +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class NormalizeAndClipActions(ConnectorV2): + """Normalizes or clips actions in the input data (coming from the RLModule). + + Note: This is one of the default module-to-env ConnectorV2 pieces that + are added automatically by RLlib into every module-to-env connector pipeline, + unless `config.add_default_connectors_to_module_to_env_pipeline` is set to + False. + + The default module-to-env connector pipeline is: + [ + GetActions, + TensorToNumpy, + UnBatchToIndividualItems, + ModuleToAgentUnmapping, # only in multi-agent setups! + RemoveSingleTsTimeRankFromBatch, + + [0 or more user defined ConnectorV2 pieces], + + NormalizeAndClipActions, + ListifyDataForVectorEnv, + ] + + This ConnectorV2: + - Deep copies the Columns.ACTIONS in the incoming `data` into a new column: + Columns.ACTIONS_FOR_ENV. + - Loops through the Columns.ACTIONS in the incoming `data` and normalizes or clips + these depending on the c'tor settings in `config.normalize_actions` and + `config.clip_actions`. + - Only applies to envs with Box action spaces. + + Normalizing is the process of mapping NN-outputs (which are usually small + numbers, e.g. between -1.0 and 1.0) to the bounds defined by the action-space. + Normalizing helps the NN to learn faster in environments with large ranges between + `low` and `high` bounds or skewed action bounds (e.g. Box(-3000.0, 1.0, ...)). + + Clipping clips the actions computed by the NN (and sampled from a distribution) + between the bounds defined by the action-space. Note that clipping is only performed + if `normalize_actions` is False. + """ + + @override(ConnectorV2) + def recompute_output_action_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + self._action_space_struct = get_base_struct_from_space(input_action_space) + return input_action_space + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + normalize_actions: bool, + clip_actions: bool, + **kwargs, + ): + """Initializes a DefaultModuleToEnv (connector piece) instance. + + Args: + normalize_actions: If True, actions coming from the RLModule's distribution + (or are directly computed by the RLModule w/o sampling) will + be assumed 0.0 centered with a small stddev (only affecting Box + components) and thus be unsquashed (and clipped, just in case) to the + bounds of the env's action space. For example, if the action space of + the environment is `Box(-2.0, -0.5, (1,))`, the model outputs + mean and stddev as 0.1 and exp(0.2), and we sample an action of 0.9 + from the resulting distribution, then this 0.9 will be unsquashed into + the [-2.0 -0.5] interval. If - after unsquashing - the action still + breaches the action space, it will simply be clipped. + clip_actions: If True, actions coming from the RLModule's distribution + (or are directly computed by the RLModule w/o sampling) will be clipped + such that they fit into the env's action space's bounds. + For example, if the action space of the environment is + `Box(-0.5, 0.5, (1,))`, the model outputs + mean and stddev as 0.1 and exp(0.2), and we sample an action of 0.9 + from the resulting distribution, then this 0.9 will be clipped to 0.5 + to fit into the [-0.5 0.5] interval. + """ + self._action_space_struct = None + + super().__init__(input_observation_space, input_action_space, **kwargs) + + self.normalize_actions = normalize_actions + self.clip_actions = clip_actions + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + """Based on settings, will normalize (unsquash) and/or clip computed actions. + + This is such that the final actions (to be sent to the env) match the + environment's action space and thus don't lead to an error. + """ + + def _unsquash_or_clip(action_for_env, env_id, agent_id, module_id): + if agent_id is not None: + struct = self._action_space_struct[agent_id] + else: + struct = self._action_space_struct + + if self.normalize_actions: + return unsquash_action(action_for_env, struct) + else: + return clip_action(action_for_env, struct) + + # Normalize or clip this new actions_for_env column, leaving the originally + # computed/sampled actions intact. + if self.normalize_actions or self.clip_actions: + # Copy actions into separate column, just to go to the env. + batch[Columns.ACTIONS_FOR_ENV] = copy.deepcopy(batch[Columns.ACTIONS]) + self.foreach_batch_item_change_in_place( + batch=batch, + column=Columns.ACTIONS_FOR_ENV, + func=_unsquash_or_clip, + ) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/remove_single_ts_time_rank_from_batch.py b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/remove_single_ts_time_rank_from_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..3e3b200ebc071b38825004fab89002e9693a5a98 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/remove_single_ts_time_rank_from_batch.py @@ -0,0 +1,70 @@ +from typing import Any, Dict, List, Optional + +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class RemoveSingleTsTimeRankFromBatch(ConnectorV2): + """ + Note: This is one of the default module-to-env ConnectorV2 pieces that + are added automatically by RLlib into every module-to-env connector pipeline, + unless `config.add_default_connectors_to_module_to_env_pipeline` is set to + False. + + The default module-to-env connector pipeline is: + [ + GetActions, + TensorToNumpy, + UnBatchToIndividualItems, + ModuleToAgentUnmapping, # only in multi-agent setups! + RemoveSingleTsTimeRankFromBatch, + + [0 or more user defined ConnectorV2 pieces], + + NormalizeAndClipActions, + ListifyDataForVectorEnv, + ] + + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Optional[Dict[str, Any]], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # If single ts time-rank had not been added, early out. + if shared_data is None or not shared_data.get("_added_single_ts_time_rank"): + return batch + + def _remove_single_ts(item, eps_id, aid, mid): + # Only remove time-rank for modules that are statefule (only for those, a + # timerank has been added). + if mid is None or rl_module[mid].is_stateful(): + return tree.map_structure(lambda s: np.squeeze(s, axis=0), item) + return item + + for column, column_data in batch.copy().items(): + # Skip state_out (doesn't have a time rank). + if column == Columns.STATE_OUT: + continue + self.foreach_batch_item_change_in_place( + batch, + column=column, + func=_remove_single_ts, + ) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/unbatch_to_individual_items.py b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/unbatch_to_individual_items.py new file mode 100644 index 0000000000000000000000000000000000000000..b8b84a9f92f20117451efaf7fd63325a7c7cca58 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/module_to_env/unbatch_to_individual_items.py @@ -0,0 +1,110 @@ +from collections import defaultdict +from typing import Any, Dict, List, Optional + +import tree # pip install dm_tree + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.spaces.space_utils import unbatch as unbatch_fn +from ray.rllib.utils.typing import EpisodeType +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class UnBatchToIndividualItems(ConnectorV2): + """Unbatches the given `data` back into the individual-batch-items format. + + Note: This is one of the default module-to-env ConnectorV2 pieces that + are added automatically by RLlib into every module-to-env connector pipeline, + unless `config.add_default_connectors_to_module_to_env_pipeline` is set to + False. + + The default module-to-env connector pipeline is: + [ + GetActions, + TensorToNumpy, + UnBatchToIndividualItems, + ModuleToAgentUnmapping, # only in multi-agent setups! + RemoveSingleTsTimeRankFromBatch, + + [0 or more user defined ConnectorV2 pieces], + + NormalizeAndClipActions, + ListifyDataForVectorEnv, + ] + """ + + @override(ConnectorV2) + def __call__( + self, + *, + rl_module: RLModule, + batch: Dict[str, Any], + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + memorized_map_structure = shared_data.get("memorized_map_structure") + episode_map_structure = shared_data.get("vector_env_episodes_map", {}) + # Simple case (no structure stored): Just unbatch. + if memorized_map_structure is None: + return tree.map_structure(lambda s: unbatch_fn(s), batch) + # Single agent case: Memorized structure is a list, whose indices map to + # eps_id values. + elif isinstance(memorized_map_structure, list): + for column, column_data in batch.copy().items(): + column_data = unbatch_fn(column_data) + new_column_data = defaultdict(list) + for i, eps_id in enumerate(memorized_map_structure): + # Keys are always tuples to resemble multi-agent keys, which + # have the structure (eps_id, agent_id, module_id). + key = (eps_id,) + new_column_data[key].append(column_data[i]) + batch[column] = dict(new_column_data) + # Multi-agent case: Memorized structure is dict mapping module_ids to lists of + # (eps_id, agent_id)-tuples, such that the original individual-items-based form + # can be constructed. + else: + for module_id, module_data in batch.copy().items(): + if module_id not in memorized_map_structure: + raise KeyError( + f"ModuleID={module_id} not found in `memorized_map_structure`!" + ) + for column, column_data in module_data.items(): + column_data = unbatch_fn(column_data) + new_column_data = defaultdict(list) + for i, (eps_id, agent_id) in enumerate( + memorized_map_structure[module_id] + ): + key = (eps_id, agent_id, module_id) + + # Check, if an agent episode is already done. For this we need + # to get the corresponding episode in the `EnvRunner`s list of + # episodes. + eps_id = episode_map_structure.get(eps_id, eps_id) + episode = next( + (eps for eps in episodes if eps.id_ == eps_id), None + ) + + if episode is None: + raise ValueError( + f"No episode found that matches the ID={eps_id}. Check " + "shared_data['vector_env_episodes_map'] for a missing ", + "mapping.", + ) + # If an episode has not just started and the agent's episode + # is done do not return data. + # This should not be `True` for new `MultiAgentEpisode`s. + if ( + episode.agent_episodes + and episode.agent_episodes[agent_id].is_done + and not episode.is_done + ): + continue + + new_column_data[key].append(column_data[i]) + module_data[column] = dict(new_column_data) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/registry.py b/lib/python3.12/site-packages/ray/rllib/connectors/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..2f51aa6f446a14ac40f7dee7c0f5175ccc981b61 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/registry.py @@ -0,0 +1,45 @@ +"""Registry of connector names for global access.""" +from typing import Any + +from ray.rllib.connectors.connector import Connector, ConnectorContext +from ray.rllib.utils.annotations import OldAPIStack + +ALL_CONNECTORS = dict() + + +@OldAPIStack +def register_connector(name: str, cls: Connector): + """Register a connector for use with RLlib. + + Args: + name: Name to register. + cls: Callable that creates an env. + """ + if name in ALL_CONNECTORS: + return + + if not issubclass(cls, Connector): + raise TypeError("Can only register Connector type.", cls) + + # Record it in local registry in case we need to register everything + # again in the global registry, for example in the event of cluster + # restarts. + ALL_CONNECTORS[name] = cls + + +@OldAPIStack +def get_connector(name: str, ctx: ConnectorContext, params: Any = None) -> Connector: + # TODO(jungong) : switch the order of parameters man!! + """Get a connector by its name and serialized config. + + Args: + name: name of the connector. + ctx: Connector context. + params: serialized parameters of the connector. + + Returns: + Constructed connector. + """ + if name not in ALL_CONNECTORS: + raise NameError("connector not found.", name) + return ALL_CONNECTORS[name].from_state(ctx, params) diff --git a/lib/python3.12/site-packages/ray/rllib/connectors/util.py b/lib/python3.12/site-packages/ray/rllib/connectors/util.py new file mode 100644 index 0000000000000000000000000000000000000000..e7824a075d3834c5eb6f780eb02b20dcbb46fc04 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/connectors/util.py @@ -0,0 +1,170 @@ +import logging +from typing import TYPE_CHECKING, Any, Tuple + +from ray.rllib.connectors.action.clip import ClipActionsConnector +from ray.rllib.connectors.action.immutable import ImmutableActionsConnector +from ray.rllib.connectors.action.lambdas import ConvertToNumpyConnector +from ray.rllib.connectors.action.normalize import NormalizeActionsConnector +from ray.rllib.connectors.action.pipeline import ActionConnectorPipeline +from ray.rllib.connectors.agent.clip_reward import ClipRewardAgentConnector +from ray.rllib.connectors.agent.mean_std_filter import ( + ConcurrentMeanStdObservationFilterAgentConnector, + MeanStdObservationFilterAgentConnector, +) +from ray.rllib.connectors.agent.obs_preproc import ObsPreprocessorConnector +from ray.rllib.connectors.agent.pipeline import AgentConnectorPipeline +from ray.rllib.connectors.agent.state_buffer import StateBufferConnector +from ray.rllib.connectors.agent.synced_filter import SyncedFilterAgentConnector +from ray.rllib.connectors.agent.view_requirement import ViewRequirementAgentConnector +from ray.rllib.connectors.connector import Connector, ConnectorContext +from ray.rllib.connectors.registry import get_connector +from ray.rllib.utils.annotations import OldAPIStack + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.policy.policy import Policy + +logger = logging.getLogger(__name__) + + +def __preprocessing_enabled(config: "AlgorithmConfig"): + if config._disable_preprocessor_api: + return False + # Same conditions as in RolloutWorker.__init__. + if config.is_atari and config.preprocessor_pref == "deepmind": + return False + if config.preprocessor_pref is None: + return False + return True + + +def __clip_rewards(config: "AlgorithmConfig"): + # Same logic as in RolloutWorker.__init__. + # We always clip rewards for Atari games. + return config.clip_rewards or config.is_atari + + +@OldAPIStack +def get_agent_connectors_from_config( + ctx: ConnectorContext, + config: "AlgorithmConfig", +) -> AgentConnectorPipeline: + connectors = [] + + clip_rewards = __clip_rewards(config) + if clip_rewards is True: + connectors.append(ClipRewardAgentConnector(ctx, sign=True)) + elif type(clip_rewards) is float: + connectors.append(ClipRewardAgentConnector(ctx, limit=abs(clip_rewards))) + + if __preprocessing_enabled(config): + connectors.append(ObsPreprocessorConnector(ctx)) + + # Filters should be after observation preprocessing + filter_connector = get_synced_filter_connector( + ctx, + ) + # Configuration option "NoFilter" results in `filter_connector==None`. + if filter_connector: + connectors.append(filter_connector) + + connectors.extend( + [ + StateBufferConnector(ctx), + ViewRequirementAgentConnector(ctx), + ] + ) + + return AgentConnectorPipeline(ctx, connectors) + + +@OldAPIStack +def get_action_connectors_from_config( + ctx: ConnectorContext, + config: "AlgorithmConfig", +) -> ActionConnectorPipeline: + """Default list of action connectors to use for a new policy. + + Args: + ctx: context used to create connectors. + config: The AlgorithmConfig object. + """ + connectors = [ConvertToNumpyConnector(ctx)] + if config.get("normalize_actions", False): + connectors.append(NormalizeActionsConnector(ctx)) + if config.get("clip_actions", False): + connectors.append(ClipActionsConnector(ctx)) + connectors.append(ImmutableActionsConnector(ctx)) + return ActionConnectorPipeline(ctx, connectors) + + +@OldAPIStack +def create_connectors_for_policy(policy: "Policy", config: "AlgorithmConfig"): + """Util to create agent and action connectors for a Policy. + + Args: + policy: Policy instance. + config: Algorithm config dict. + """ + ctx: ConnectorContext = ConnectorContext.from_policy(policy) + + assert ( + policy.agent_connectors is None and policy.action_connectors is None + ), "Can not create connectors for a policy that already has connectors." + + policy.agent_connectors = get_agent_connectors_from_config(ctx, config) + policy.action_connectors = get_action_connectors_from_config(ctx, config) + + logger.info("Using connectors:") + logger.info(policy.agent_connectors.__str__(indentation=4)) + logger.info(policy.action_connectors.__str__(indentation=4)) + + +@OldAPIStack +def restore_connectors_for_policy( + policy: "Policy", connector_config: Tuple[str, Tuple[Any]] +) -> Connector: + """Util to create connector for a Policy based on serialized config. + + Args: + policy: Policy instance. + connector_config: Serialized connector config. + """ + ctx: ConnectorContext = ConnectorContext.from_policy(policy) + name, params = connector_config + return get_connector(name, ctx, params) + + +# We need this filter selection mechanism temporarily to remain compatible to old API +@OldAPIStack +def get_synced_filter_connector(ctx: ConnectorContext): + filter_specifier = ctx.config.get("observation_filter") + if filter_specifier == "MeanStdFilter": + return MeanStdObservationFilterAgentConnector(ctx, clip=None) + elif filter_specifier == "ConcurrentMeanStdFilter": + return ConcurrentMeanStdObservationFilterAgentConnector(ctx, clip=None) + elif filter_specifier == "NoFilter": + return None + else: + raise Exception("Unknown observation_filter: " + str(filter_specifier)) + + +@OldAPIStack +def maybe_get_filters_for_syncing(rollout_worker, policy_id): + # As long as the historic filter synchronization mechanism is in + # place, we need to put filters into self.filters so that they get + # synchronized + policy = rollout_worker.policy_map[policy_id] + if not policy.agent_connectors: + return + + filter_connectors = policy.agent_connectors[SyncedFilterAgentConnector] + # There can only be one filter at a time + if not filter_connectors: + return + + assert len(filter_connectors) == 1, ( + "ConnectorPipeline has multiple connectors of type " + "SyncedFilterAgentConnector but can only have one." + ) + rollout_worker.filters[policy_id] = filter_connectors[0].filter diff --git a/lib/python3.12/site-packages/ray/rllib/env/__init__.py b/lib/python3.12/site-packages/ray/rllib/env/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..247096892b304226eca06326f6dabe02862910be --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/__init__.py @@ -0,0 +1,36 @@ +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.env.env_context import EnvContext +from ray.rllib.env.external_env import ExternalEnv +from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.policy_client import PolicyClient +from ray.rllib.env.remote_base_env import RemoteBaseEnv +from ray.rllib.env.vector_env import VectorEnv +from ray.rllib.env.wrappers.dm_control_wrapper import DMCEnv +from ray.rllib.env.wrappers.dm_env_wrapper import DMEnv +from ray.rllib.env.wrappers.group_agents_wrapper import GroupAgentsWrapper +from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv, PettingZooEnv +from ray.rllib.env.wrappers.unity3d_env import Unity3DEnv + +INPUT_ENV_SPACES = "__env__" +INPUT_ENV_SINGLE_SPACES = "__env_single__" + + +__all__ = [ + "BaseEnv", + "DMEnv", + "DMCEnv", + "EnvContext", + "ExternalEnv", + "ExternalMultiAgentEnv", + "GroupAgentsWrapper", + "MultiAgentEnv", + "PettingZooEnv", + "ParallelPettingZooEnv", + "PolicyClient", + "RemoteBaseEnv", + "Unity3DEnv", + "VectorEnv", + "INPUT_ENV_SPACES", + "INPUT_ENV_SINGLE_SPACES", +] diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97bc64cca03cb8a1e911955e670f3657d2518995 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/base_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/base_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63c996cbec0d099060accfd1e5931d9ce458569e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/base_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_context.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93e6ff3af6cef4447735a096b1cc0be52a2d17e0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_errors.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_errors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8dcb48e247d4aaf7af7e5d425ee6e67b286be22 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_errors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa5450e0cfedf4ec84a283ea7d4bce18afa0d0c8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_runner_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_runner_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..298b8a503dadcaeac547454cfc201db260b51119 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/env_runner_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/external_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/external_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..302424c8a4e8b5871a71ffc5a1417663341f6faf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/external_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/external_multi_agent_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/external_multi_agent_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f0e07ba02c0292bd3d57970d3f188319c23e26c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/external_multi_agent_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/multi_agent_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/multi_agent_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..995f0cd32d6a9b85ce4fa9775924e0ebd0b7b2de Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/multi_agent_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/multi_agent_env_runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/multi_agent_env_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6ff6eeb10441026f344419c24e155109641a0c6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/multi_agent_env_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/policy_client.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/policy_client.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e48be1b1729e2cdf0f5300d40452baed4f0ef1c9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/policy_client.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/remote_base_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/remote_base_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfee40d88102ef2951391fe3e30c8895ed930058 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/remote_base_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/single_agent_env_runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/single_agent_env_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d011645f614bbb5002278fc390d5ae15ee45e12a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/single_agent_env_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/single_agent_episode.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/single_agent_episode.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..140c202fddbbd27b9f7daeb5b45376cddf679811 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/single_agent_episode.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/tcp_client_inference_env_runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/tcp_client_inference_env_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf00d4475c544e3c3dbd0c13402034971d8905ac Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/tcp_client_inference_env_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/__pycache__/vector_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/vector_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adfe38136c3fd5d8e70549f139a9efae44ed6fb9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/__pycache__/vector_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/base_env.py b/lib/python3.12/site-packages/ray/rllib/env/base_env.py new file mode 100644 index 0000000000000000000000000000000000000000..1467472eead85216b55a4c819e1598272b0c7a54 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/base_env.py @@ -0,0 +1,429 @@ +import logging +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union + +import gymnasium as gym + +import ray +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import AgentID, EnvID, EnvType, MultiEnvDict + +if TYPE_CHECKING: + from ray.rllib.evaluation.rollout_worker import RolloutWorker + +ASYNC_RESET_RETURN = "async_reset_return" + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class BaseEnv: + """The lowest-level env interface used by RLlib for sampling. + + BaseEnv models multiple agents executing asynchronously in multiple + vectorized sub-environments. A call to `poll()` returns observations from + ready agents keyed by their sub-environment ID and agent IDs, and + actions for those agents can be sent back via `send_actions()`. + + All other RLlib supported env types can be converted to BaseEnv. + RLlib handles these conversions internally in RolloutWorker, for example: + + gym.Env => rllib.VectorEnv => rllib.BaseEnv + rllib.MultiAgentEnv (is-a gym.Env) => rllib.VectorEnv => rllib.BaseEnv + rllib.ExternalEnv => rllib.BaseEnv + + .. testcode:: + :skipif: True + + MyBaseEnv = ... + env = MyBaseEnv() + obs, rewards, terminateds, truncateds, infos, off_policy_actions = ( + env.poll() + ) + print(obs) + + env.send_actions({ + "env_0": { + "car_0": 0, + "car_1": 1, + }, ... + }) + obs, rewards, terminateds, truncateds, infos, off_policy_actions = ( + env.poll() + ) + print(obs) + + print(terminateds) + + .. testoutput:: + + { + "env_0": { + "car_0": [2.4, 1.6], + "car_1": [3.4, -3.2], + }, + "env_1": { + "car_0": [8.0, 4.1], + }, + "env_2": { + "car_0": [2.3, 3.3], + "car_1": [1.4, -0.2], + "car_3": [1.2, 0.1], + }, + } + { + "env_0": { + "car_0": [4.1, 1.7], + "car_1": [3.2, -4.2], + }, ... + } + { + "env_0": { + "__all__": False, + "car_0": False, + "car_1": True, + }, ... + } + + """ + + def to_base_env( + self, + make_env: Optional[Callable[[int], EnvType]] = None, + num_envs: int = 1, + remote_envs: bool = False, + remote_env_batch_wait_ms: int = 0, + restart_failed_sub_environments: bool = False, + ) -> "BaseEnv": + """Converts an RLlib-supported env into a BaseEnv object. + + Supported types for the `env` arg are gym.Env, BaseEnv, + VectorEnv, MultiAgentEnv, ExternalEnv, or ExternalMultiAgentEnv. + + The resulting BaseEnv is always vectorized (contains n + sub-environments) to support batched forward passes, where n may also + be 1. BaseEnv also supports async execution via the `poll` and + `send_actions` methods and thus supports external simulators. + + TODO: Support gym3 environments, which are already vectorized. + + Args: + env: An already existing environment of any supported env type + to convert/wrap into a BaseEnv. Supported types are gym.Env, + BaseEnv, VectorEnv, MultiAgentEnv, ExternalEnv, and + ExternalMultiAgentEnv. + make_env: A callable taking an int as input (which indicates the + number of individual sub-environments within the final + vectorized BaseEnv) and returning one individual + sub-environment. + num_envs: The number of sub-environments to create in the + resulting (vectorized) BaseEnv. The already existing `env` + will be one of the `num_envs`. + remote_envs: Whether each sub-env should be a @ray.remote actor. + You can set this behavior in your config via the + `remote_worker_envs=True` option. + remote_env_batch_wait_ms: The wait time (in ms) to poll remote + sub-environments for, if applicable. Only used if + `remote_envs` is True. + policy_config: Optional policy config dict. + + Returns: + The resulting BaseEnv object. + """ + return self + + def poll( + self, + ) -> Tuple[ + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + ]: + """Returns observations from ready agents. + + All return values are two-level dicts mapping from EnvID to dicts + mapping from AgentIDs to (observation/reward/etc..) values. + The number of agents and sub-environments may vary over time. + + Returns: + Tuple consisting of: + New observations for each ready agent. + Reward values for each ready agent. If the episode is just started, + the value will be None. + Terminated values for each ready agent. The special key "__all__" is used to + indicate episode termination. + Truncated values for each ready agent. The special key "__all__" + is used to indicate episode truncation. + Info values for each ready agent. + Agents may take off-policy actions, in which case, there will be an entry + in this dict that contains the taken action. There is no need to + `send_actions()` for agents that have already chosen off-policy actions. + """ + raise NotImplementedError + + def send_actions(self, action_dict: MultiEnvDict) -> None: + """Called to send actions back to running agents in this env. + + Actions should be sent for each ready agent that returned observations + in the previous poll() call. + + Args: + action_dict: Actions values keyed by env_id and agent_id. + """ + raise NotImplementedError + + def try_reset( + self, + env_id: Optional[EnvID] = None, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Tuple[Optional[MultiEnvDict], Optional[MultiEnvDict]]: + """Attempt to reset the sub-env with the given id or all sub-envs. + + If the environment does not support synchronous reset, a tuple of + (ASYNC_RESET_REQUEST, ASYNC_RESET_REQUEST) can be returned here. + + Note: A MultiAgentDict is returned when using the deprecated wrapper + classes such as `ray.rllib.env.base_env._MultiAgentEnvToBaseEnv`, + however for consistency with the poll() method, a `MultiEnvDict` is + returned from the new wrapper classes, such as + `ray.rllib.env.multi_agent_env.MultiAgentEnvWrapper`. + + Args: + env_id: The sub-environment's ID if applicable. If None, reset + the entire Env (i.e. all sub-environments). + seed: The seed to be passed to the sub-environment(s) when + resetting it. If None, will not reset any existing PRNG. If you pass an + integer, the PRNG will be reset even if it already exists. + options: An options dict to be passed to the sub-environment(s) when + resetting it. + + Returns: + A tuple consisting of a) the reset (multi-env/multi-agent) observation + dict and b) the reset (multi-env/multi-agent) infos dict. Returns the + (ASYNC_RESET_REQUEST, ASYNC_RESET_REQUEST) tuple, if not supported. + """ + return None, None + + def try_restart(self, env_id: Optional[EnvID] = None) -> None: + """Attempt to restart the sub-env with the given id or all sub-envs. + + This could result in the sub-env being completely removed (gc'd) and recreated. + + Args: + env_id: The sub-environment's ID, if applicable. If None, restart + the entire Env (i.e. all sub-environments). + """ + return None + + def get_sub_environments(self, as_dict: bool = False) -> Union[List[EnvType], dict]: + """Return a reference to the underlying sub environments, if any. + + Args: + as_dict: If True, return a dict mapping from env_id to env. + + Returns: + List or dictionary of the underlying sub environments or [] / {}. + """ + if as_dict: + return {} + return [] + + def get_agent_ids(self) -> Set[AgentID]: + """Return the agent ids for the sub_environment. + + Returns: + All agent ids for each the environment. + """ + return {} + + def try_render(self, env_id: Optional[EnvID] = None) -> None: + """Tries to render the sub-environment with the given id or all. + + Args: + env_id: The sub-environment's ID, if applicable. + If None, renders the entire Env (i.e. all sub-environments). + """ + + # By default, do nothing. + pass + + def stop(self) -> None: + """Releases all resources used.""" + + # Try calling `close` on all sub-environments. + for env in self.get_sub_environments(): + if hasattr(env, "close"): + env.close() + + @property + def observation_space(self) -> gym.Space: + """Returns the observation space for each agent. + + Note: samples from the observation space need to be preprocessed into a + `MultiEnvDict` before being used by a policy. + + Returns: + The observation space for each environment. + """ + raise NotImplementedError + + @property + def action_space(self) -> gym.Space: + """Returns the action space for each agent. + + Note: samples from the action space need to be preprocessed into a + `MultiEnvDict` before being passed to `send_actions`. + + Returns: + The observation space for each environment. + """ + raise NotImplementedError + + def last( + self, + ) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, MultiEnvDict, MultiEnvDict]: + """Returns the last observations, rewards, done- truncated flags and infos ... + + that were returned by the environment. + + Returns: + The last observations, rewards, done- and truncated flags, and infos + for each sub-environment. + """ + logger.warning("last has not been implemented for this environment.") + return {}, {}, {}, {}, {} + + +# Fixed agent identifier when there is only the single agent in the env +_DUMMY_AGENT_ID = "agent0" + + +@OldAPIStack +def with_dummy_agent_id( + env_id_to_values: Dict[EnvID, Any], dummy_id: "AgentID" = _DUMMY_AGENT_ID +) -> MultiEnvDict: + ret = {} + for (env_id, value) in env_id_to_values.items(): + # If the value (e.g. the observation) is an Exception, publish this error + # under the env ID so the caller of `poll()` knows that the entire episode + # (sub-environment) has crashed. + ret[env_id] = value if isinstance(value, Exception) else {dummy_id: value} + return ret + + +@OldAPIStack +def convert_to_base_env( + env: EnvType, + make_env: Callable[[int], EnvType] = None, + num_envs: int = 1, + remote_envs: bool = False, + remote_env_batch_wait_ms: int = 0, + worker: Optional["RolloutWorker"] = None, + restart_failed_sub_environments: bool = False, +) -> "BaseEnv": + """Converts an RLlib-supported env into a BaseEnv object. + + Supported types for the `env` arg are gym.Env, BaseEnv, + VectorEnv, MultiAgentEnv, ExternalEnv, or ExternalMultiAgentEnv. + + The resulting BaseEnv is always vectorized (contains n + sub-environments) to support batched forward passes, where n may also + be 1. BaseEnv also supports async execution via the `poll` and + `send_actions` methods and thus supports external simulators. + + TODO: Support gym3 environments, which are already vectorized. + + Args: + env: An already existing environment of any supported env type + to convert/wrap into a BaseEnv. Supported types are gym.Env, + BaseEnv, VectorEnv, MultiAgentEnv, ExternalEnv, and + ExternalMultiAgentEnv. + make_env: A callable taking an int as input (which indicates the + number of individual sub-environments within the final + vectorized BaseEnv) and returning one individual + sub-environment. + num_envs: The number of sub-environments to create in the + resulting (vectorized) BaseEnv. The already existing `env` + will be one of the `num_envs`. + remote_envs: Whether each sub-env should be a @ray.remote actor. + You can set this behavior in your config via the + `remote_worker_envs=True` option. + remote_env_batch_wait_ms: The wait time (in ms) to poll remote + sub-environments for, if applicable. Only used if + `remote_envs` is True. + worker: An optional RolloutWorker that owns the env. This is only + used if `remote_worker_envs` is True in your config and the + `on_sub_environment_created` custom callback needs to be called + on each created actor. + restart_failed_sub_environments: If True and any sub-environment (within + a vectorized env) throws any error during env stepping, the + Sampler will try to restart the faulty sub-environment. This is done + without disturbing the other (still intact) sub-environment and without + the RolloutWorker crashing. + + Returns: + The resulting BaseEnv object. + """ + + from ray.rllib.env.external_env import ExternalEnv + from ray.rllib.env.multi_agent_env import MultiAgentEnv + from ray.rllib.env.remote_base_env import RemoteBaseEnv + from ray.rllib.env.vector_env import VectorEnv, VectorEnvWrapper + + if remote_envs and num_envs == 1: + raise ValueError( + "Remote envs only make sense to use if num_envs > 1 " + "(i.e. environment vectorization is enabled)." + ) + + # Given `env` has a `to_base_env` method -> Call that to convert to a BaseEnv type. + if isinstance(env, (BaseEnv, MultiAgentEnv, VectorEnv, ExternalEnv)): + return env.to_base_env( + make_env=make_env, + num_envs=num_envs, + remote_envs=remote_envs, + remote_env_batch_wait_ms=remote_env_batch_wait_ms, + restart_failed_sub_environments=restart_failed_sub_environments, + ) + # `env` is not a BaseEnv yet -> Need to convert/vectorize. + else: + # Sub-environments are ray.remote actors: + if remote_envs: + # Determine, whether the already existing sub-env (could + # be a ray.actor) is multi-agent or not. + multiagent = ( + ray.get(env._is_multi_agent.remote()) + if hasattr(env, "_is_multi_agent") + else False + ) + env = RemoteBaseEnv( + make_env, + num_envs, + multiagent=multiagent, + remote_env_batch_wait_ms=remote_env_batch_wait_ms, + existing_envs=[env], + worker=worker, + restart_failed_sub_environments=restart_failed_sub_environments, + ) + # Sub-environments are not ray.remote actors. + else: + # Convert gym.Env to VectorEnv ... + env = VectorEnv.vectorize_gym_envs( + make_env=make_env, + existing_envs=[env], + num_envs=num_envs, + action_space=env.action_space, + observation_space=env.observation_space, + restart_failed_sub_environments=restart_failed_sub_environments, + ) + # ... then the resulting VectorEnv to a BaseEnv. + env = VectorEnvWrapper(env) + + # Make sure conversion went well. + assert isinstance(env, BaseEnv), env + + return env diff --git a/lib/python3.12/site-packages/ray/rllib/env/env_context.py b/lib/python3.12/site-packages/ray/rllib/env/env_context.py new file mode 100644 index 0000000000000000000000000000000000000000..2dbdae17c6add842021b1b33a4e476c15bfcc86d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/env_context.py @@ -0,0 +1,128 @@ +import copy +from typing import Optional + +from ray.rllib.utils.typing import EnvConfigDict +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class EnvContext(dict): + """Wraps env configurations to include extra rllib metadata. + + These attributes can be used to parameterize environments per process. + For example, one might use `worker_index` to control which data file an + environment reads in on initialization. + + RLlib auto-sets these attributes when constructing registered envs. + """ + + def __init__( + self, + env_config: EnvConfigDict, + worker_index: int, + vector_index: int = 0, + remote: bool = False, + num_workers: Optional[int] = None, + recreated_worker: bool = False, + ): + """Initializes an EnvContext instance. + + Args: + env_config: The env's configuration defined under the + "env_config" key in the Algorithm's config. + worker_index: When there are multiple workers created, this + uniquely identifies the worker the env is created in. + 0 for local worker, >0 for remote workers. + vector_index: When there are multiple envs per worker, this + uniquely identifies the env index within the worker. + Starts from 0. + remote: Whether individual sub-environments (in a vectorized + env) should be @ray.remote actors or not. + num_workers: The total number of (remote) workers in the set. + 0 if only a local worker exists. + recreated_worker: Whether the worker that holds this env is a recreated one. + This means that it replaced a previous (failed) worker when + `restart_failed_env_runners=True` in the Algorithm's config. + """ + # Store the env_config in the (super) dict. + dict.__init__(self, env_config) + + # Set some metadata attributes. + self.worker_index = worker_index + self.vector_index = vector_index + self.remote = remote + self.num_workers = num_workers + self.recreated_worker = recreated_worker + + def copy_with_overrides( + self, + env_config: Optional[EnvConfigDict] = None, + worker_index: Optional[int] = None, + vector_index: Optional[int] = None, + remote: Optional[bool] = None, + num_workers: Optional[int] = None, + recreated_worker: Optional[bool] = None, + ) -> "EnvContext": + """Returns a copy of this EnvContext with some attributes overridden. + + Args: + env_config: Optional env config to use. None for not overriding + the one from the source (self). + worker_index: Optional worker index to use. None for not + overriding the one from the source (self). + vector_index: Optional vector index to use. None for not + overriding the one from the source (self). + remote: Optional remote setting to use. None for not overriding + the one from the source (self). + num_workers: Optional num_workers to use. None for not overriding + the one from the source (self). + recreated_worker: Optional flag, indicating, whether the worker that holds + the env is a recreated one. This means that it replaced a previous + (failed) worker when `restart_failed_env_runners=True` in the + Algorithm's config. + + Returns: + A new EnvContext object as a copy of self plus the provided + overrides. + """ + return EnvContext( + copy.deepcopy(env_config) if env_config is not None else self, + worker_index if worker_index is not None else self.worker_index, + vector_index if vector_index is not None else self.vector_index, + remote if remote is not None else self.remote, + num_workers if num_workers is not None else self.num_workers, + recreated_worker if recreated_worker is not None else self.recreated_worker, + ) + + def set_defaults(self, defaults: dict) -> None: + """Sets missing keys of self to the values given in `defaults`. + + If `defaults` contains keys that already exist in self, don't override + the values with these defaults. + + Args: + defaults: The key/value pairs to add to self, but only for those + keys in `defaults` that don't exist yet in self. + + .. testcode:: + :skipif: True + + from ray.rllib.env.env_context import EnvContext + env_ctx = EnvContext({"a": 1, "b": 2}, worker_index=0) + env_ctx.set_defaults({"a": -42, "c": 3}) + print(env_ctx) + + .. testoutput:: + + {"a": 1, "b": 2, "c": 3} + """ + for key, value in defaults.items(): + if key not in self: + self[key] = value + + def __str__(self): + return ( + super().__str__()[:-1] + + f", worker={self.worker_index}/{self.num_workers}, " + f"vector_idx={self.vector_index}, remote={self.remote}" + "}" + ) diff --git a/lib/python3.12/site-packages/ray/rllib/env/env_errors.py b/lib/python3.12/site-packages/ray/rllib/env/env_errors.py new file mode 100644 index 0000000000000000000000000000000000000000..cb52892db1da3b70e46cf1a6cec2fa4488ab6cd8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/env_errors.py @@ -0,0 +1,18 @@ +"""Error classes for RLlib environment operations.""" + +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class StepFailedRecreateEnvError(Exception): + """An exception that signals that the environment step failed and the environment needs to be reset. + + This exception may be raised by the environment's `step` method. + It is then caught by the `EnvRunner` and the environment is reset. + This can be useful if your environment is unstable, regularely crashing in a certain way. + For example, if you connect to an external simulator that you have little control over. + You can detect such crashes in your step method and throw this error to not log the error. + Use this with caution, as it may lead to infinite loops of resetting the environment. + """ + + pass diff --git a/lib/python3.12/site-packages/ray/rllib/env/env_runner.py b/lib/python3.12/site-packages/ray/rllib/env/env_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..c407f02c42cb8234c016e302e4cd45b43f7d09ba --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/env_runner.py @@ -0,0 +1,296 @@ +import abc +import logging +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple + +import gymnasium as gym +import tree # pip install dm_tree + +import ray +from ray.rllib.core import COMPONENT_RL_MODULE +from ray.rllib.env.env_errors import StepFailedRecreateEnvError +from ray.rllib.utils.actor_manager import FaultAwareApply +from ray.rllib.utils.debug import update_global_seed_if_necessary +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import ENV_RESET_TIMER, ENV_STEP_TIMER +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.rllib.utils.typing import StateDict, TensorType +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.metrics import Counter + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +logger = logging.getLogger("ray.rllib") + +tf1, tf, _ = try_import_tf() + +ENV_RESET_FAILURE = "env_reset_failure" +ENV_STEP_FAILURE = "env_step_failure" +NUM_ENV_STEP_FAILURES_LIFETIME = "num_env_step_failures" + + +# TODO (sven): As soon as RolloutWorker is no longer supported, make this base class +# a Checkpointable. Currently, only some of its subclasses are Checkpointables. +@PublicAPI(stability="alpha") +class EnvRunner(FaultAwareApply, metaclass=abc.ABCMeta): + """Base class for distributed RL-style data collection from an environment. + + The EnvRunner API's core functionalities can be summarized as: + - Gets configured via passing a AlgorithmConfig object to the constructor. + Normally, subclasses of EnvRunner then construct their own environment (possibly + vectorized) copies and RLModules/Policies and use the latter to step through the + environment in order to collect training data. + - Clients of EnvRunner can use the `sample()` method to collect data for training + from the environment(s). + - EnvRunner offers parallelism via creating n remote Ray Actors based on this class. + Use `ray.remote([resources])(EnvRunner)` method to create the corresponding Ray + remote class. Then instantiate n Actors using the Ray `[ctor].remote(...)` syntax. + - EnvRunner clients can get information about the server/node on which the + individual Actors are running. + """ + + def __init__(self, *, config: "AlgorithmConfig", **kwargs): + """Initializes an EnvRunner instance. + + Args: + config: The AlgorithmConfig to use to setup this EnvRunner. + **kwargs: Forward compatibility kwargs. + """ + self.config: AlgorithmConfig = config.copy(copy_frozen=False) + + # Get the worker index on which this instance is running. + + # TODO (sven): We should make these c'tor named args. + self.worker_index: int = kwargs.get("worker_index") + self.num_workers: int = kwargs.get("num_workers", self.config.num_env_runners) + + self.env = None + # Create a MetricsLogger object for logging custom stats. + self.metrics: MetricsLogger = MetricsLogger() + + super().__init__() + + # This eager check is necessary for certain all-framework tests + # that use tf's eager_mode() context generator. + if ( + tf1 + and (self.config.framework_str == "tf2" or config.enable_tf1_exec_eagerly) + and not tf1.executing_eagerly() + ): + tf1.enable_eager_execution() + + # Determine actual seed for this particular worker based on worker index AND + # whether it's an eval worker. + self._seed: Optional[int] = None + if self.config.seed is not None: + self._seed = int( + self.config.seed + + (self.worker_index or 0) + # Eval workers get a +1M seed. + + (1e6 * self.config.in_evaluation) + ) + # Seed everything (random, numpy, torch, tf), if `seed` is provided. + update_global_seed_if_necessary( + framework=self.config.framework_str, + seed=self._seed, + ) + + # Ray metrics + self._metrics_num_try_env_step = Counter( + name="rllib_env_runner_num_try_env_step_counter", + description="Number of env.step() calls attempted in this Env Runner.", + tag_keys=("rllib",), + ) + self._metrics_num_try_env_step.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + self._metrics_num_env_steps_sampled = Counter( + name="rllib_env_runner_num_env_steps_sampled_counter", + description="Number of env steps sampled in this Env Runner.", + tag_keys=("rllib",), + ) + self._metrics_num_env_steps_sampled.set_default_tags( + {"rllib": self.__class__.__name__} + ) + + @abc.abstractmethod + def assert_healthy(self): + """Checks that self.__init__() has been completed properly. + + Useful in case an `EnvRunner` is run as @ray.remote (Actor) and the owner + would like to make sure the Ray Actor has been properly initialized. + + Raises: + AssertionError: If the EnvRunner Actor has NOT been properly initialized. + """ + + # TODO: Make this an abstract method that must be implemented. + def make_env(self): + """Creates the RL environment for this EnvRunner and assigns it to `self.env`. + + Note that users should be able to change the EnvRunner's config (e.g. change + `self.config.env_config`) and then call this method to create new environments + with the updated configuration. + It should also be called after a failure of an earlier env in order to clean up + the existing env (for example `close()` it), re-create a new one, and then + continue sampling with that new env. + """ + pass + + # TODO: Make this an abstract method that must be implemented. + def make_module(self): + """Creates the RLModule for this EnvRunner and assigns it to `self.module`. + + Note that users should be able to change the EnvRunner's config (e.g. change + `self.config.rl_module_spec`) and then call this method to create a new RLModule + with the updated configuration. + """ + pass + + @abc.abstractmethod + def sample(self, **kwargs) -> Any: + """Returns experiences (of any form) sampled from this EnvRunner. + + The exact nature and size of collected data are defined via the EnvRunner's + config and may be overridden by the given arguments. + + Args: + **kwargs: Forward compatibility kwargs. + + Returns: + The collected experience in any form. + """ + + # TODO (sven): Make this an abstract method that must be overridden. + def get_metrics(self) -> Any: + """Returns metrics (in any form) of the thus far collected, completed episodes. + + Returns: + Metrics of any form. + """ + pass + + @DeveloperAPI + def sample_get_state_and_metrics( + self, + ) -> Tuple[ray.ObjectRef, StateDict, StateDict]: + """Convenience method for fast, async algorithms. + + Use this in Algorithms that need to sample Episode lists as ray.ObjectRef, but + also require (in the same remote call) the metrics and the EnvRunner states, + except for the module weights. + """ + _episodes = self.sample() + # Get the EnvRunner's connector states. + _connector_states = self.get_state(not_components=COMPONENT_RL_MODULE) + _metrics = self.get_metrics() + # Return episode lists by reference so we don't have to send them to the + # main algo process, but to the Aggregator- or Learner actors directly. + return ray.put(_episodes), _connector_states, _metrics + + @abc.abstractmethod + def get_spaces(self) -> Dict[str, Tuple[gym.Space, gym.Space]]: + """Returns a dict mapping ModuleIDs to 2-tuples of obs- and action space.""" + + def stop(self) -> None: + """Releases all resources used by this EnvRunner. + + For example, when using a gym.Env in this EnvRunner, you should make sure + that its `close()` method is called. + """ + pass + + def __del__(self) -> None: + """If this Actor is deleted, clears all resources used by it.""" + pass + + def _try_env_reset( + self, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Tuple[Any, Any]: + """Tries resetting the env and - if an error occurs - handles it gracefully. + + Args: + seed: An optional seed (int) to be passed to the Env.reset() call. + options: An optional options-dict to be passed to the Env.reset() call. + + Returns: + The results of calling `Env.reset()`, which is a tuple of observations and + info dicts. + + Raises: + Exception: In case `config.restart_failed_sub_environments` is False and + `Env.reset()` resulted in an error. + """ + # Try to reset. + try: + with self.metrics.log_time(ENV_RESET_TIMER): + obs, infos = self.env.reset(seed=seed, options=options) + # Everything ok -> return. + return obs, infos + # Error. + except Exception as e: + # If user wants to simply restart the env -> recreate env and try again + # (calling this method recursively until success). + if self.config.restart_failed_sub_environments: + logger.exception( + "Resetting the env resulted in an error! The original error " + f"is: {e.args[0]}" + ) + # Recreate the env and simply try again. + self.make_env() + return self._try_env_reset(seed=seed, options=options) + else: + raise e + + def _try_env_step(self, actions): + """Tries stepping the env and - if an error occurs - handles it gracefully.""" + try: + with self.metrics.log_time(ENV_STEP_TIMER): + results = self.env.step(actions) + self._log_env_steps(metric=self._metrics_num_try_env_step, num_steps=1) + + return results + except Exception as e: + self.metrics.log_value(NUM_ENV_STEP_FAILURES_LIFETIME, 1, reduce="sum") + + # @OldAPIStack (config.restart_failed_sub_environments) + if self.config.restart_failed_sub_environments: + if not isinstance(e, StepFailedRecreateEnvError): + logger.exception( + "Stepping the env resulted in an error! The original error " + f"is: {e}" + ) + # Recreate the env. + self.make_env() + # And return that the stepping failed. The caller will then handle + # specific cleanup operations (for example discarding thus-far collected + # data and repeating the step attempt). + return ENV_STEP_FAILURE + else: + if isinstance(e, StepFailedRecreateEnvError): + raise ValueError( + "Environment raised StepFailedRecreateEnvError but config.restart_failed_sub_environments is False." + ) from e + raise e + + def _convert_to_tensor(self, struct) -> TensorType: + """Converts structs to a framework-specific tensor.""" + + if self.config.framework_str == "torch": + return convert_to_torch_tensor(struct) + else: + return tree.map_structure(tf.convert_to_tensor, struct) + + def _log_env_steps(self, metric: Counter, num_steps: int) -> None: + if num_steps > 0: + metric.inc(value=num_steps) + else: + logger.warning( + f"RLlib {self.__class__.__name__}: Skipping Prometheus logging for metric '{metric.info['name']}'. " + f"Received num_steps={num_steps}, but the number of steps must be greater than 0." + ) diff --git a/lib/python3.12/site-packages/ray/rllib/env/env_runner_group.py b/lib/python3.12/site-packages/ray/rllib/env/env_runner_group.py new file mode 100644 index 0000000000000000000000000000000000000000..83203907b09f6a6ec24ef36eda481df6b129a05d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/env_runner_group.py @@ -0,0 +1,1303 @@ +import importlib.util +import logging +import os +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + Dict, + List, + Optional, + Tuple, + Type, + TypeVar, + Union, +) + +import gymnasium as gym + +import ray +from ray._common.deprecation import ( + DEPRECATED_VALUE, + deprecation_warning, +) +from ray.actor import ActorHandle +from ray.exceptions import RayActorError +from ray.rllib.core import ( + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_LEARNER, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + COMPONENT_RL_MODULE, +) +from ray.rllib.core.learner import LearnerGroup +from ray.rllib.core.rl_module import validate_module_id +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.env.env_context import EnvContext +from ray.rllib.env.env_runner import EnvRunner +from ray.rllib.evaluation.rollout_worker import RolloutWorker +from ray.rllib.offline import get_dataset_and_shards +from ray.rllib.policy.policy import Policy, PolicyState +from ray.rllib.utils.actor_manager import FaultTolerantActorManager +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME, WEIGHTS_SEQ_NO +from ray.rllib.utils.typing import ( + AgentID, + EnvCreator, + EnvType, + EpisodeID, + PartialAlgorithmConfigDict, + PolicyID, + SampleBatchType, + TensorType, +) +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + +# Generic type var for foreach_* methods. +T = TypeVar("T") + + +@DeveloperAPI +class EnvRunnerGroup: + """Set of EnvRunners with n @ray.remote workers and zero or one local worker. + + Where: n >= 0. + """ + + def __init__( + self, + *, + env_creator: Optional[EnvCreator] = None, + validate_env: Optional[Callable[[EnvType], None]] = None, + default_policy_class: Optional[Type[Policy]] = None, + config: Optional["AlgorithmConfig"] = None, + local_env_runner: bool = True, + logdir: Optional[str] = None, + _setup: bool = True, + tune_trial_id: Optional[str] = None, + pg_offset: int = 0, + # Deprecated args. + num_env_runners: Optional[int] = None, + num_workers=DEPRECATED_VALUE, + local_worker=DEPRECATED_VALUE, + ): + """Initializes a EnvRunnerGroup instance. + + Args: + env_creator: Function that returns env given env config. + validate_env: Optional callable to validate the generated + environment (only on worker=0). This callable should raise + an exception if the environment is invalid. + default_policy_class: An optional default Policy class to use inside + the (multi-agent) `policies` dict. In case the PolicySpecs in there + have no class defined, use this `default_policy_class`. + If None, PolicySpecs will be using the Algorithm's default Policy + class. + config: Optional AlgorithmConfig (or config dict). + local_env_runner: Whether to create a local (non @ray.remote) EnvRunner + in the returned set as well (default: True). If `num_env_runners` + is 0, always create a local EnvRunner. + logdir: Optional logging directory for workers. + _setup: Whether to actually set up workers. This is only for testing. + tune_trial_id: The Ray Tune trial ID, if this EnvRunnerGroup is part of + an Algorithm run as a Tune trial. None, otherwise. + """ + if num_workers != DEPRECATED_VALUE or local_worker != DEPRECATED_VALUE: + deprecation_warning( + old="WorkerSet(num_workers=..., local_worker=...)", + new="EnvRunnerGroup(num_env_runners=..., local_env_runner=...)", + error=True, + ) + + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + # Make sure `config` is an AlgorithmConfig object. + if not config: + config = AlgorithmConfig() + elif isinstance(config, dict): + config = AlgorithmConfig.from_dict(config) + + self._env_creator = env_creator + self._policy_class = default_policy_class + self._remote_config = config + self._remote_config_obj_ref = ray.put(self._remote_config) + self._remote_args = { + "num_cpus": self._remote_config.num_cpus_per_env_runner, + "num_gpus": self._remote_config.num_gpus_per_env_runner, + "resources": self._remote_config.custom_resources_per_env_runner, + "max_restarts": ( + config.max_num_env_runner_restarts + if config.restart_failed_env_runners + else 0 + ), + } + self._tune_trial_id = tune_trial_id + self._pg_offset = pg_offset + + # Set the EnvRunner subclass to be used as "workers". Default: RolloutWorker. + self.env_runner_cls = config.env_runner_cls + if self.env_runner_cls is None: + if config.enable_env_runner_and_connector_v2: + # If experiences should be recorded, use the ` + # OfflineSingleAgentEnvRunner`. + if config.output: + # No multi-agent support. + if config.is_multi_agent: + raise ValueError("Multi-agent recording is not supported, yet.") + # Otherwise, load the single-agent env runner for + # recording. + else: + from ray.rllib.offline.offline_env_runner import ( + OfflineSingleAgentEnvRunner, + ) + + self.env_runner_cls = OfflineSingleAgentEnvRunner + else: + if config.is_multi_agent: + from ray.rllib.env.multi_agent_env_runner import ( + MultiAgentEnvRunner, + ) + + self.env_runner_cls = MultiAgentEnvRunner + else: + from ray.rllib.env.single_agent_env_runner import ( + SingleAgentEnvRunner, + ) + + self.env_runner_cls = SingleAgentEnvRunner + else: + self.env_runner_cls = RolloutWorker + + self._logdir = logdir + self._ignore_ray_errors_on_env_runners = ( + config.ignore_env_runner_failures or config.restart_failed_env_runners + ) + + # Create remote worker manager. + # ID=0 is used by the local worker. + # Starting remote workers from ID=1 to avoid conflicts. + self._worker_manager = FaultTolerantActorManager( + max_remote_requests_in_flight_per_actor=( + config.max_requests_in_flight_per_env_runner + ), + init_id=1, + ) + + if _setup: + try: + self._setup( + validate_env=validate_env, + config=config, + num_env_runners=( + num_env_runners + if num_env_runners is not None + else config.num_env_runners + ), + local_env_runner=local_env_runner, + ) + # EnvRunnerGroup creation possibly fails, if some (remote) workers cannot + # be initialized properly (due to some errors in the EnvRunners's + # constructor). + except RayActorError as e: + # In case of an actor (remote worker) init failure, the remote worker + # may still exist and will be accessible, however, e.g. calling + # its `sample.remote()` would result in strange "property not found" + # errors. + if e.actor_init_failed: + # Raise the original error here that the EnvRunners raised + # during its construction process. This is to enforce transparency + # for the user (better to understand the real reason behind the + # failure). + # - e.args[0]: The RayTaskError (inside the caught RayActorError). + # - e.args[0].args[2]: The original Exception (e.g. a ValueError due + # to a config mismatch) thrown inside the actor. + raise e.args[0].args[2] + # In any other case, raise the RayActorError as-is. + else: + raise e + + def _setup( + self, + *, + validate_env: Optional[Callable[[EnvType], None]] = None, + config: Optional["AlgorithmConfig"] = None, + num_env_runners: int = 0, + local_env_runner: bool = True, + ): + """Sets up an EnvRunnerGroup instance. + Args: + validate_env: Optional callable to validate the generated + environment (only on worker=0). + config: Optional dict that extends the common config of + the Algorithm class. + num_env_runners: Number of remote EnvRunner workers to create. + local_env_runner: Whether to create a local (non @ray.remote) EnvRunner + in the returned set as well (default: True). If `num_env_runners` + is 0, always create a local EnvRunner. + """ + # Force a local worker if num_env_runners == 0 (no remote workers). + # Otherwise, this EnvRunnerGroup would be empty. + self._local_env_runner = None + if num_env_runners == 0: + local_env_runner = True + # Create a local (learner) version of the config for the local worker. + # The only difference is the tf_session_args, which - for the local worker - + # will be `config.tf_session_args` updated/overridden with + # `config.local_tf_session_args`. + local_tf_session_args = config.tf_session_args.copy() + local_tf_session_args.update(config.local_tf_session_args) + self._local_config = config.copy(copy_frozen=False).framework( + tf_session_args=local_tf_session_args + ) + + if config.input_ == "dataset": + # Create the set of dataset readers to be shared by all the + # rollout workers. + self._ds, self._ds_shards = get_dataset_and_shards(config, num_env_runners) + else: + self._ds = None + self._ds_shards = None + + # Create a number of @ray.remote workers. + self.add_workers( + num_env_runners, + validate=config.validate_env_runners_after_construction, + ) + + # If num_env_runners > 0 and we don't have an env on the local worker, + # get the observation- and action spaces for each policy from + # the first remote worker (which does have an env). + if ( + local_env_runner + and self._worker_manager.num_actors() > 0 + and not config.create_env_on_local_worker + and (not config.observation_space or not config.action_space) + ): + spaces = self.get_spaces() + else: + spaces = None + + # Create a local worker, if needed. + if local_env_runner: + self._local_env_runner = self._make_worker( + env_creator=self._env_creator, + validate_env=validate_env, + worker_index=0, + num_workers=num_env_runners, + config=self._local_config, + spaces=spaces, + ) + + def get_spaces(self): + """Infer observation and action spaces from one (local or remote) EnvRunner. + + Returns: + A dict mapping from ModuleID to a 2-tuple containing obs- and action-space. + """ + # Get ID of the first remote worker. + remote_worker_ids = ( + [self._worker_manager.actor_ids()[0]] + if self._worker_manager.actor_ids() + else [] + ) + + spaces = self.foreach_env_runner( + lambda env_runner: env_runner.get_spaces(), + remote_worker_ids=remote_worker_ids, + local_env_runner=not remote_worker_ids, + )[0] + + logger.info( + "Inferred observation/action spaces from remote " + f"worker (local worker has no env): {spaces}" + ) + + return spaces + + @property + def local_env_runner(self) -> EnvRunner: + """Returns the local EnvRunner.""" + return self._local_env_runner + + def healthy_env_runner_ids(self) -> List[int]: + """Returns the list of remote worker IDs.""" + return self._worker_manager.healthy_actor_ids() + + def healthy_worker_ids(self) -> List[int]: + """Returns the list of remote worker IDs.""" + return self.healthy_env_runner_ids() + + def num_remote_env_runners(self) -> int: + """Returns the number of remote EnvRunners.""" + return self._worker_manager.num_actors() + + def num_remote_workers(self) -> int: + """Returns the number of remote EnvRunners.""" + return self.num_remote_env_runners() + + def num_healthy_remote_env_runners(self) -> int: + """Returns the number of healthy remote workers.""" + return self._worker_manager.num_healthy_actors() + + def num_healthy_remote_workers(self) -> int: + """Returns the number of healthy remote workers.""" + return self.num_healthy_remote_env_runners() + + def num_healthy_env_runners(self) -> int: + """Returns the number of all healthy workers, including the local worker.""" + return int(bool(self._local_env_runner)) + self.num_healthy_remote_workers() + + def num_healthy_workers(self) -> int: + """Returns the number of all healthy workers, including the local worker.""" + return self.num_healthy_env_runners() + + def num_in_flight_async_reqs(self, tag: Optional[str] = None) -> int: + """Returns the number of in-flight async requests.""" + return self._worker_manager.num_outstanding_async_reqs(tag=tag) + + def num_remote_worker_restarts(self) -> int: + """Total number of times managed remote workers have been restarted.""" + return self._worker_manager.total_num_restarts() + + def sync_env_runner_states( + self, + *, + config: "AlgorithmConfig", + from_worker: Optional[EnvRunner] = None, + env_steps_sampled: Optional[int] = None, + connector_states: Optional[List[Dict[str, Any]]] = None, + rl_module_state: Optional[Dict[str, Any]] = None, + env_runner_indices_to_update: Optional[List[int]] = None, + env_to_module=None, + module_to_env=None, + ) -> None: + """Synchronizes the connectors of this EnvRunnerGroup's EnvRunners. + + The exact procedure works as follows: + - If `from_worker` is None, set `from_worker=self.local_env_runner`. + - If `config.use_worker_filter_stats` is True, gather all remote EnvRunners' + ConnectorV2 states. Otherwise, only use the ConnectorV2 states of `from_worker`. + - Merge all gathered states into one resulting state. + - Broadcast the resulting state back to all remote EnvRunners AND the local + EnvRunner. + + Args: + config: The AlgorithmConfig object to use to determine, in which + direction(s) we need to synch and what the timeouts are. + from_worker: The EnvRunner from which to synch. If None, will use the local + worker of this EnvRunnerGroup. + env_steps_sampled: The total number of env steps taken thus far by all + workers combined. Used to broadcast this number to all remote workers + if `update_worker_filter_stats` is True in `config`. + env_runner_indices_to_update: The indices of those EnvRunners to update + with the merged state. Use None (default) to update all remote + EnvRunners. + """ + from_worker = from_worker or self.local_env_runner + + merge = ( + not config.enable_env_runner_and_connector_v2 + and config.use_worker_filter_stats + ) or ( + config.enable_env_runner_and_connector_v2 + and ( + config.merge_env_runner_states is True + or ( + config.merge_env_runner_states == "training_only" + and not config.in_evaluation + ) + ) + ) + broadcast = ( + not config.enable_env_runner_and_connector_v2 + and config.update_worker_filter_stats + ) or ( + config.enable_env_runner_and_connector_v2 + and config.broadcast_env_runner_states + ) + + # Early out if the number of (healthy) remote workers is 0. In this case, the + # local worker is the only operating worker and thus of course always holds + # the reference connector state. + if self.num_healthy_remote_workers() == 0 and self.local_env_runner: + self.local_env_runner.set_state( + { + **( + {NUM_ENV_STEPS_SAMPLED_LIFETIME: env_steps_sampled} + if env_steps_sampled is not None + else {} + ), + **(rl_module_state or {}), + } + ) + return + + # Also early out, if we don't merge AND don't broadcast. + if not merge and not broadcast: + return + + # Use states from all remote EnvRunners. + if merge: + if connector_states == []: + env_runner_states = {} + else: + if connector_states is None: + connector_states = self.foreach_env_runner( + lambda w: w.get_state( + components=[ + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + ] + ), + local_env_runner=False, + timeout_seconds=( + config.sync_filters_on_rollout_workers_timeout_s + ), + ) + env_to_module_states = [ + s[COMPONENT_ENV_TO_MODULE_CONNECTOR] + for s in connector_states + if COMPONENT_ENV_TO_MODULE_CONNECTOR in s + ] + module_to_env_states = [ + s[COMPONENT_MODULE_TO_ENV_CONNECTOR] + for s in connector_states + if COMPONENT_MODULE_TO_ENV_CONNECTOR in s + ] + + if ( + self.local_env_runner is not None + and hasattr(self.local_env_runner, "_env_to_module") + and hasattr(self.local_env_runner, "_module_to_env") + ): + assert env_to_module is None + env_to_module = self.local_env_runner._env_to_module + assert module_to_env is None + module_to_env = self.local_env_runner._module_to_env + + env_runner_states = {} + if env_to_module_states: + env_runner_states.update( + { + COMPONENT_ENV_TO_MODULE_CONNECTOR: ( + env_to_module.merge_states(env_to_module_states) + ), + } + ) + if module_to_env_states: + env_runner_states.update( + { + COMPONENT_MODULE_TO_ENV_CONNECTOR: ( + module_to_env.merge_states(module_to_env_states) + ), + } + ) + # Ignore states from remote EnvRunners (use the current `from_worker` states + # only). + else: + if from_worker is None: + env_runner_states = { + COMPONENT_ENV_TO_MODULE_CONNECTOR: env_to_module.get_state(), + COMPONENT_MODULE_TO_ENV_CONNECTOR: module_to_env.get_state(), + } + else: + env_runner_states = from_worker.get_state( + components=[ + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + ] + ) + + # Update the global number of environment steps, if necessary. + # Make sure to divide by the number of env runners (such that each EnvRunner + # knows (roughly) its own(!) lifetime count and can infer the global lifetime + # count from it). + if env_steps_sampled is not None: + env_runner_states[NUM_ENV_STEPS_SAMPLED_LIFETIME] = env_steps_sampled // ( + config.num_env_runners or 1 + ) + + # If we do NOT want remote EnvRunners to get their Connector states updated, + # only update the local worker here (with all state components, except the model + # weights) and then remove the connector components. + if not broadcast: + if self.local_env_runner is not None: + self.local_env_runner.set_state(env_runner_states) + else: + env_to_module.set_state( + env_runner_states.get(COMPONENT_ENV_TO_MODULE_CONNECTOR), {} + ) + module_to_env.set_state( + env_runner_states.get(COMPONENT_MODULE_TO_ENV_CONNECTOR), {} + ) + env_runner_states.pop(COMPONENT_ENV_TO_MODULE_CONNECTOR, None) + env_runner_states.pop(COMPONENT_MODULE_TO_ENV_CONNECTOR, None) + + # If there are components in the state left -> Update remote workers with these + # state components (and maybe the local worker, if it hasn't been updated yet). + if env_runner_states: + # Update the local EnvRunner, but NOT with the weights. If used at all for + # evaluation (through the user calling `self.evaluate`), RLlib would update + # the weights up front either way. + if self.local_env_runner is not None and broadcast: + self.local_env_runner.set_state(env_runner_states) + + # Send the model weights only to remote EnvRunners. + # In case the local EnvRunner is ever needed for evaluation, + # RLlib updates its weight right before such an eval step. + if rl_module_state: + env_runner_states.update(rl_module_state) + + # Broadcast updated states back to all workers. + # We explicitly don't want to fire and forget here, because this can lead to a lot of in-flight requests. + # When these pile up, object store memory can spike. + self.foreach_env_runner_async_fetch_ready( + func="set_state", + tag="set_state", + kwargs=dict(state=env_runner_states), + remote_worker_ids=env_runner_indices_to_update, + timeout_seconds=0.0, + ) + + def foreach_env_runner_async_fetch_ready( + self, + func: Union[ + Callable[[EnvRunner], T], List[Callable[[EnvRunner], T]], str, List[str] + ], + kwargs: Optional[Dict[str, Any]] = None, + tag: Optional[str] = None, + timeout_seconds: Optional[float] = 0.0, + return_obj_refs: bool = False, + mark_healthy: bool = False, + healthy_only: bool = True, + remote_worker_ids: List[int] = None, + return_actor_ids: bool = False, + ) -> List[Union[Tuple[int, T], T]]: + """Calls the given function asynchronously and returns previous results if any. + + This is a convenience function that calls the underlying actor manager's + `foreach_actor_async_fetch_ready()` method. + + """ + return self._worker_manager.foreach_actor_async_fetch_ready( + func=func, + tag=tag, + kwargs=kwargs, + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + healthy_only=healthy_only, + remote_actor_ids=remote_worker_ids, + ignore_ray_errors=self._ignore_ray_errors_on_env_runners, + return_actor_ids=return_actor_ids, + ) + + def sync_weights( + self, + policies: Optional[List[PolicyID]] = None, + from_worker_or_learner_group: Optional[Union[EnvRunner, "LearnerGroup"]] = None, + to_worker_indices: Optional[List[int]] = None, + global_vars: Optional[Dict[str, TensorType]] = None, + timeout_seconds: Optional[float] = 0.0, + inference_only: Optional[bool] = False, + ) -> None: + """Syncs model weights from the given weight source to all remote workers. + + Weight source can be either a (local) rollout worker or a learner_group. It + should just implement a `get_weights` method. + + Args: + policies: Optional list of PolicyIDs to sync weights for. + If None (default), sync weights to/from all policies. + from_worker_or_learner_group: Optional (local) EnvRunner instance or + LearnerGroup instance to sync from. If None (default), + sync from this EnvRunnerGroup's local worker. + to_worker_indices: Optional list of worker indices to sync the + weights to. If None (default), sync to all remote workers. + global_vars: An optional global vars dict to set this + worker to. If None, do not update the global_vars. + timeout_seconds: Timeout in seconds to wait for the sync weights + calls to complete. Default is 0.0 (fire-and-forget, do not wait + for any sync calls to finish). Setting this to 0.0 might significantly + improve algorithm performance, depending on the algo's `training_step` + logic. + inference_only: Sync weights with workers that keep inference-only + modules. This is needed for algorithms in the new stack that + use inference-only modules. In this case only a part of the + parameters are synced to the workers. Default is False. + """ + if self.local_env_runner is None and from_worker_or_learner_group is None: + raise TypeError( + "No `local_env_runner` in EnvRunnerGroup! Must provide " + "`from_worker_or_learner_group` arg in `sync_weights()`!" + ) + + # Only sync if we have remote workers or `from_worker_or_trainer` is provided. + rl_module_state = None + if self.num_remote_workers() or from_worker_or_learner_group is not None: + weights_src = ( + from_worker_or_learner_group + if from_worker_or_learner_group is not None + else self.local_env_runner + ) + + if weights_src is None: + raise ValueError( + "`from_worker_or_trainer` is None. In this case, EnvRunnerGroup " + "should have local_env_runner. But local_env_runner is also None." + ) + + modules = ( + [COMPONENT_RL_MODULE + "/" + p for p in policies] + if policies is not None + else [COMPONENT_RL_MODULE] + ) + # LearnerGroup has-a Learner, which has-a RLModule. + if isinstance(weights_src, LearnerGroup): + rl_module_state = weights_src.get_state( + components=[COMPONENT_LEARNER + "/" + m for m in modules], + inference_only=inference_only, + )[COMPONENT_LEARNER] + # EnvRunner (new API stack). + elif self._remote_config.enable_env_runner_and_connector_v2: + # EnvRunner (remote) has-a RLModule. + # TODO (sven): Replace this with a new ActorManager API: + # try_remote_request_till_success("get_state") -> tuple(int, + # remoteresult) + # `weights_src` could be the ActorManager, then. Then RLlib would know + # that it has to ping the manager to try all healthy actors until the + # first returns something. + if isinstance(weights_src, ray.actor.ActorHandle): + rl_module_state = ray.get( + weights_src.get_state.remote( + components=modules, + inference_only=inference_only, + ) + ) + # EnvRunner (local) has-a RLModule. + else: + rl_module_state = weights_src.get_state( + components=modules, + inference_only=inference_only, + ) + # RolloutWorker (old API stack). + else: + rl_module_state = weights_src.get_weights( + policies=policies, + inference_only=inference_only, + ) + + if self._remote_config.enable_env_runner_and_connector_v2: + + # Make sure `rl_module_state` only contains the weights and the + # weight seq no, nothing else. + rl_module_state = { + k: v + for k, v in rl_module_state.items() + if k in [COMPONENT_RL_MODULE, WEIGHTS_SEQ_NO] + } + + # Move weights to the object store to avoid having to make n pickled + # copies of the weights dict for each worker. + rl_module_state_ref = ray.put(rl_module_state) + + # Sync to specified remote workers in this EnvRunnerGroup. + # We explicitly don't want to fire and forget here, because this can lead to a lot of in-flight requests. + # When these pile up, object store memory can spike. + self.foreach_env_runner_async_fetch_ready( + func="set_state", + tag="set_state", + kwargs=dict(state=rl_module_state_ref), + remote_worker_ids=to_worker_indices, + timeout_seconds=timeout_seconds, + ) + + else: + rl_module_state_ref = ray.put(rl_module_state) + + def _set_weights(env_runner): + env_runner.set_weights(ray.get(rl_module_state_ref), global_vars) + + # Sync to specified remote workers in this EnvRunnerGroup. + self.foreach_env_runner( + func=_set_weights, + local_env_runner=False, # Do not sync back to local worker. + remote_worker_ids=to_worker_indices, + timeout_seconds=timeout_seconds, + ) + + # If `from_worker_or_learner_group` is provided, also sync to this + # EnvRunnerGroup's local worker. + if self.local_env_runner is not None: + if from_worker_or_learner_group is not None: + if self._remote_config.enable_env_runner_and_connector_v2: + self.local_env_runner.set_state(rl_module_state) + else: + self.local_env_runner.set_weights(rl_module_state) + # If `global_vars` is provided and local worker exists -> Update its + # global_vars. + if global_vars is not None: + self.local_env_runner.set_global_vars(global_vars) + + def add_workers(self, num_workers: int, validate: bool = False) -> None: + """Creates and adds a number of remote workers to this worker set. + + Can be called several times on the same EnvRunnerGroup to add more + EnvRunners to the set. + + Args: + num_workers: The number of remote Workers to add to this + EnvRunnerGroup. + validate: Whether to validate remote workers after their construction + process. + + Raises: + RayError: If any of the constructed remote workers is not up and running + properly. + """ + old_num_workers = self._worker_manager.num_actors() + new_workers = [ + self._make_worker( + env_creator=self._env_creator, + validate_env=None, + worker_index=old_num_workers + i + 1, + num_workers=old_num_workers + num_workers, + # self._remote_config can be large + # and it's best practice to pass it by reference + # instead of value (https://docs.ray.io/en/latest/ray-core/patterns/pass-large-arg-by-value.html) + config=self._remote_config_obj_ref, + ) + for i in range(num_workers) + ] + self._worker_manager.add_actors(new_workers) + + # Validate here, whether all remote workers have been constructed properly + # and are "up and running". Establish initial states. + if validate: + for result in self._worker_manager.foreach_actor( + lambda w: w.assert_healthy() + ): + # Simiply raise the error, which will get handled by the try-except + # clause around the _setup(). + if not result.ok: + e = result.get() + if self._ignore_ray_errors_on_env_runners: + logger.error(f"Validation of EnvRunner failed! Error={str(e)}") + else: + raise e + + def reset(self, new_remote_workers: List[ActorHandle]) -> None: + """Hard overrides the remote EnvRunners in this set with the provided ones. + + Args: + new_remote_workers: A list of new EnvRunners (as `ActorHandles`) to use as + new remote workers. + """ + self._worker_manager.clear() + self._worker_manager.add_actors(new_remote_workers) + + def stop(self) -> None: + """Calls `stop` on all EnvRunners (including the local one).""" + try: + # Make sure we stop all EnvRunners, include the ones that were just + # restarted / recovered or that are tagged unhealthy (at least, we should + # try). + self.foreach_env_runner( + lambda w: w.stop(), healthy_only=False, local_env_runner=True + ) + except Exception: + logger.exception("Failed to stop workers!") + finally: + self._worker_manager.clear() + + def foreach_env_runner( + self, + func: Union[ + Callable[[EnvRunner], T], List[Callable[[EnvRunner], T]], str, List[str] + ], + *, + kwargs=None, + local_env_runner: bool = True, + healthy_only: bool = True, + remote_worker_ids: List[int] = None, + timeout_seconds: Optional[float] = None, + return_obj_refs: bool = False, + mark_healthy: bool = False, + ) -> List[T]: + """Calls the given function with each EnvRunner as its argument. + + Args: + func: The function to call for each EnvRunners. The only call argument is + the respective EnvRunner instance. + local_env_runner: Whether to apply `func` to local EnvRunner, too. + Default is True. + healthy_only: Apply `func` on known-to-be healthy EnvRunners only. + remote_worker_ids: Apply `func` on a selected set of remote EnvRunners. + Use None (default) for all remote EnvRunners. + timeout_seconds: Time to wait (in seconds) for results. Set this to 0.0 for + fire-and-forget. Set this to None (default) to wait infinitely (i.e. for + synchronous execution). + return_obj_refs: Whether to return ObjectRef instead of actual results. + Note, for fault tolerance reasons, these returned ObjectRefs should + never be resolved with ray.get() outside of this EnvRunnerGroup. + mark_healthy: Whether to mark all those EnvRunners healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + Note that EnvRunners are NOT set unhealthy, if they simply time out + (only if they return a RayActorError). + Also note that this setting is ignored if `healthy_only=True` (b/c + `mark_healthy` only affects EnvRunners that are currently tagged as + unhealthy). + + Returns: + The list of return values of all calls to `func([worker])`. + """ + assert ( + not return_obj_refs or not local_env_runner + ), "Can not return ObjectRef from local worker." + + local_result = [] + if local_env_runner and self.local_env_runner is not None: + assert kwargs is None + if isinstance(func, str): + local_result = [getattr(self.local_env_runner, func)] + else: + local_result = [func(self.local_env_runner)] + + if not self._worker_manager.actor_ids(): + return local_result + + remote_results = self._worker_manager.foreach_actor( + func, + kwargs=kwargs, + healthy_only=healthy_only, + remote_actor_ids=remote_worker_ids, + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + FaultTolerantActorManager.handle_remote_call_result_errors( + remote_results, ignore_ray_errors=self._ignore_ray_errors_on_env_runners + ) + + # With application errors handled, return good results. + remote_results = [r.get() for r in remote_results.ignore_errors()] + + return local_result + remote_results + + def foreach_env_runner_async( + self, + func: Union[ + Callable[[EnvRunner], T], List[Callable[[EnvRunner], T]], str, List[str] + ], + tag: Optional[str] = None, + *, + kwargs=None, + healthy_only: bool = True, + remote_worker_ids: List[int] = None, + ) -> int: + """Calls the given function asynchronously with each EnvRunner as the argument. + + Does not return results directly. Instead, `fetch_ready_async_reqs()` can be + used to pull results in an async manner whenever they are available. + + Args: + func: The function to call for each EnvRunners. The only call argument is + the respective EnvRunner instance. + tag: A tag to identify the results from this async call when fetching with + `fetch_ready_async_reqs()`. + kwargs: An optional kwargs dict to be passed to the remote function calls. + healthy_only: Apply `func` on known-to-be healthy EnvRunners only. + remote_worker_ids: Apply `func` on a selected set of remote EnvRunners. + + Returns: + The number of async requests that have actually been made. This is the + length of `remote_worker_ids` (or self.num_remote_workers()` if + `remote_worker_ids` is None) minus the number of requests that were NOT + made b/c a remote EnvRunner already had its + `max_remote_requests_in_flight_per_actor` counter reached for this tag. + """ + + return self._worker_manager.foreach_actor_async( + func, + tag=tag, + kwargs=kwargs, + healthy_only=healthy_only, + remote_actor_ids=remote_worker_ids, + ) + + def fetch_ready_async_reqs( + self, + *, + tags: Optional[Union[str, List[str], Tuple[str]]] = None, + timeout_seconds: Optional[float] = 0.0, + return_obj_refs: bool = False, + mark_healthy: bool = False, + ) -> List[Tuple[int, T]]: + """Get results from outstanding asynchronous requests that are ready. + + Args: + tags: Tags to identify the results from a specific async call. + If None (default), returns results from all ready async requests. + If a single string, returns results from all ready async requests with that tag. + timeout_seconds: Time to wait for results. Default is 0, meaning + those requests that are already ready. + return_obj_refs: Whether to return ObjectRef instead of actual results. + mark_healthy: Whether to mark all those workers healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + Note that workers are NOT set unhealthy, if they simply time out + (only if they return a RayActorError). + Also note that this setting is ignored if `healthy_only=True` (b/c + `mark_healthy` only affects workers that are currently tagged as + unhealthy). + + Returns: + A list of results successfully returned from outstanding remote calls, + paired with the indices of the callee workers. + """ + # Get remote results + remote_results = self._worker_manager.fetch_ready_async_reqs( + tags=tags, + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + FaultTolerantActorManager.handle_remote_call_result_errors( + remote_results, + ignore_ray_errors=self._ignore_ray_errors_on_env_runners, + ) + + return [(r.actor_id, r.get()) for r in remote_results.ignore_errors()] + + @OldAPIStack + def foreach_env(self, func: Callable[[EnvType], List[T]]) -> List[List[T]]: + """Calls `func` with all workers' sub-environments as args. + + An "underlying sub environment" is a single clone of an env within + a vectorized environment. + `func` takes a single underlying sub environment as arg, e.g. a + gym.Env object. + + Args: + func: A function - taking an EnvType (normally a gym.Env object) + as arg and returning a list of lists of return values, one + value per underlying sub-environment per each worker. + + Returns: + The list (workers) of lists (sub environments) of results. + """ + return list( + self.foreach_env_runner( + lambda w: w.foreach_env(func), + local_env_runner=True, + ) + ) + + @OldAPIStack + def foreach_env_with_context( + self, func: Callable[[BaseEnv, EnvContext], List[T]] + ) -> List[List[T]]: + """Calls `func` with all workers' sub-environments and env_ctx as args. + + An "underlying sub environment" is a single clone of an env within + a vectorized environment. + `func` takes a single underlying sub environment and the env_context + as args. + + Args: + func: A function - taking a BaseEnv object and an EnvContext as + arg - and returning a list of lists of return values over envs + of the worker. + + Returns: + The list (1 item per workers) of lists (1 item per sub-environment) + of results. + """ + return list( + self.foreach_env_runner( + lambda w: w.foreach_env_with_context(func), + local_env_runner=True, + ) + ) + + def probe_unhealthy_env_runners(self) -> List[int]: + """Checks for unhealthy workers and tries restoring their states. + + Returns: + List of IDs of the workers that were restored. + """ + return self._worker_manager.probe_unhealthy_actors( + timeout_seconds=self._remote_config.env_runner_health_probe_timeout_s, + mark_healthy=True, + ) + + @OldAPIStack + def foreach_policy(self, func: Callable[[Policy, PolicyID], T]) -> List[T]: + """Calls `func` with each worker's (policy, PolicyID) tuple. + + Note that in the multi-agent case, each worker may have more than one + policy. + + Args: + func: A function - taking a Policy and its ID - that is + called on all workers' Policies. + + Returns: + The list of return values of func over all workers' policies. The + length of this list is: + (num_workers + 1 (local-worker)) * + [num policies in the multi-agent config dict]. + The local workers' results are first, followed by all remote + workers' results + """ + results = [] + for r in self.foreach_env_runner( + lambda w: w.foreach_policy(func), local_env_runner=True + ): + results.extend(r) + return results + + @OldAPIStack + def foreach_policy_to_train(self, func: Callable[[Policy, PolicyID], T]) -> List[T]: + """Apply `func` to all workers' Policies iff in `policies_to_train`. + + Args: + func: A function - taking a Policy and its ID - that is + called on all workers' Policies, for which + `worker.is_policy_to_train()` returns True. + + Returns: + List[any]: The list of n return values of all + `func([trainable policy], [ID])`-calls. + """ + results = [] + for r in self.foreach_env_runner( + lambda w: w.foreach_policy_to_train(func), local_env_runner=True + ): + results.extend(r) + return results + + @OldAPIStack + def is_policy_to_train( + self, policy_id: PolicyID, batch: Optional[SampleBatchType] = None + ) -> bool: + """Whether given PolicyID (optionally inside some batch) is trainable.""" + if self.local_env_runner: + if self.local_env_runner.is_policy_to_train is None: + return True + return self.local_env_runner.is_policy_to_train(policy_id, batch) + else: + raise NotImplementedError + + @OldAPIStack + def add_policy( + self, + policy_id: PolicyID, + policy_cls: Optional[Type[Policy]] = None, + policy: Optional[Policy] = None, + *, + observation_space: Optional[gym.spaces.Space] = None, + action_space: Optional[gym.spaces.Space] = None, + config: Optional[Union["AlgorithmConfig", PartialAlgorithmConfigDict]] = None, + policy_state: Optional[PolicyState] = None, + policy_mapping_fn: Optional[Callable[[AgentID, EpisodeID], PolicyID]] = None, + policies_to_train: Optional[ + Union[ + Collection[PolicyID], + Callable[[PolicyID, Optional[SampleBatchType]], bool], + ] + ] = None, + module_spec: Optional[RLModuleSpec] = None, + # Deprecated. + workers: Optional[List[Union[EnvRunner, ActorHandle]]] = DEPRECATED_VALUE, + ) -> None: + """Adds a policy to this EnvRunnerGroup's workers or a specific list of workers. + + Args: + policy_id: ID of the policy to add. + policy_cls: The Policy class to use for constructing the new Policy. + Note: Only one of `policy_cls` or `policy` must be provided. + policy: The Policy instance to add to this EnvRunnerGroup. If not None, the + given Policy object will be directly inserted into the + local worker and clones of that Policy will be created on all remote + workers. + Note: Only one of `policy_cls` or `policy` must be provided. + observation_space: The observation space of the policy to add. + If None, try to infer this space from the environment. + action_space: The action space of the policy to add. + If None, try to infer this space from the environment. + config: The config object or overrides for the policy to add. + policy_state: Optional state dict to apply to the new + policy instance, right after its construction. + policy_mapping_fn: An optional (updated) policy mapping function + to use from here on. Note that already ongoing episodes will + not change their mapping but will use the old mapping till + the end of the episode. + policies_to_train: An optional list of policy IDs to be trained + or a callable taking PolicyID and SampleBatchType and + returning a bool (trainable or not?). + If None, will keep the existing setup in place. Policies, + whose IDs are not in the list (or for which the callable + returns False) will not be updated. + module_spec: In the new RLModule API we need to pass in the module_spec for + the new module that is supposed to be added. Knowing the policy spec is + not sufficient. + workers: A list of EnvRunner/ActorHandles (remote + EnvRunners) to add this policy to. If defined, will only + add the given policy to these workers. + + Raises: + KeyError: If the given `policy_id` already exists in this EnvRunnerGroup. + """ + if self.local_env_runner and policy_id in self.local_env_runner.policy_map: + raise KeyError( + f"Policy ID '{policy_id}' already exists in policy map! " + "Make sure you use a Policy ID that has not been taken yet." + " Policy IDs that are already in your policy map: " + f"{list(self.local_env_runner.policy_map.keys())}" + ) + + if workers is not DEPRECATED_VALUE: + deprecation_warning( + old="EnvRunnerGroup.add_policy(.., workers=..)", + help=( + "The `workers` argument to `EnvRunnerGroup.add_policy()` is " + "deprecated! Please do not use it anymore." + ), + error=True, + ) + + if (policy_cls is None) == (policy is None): + raise ValueError( + "Only one of `policy_cls` or `policy` must be provided to " + "staticmethod: `EnvRunnerGroup.add_policy()`!" + ) + validate_module_id(policy_id, error=False) + + # Policy instance not provided: Use the information given here. + if policy_cls is not None: + new_policy_instance_kwargs = dict( + policy_id=policy_id, + policy_cls=policy_cls, + observation_space=observation_space, + action_space=action_space, + config=config, + policy_state=policy_state, + policy_mapping_fn=policy_mapping_fn, + policies_to_train=list(policies_to_train) + if policies_to_train + else None, + module_spec=module_spec, + ) + # Policy instance provided: Create clones of this very policy on the different + # workers (copy all its properties here for the calls to add_policy on the + # remote workers). + else: + new_policy_instance_kwargs = dict( + policy_id=policy_id, + policy_cls=type(policy), + observation_space=policy.observation_space, + action_space=policy.action_space, + config=policy.config, + policy_state=policy.get_state(), + policy_mapping_fn=policy_mapping_fn, + policies_to_train=list(policies_to_train) + if policies_to_train + else None, + module_spec=module_spec, + ) + + def _create_new_policy_fn(worker): + # `foreach_env_runner` function: Adds the policy to the worker (and + # maybe changes its policy_mapping_fn - if provided here). + worker.add_policy(**new_policy_instance_kwargs) + + if self.local_env_runner is not None: + # Add policy directly by (already instantiated) object. + if policy is not None: + self.local_env_runner.add_policy( + policy_id=policy_id, + policy=policy, + policy_mapping_fn=policy_mapping_fn, + policies_to_train=policies_to_train, + module_spec=module_spec, + ) + # Add policy by constructor kwargs. + else: + self.local_env_runner.add_policy(**new_policy_instance_kwargs) + + # Add the policy to all remote workers. + self.foreach_env_runner(_create_new_policy_fn, local_env_runner=False) + + def _make_worker( + self, + *, + env_creator: EnvCreator, + validate_env: Optional[Callable[[EnvType], None]], + worker_index: int, + num_workers: int, + recreated_worker: bool = False, + config: "AlgorithmConfig", + spaces: Optional[ + Dict[PolicyID, Tuple[gym.spaces.Space, gym.spaces.Space]] + ] = None, + ) -> Union[EnvRunner, ActorHandle]: + kwargs = dict( + env_creator=env_creator, + validate_env=validate_env, + default_policy_class=self._policy_class, + config=config, + worker_index=worker_index, + num_workers=num_workers, + recreated_worker=recreated_worker, + log_dir=self._logdir, + spaces=spaces, + dataset_shards=self._ds_shards, + tune_trial_id=self._tune_trial_id, + ) + + if worker_index == 0: + return self.env_runner_cls(**kwargs) + + pg_bundle_idx = ( + -1 + if ray.util.get_current_placement_group() is None + else self._pg_offset + worker_index + ) + return ( + ray.remote(**self._remote_args)(self.env_runner_cls) + .options(placement_group_bundle_index=pg_bundle_idx) + .remote(**kwargs) + ) + + @classmethod + def _valid_module(cls, class_path): + del cls + if ( + isinstance(class_path, str) + and not os.path.isfile(class_path) + and "." in class_path + ): + module_path, class_name = class_path.rsplit(".", 1) + try: + spec = importlib.util.find_spec(module_path) + if spec is not None: + return True + except (ModuleNotFoundError, ValueError): + print( + f"module {module_path} not found while trying to get " + f"input {class_path}" + ) + return False diff --git a/lib/python3.12/site-packages/ray/rllib/env/external/__init__.py b/lib/python3.12/site-packages/ray/rllib/env/external/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9264f89c978e3ee878fb87a7c27f8406d5adcd15 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/external/__init__.py @@ -0,0 +1,11 @@ +from ray.rllib.env.external.rllink import ( + RLlink, + get_rllink_message, + send_rllink_message, +) + +__all__ = [ + "get_rllink_message", + "send_rllink_message", + "RLlink", +] diff --git a/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b6a9d01d4d7f29f23f6e192e34065bf92ab650d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/env_runner_server_for_external_inference.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/env_runner_server_for_external_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b041ce61a28c82a06a2b72738ab64ba29636bc05 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/env_runner_server_for_external_inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/rllink.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/rllink.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa6c43377293aa9d2857ca5bf11a8e75b746be1d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/external/__pycache__/rllink.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/external/env_runner_server_for_external_inference.py b/lib/python3.12/site-packages/ray/rllib/env/external/env_runner_server_for_external_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..59b55a590fcba76b4d3b66d45fca77b1bb8573bc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/external/env_runner_server_for_external_inference.py @@ -0,0 +1,368 @@ +import pickle +import socket +import threading +import time +from collections import defaultdict +from typing import Collection, DefaultDict, List, Optional, Union + +from ray.rllib.core import ( + COMPONENT_RL_MODULE, + DEFAULT_AGENT_ID, + DEFAULT_MODULE_ID, +) +from ray.rllib.env import INPUT_ENV_SPACES +from ray.rllib.env.env_runner import EnvRunner +from ray.rllib.env.external.rllink import ( + RLlink, + get_rllink_message, + send_rllink_message, +) +from ray.rllib.env.single_agent_env_runner import SingleAgentEnvRunner +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils.annotations import override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + EPISODE_DURATION_SEC_MEAN, + EPISODE_LEN_MAX, + EPISODE_LEN_MEAN, + EPISODE_LEN_MIN, + EPISODE_RETURN_MAX, + EPISODE_RETURN_MEAN, + EPISODE_RETURN_MIN, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.typing import EpisodeID, StateDict +from ray.util.annotations import DeveloperAPI + +torch, _ = try_import_torch() + + +@DeveloperAPI +class EnvRunnerServerForExternalInference(EnvRunner, Checkpointable): + """An EnvRunner communicating with an external env through a TCP socket. + + This implementation assumes: + - Only one external client ever connects to this env runner. + - The external client owns the connector pipelines (env-to-module and module-to-env) + as well as the RLModule and thus performs inference locally. Samples are sent in + bulk as lists of RLlib episodes once a certain number of timesteps has been executed + on the client's side. + - A copy of the RLModule is kept at all times on this EnvRunner, but is never used + for inference, only as a weights container. + TODO (sven): The above might be inefficient as we have to store basically two + models, one in this EnvRunner, one in the env (as ONNX). + - As a consequence, there are no environment and no connectors on this env runner. + The external env is responsible for generating all the data to create episodes. + """ + + @override(EnvRunner) + def __init__(self, *, config, **kwargs): + """ + Initializes an EnvRunnerServerForExternalInference instance. + + Args: + config: The AlgorithmConfig to use for setup. + + Keyword Args: + port: The base port number. The server socket is then actually bound to + `port` + self.worker_index. + """ + super().__init__(config=config, **kwargs) + + self.worker_index: int = kwargs.get("worker_index", 0) + + self._weights_seq_no = 0 + + # Build the module from its spec. + module_spec = self.config.get_rl_module_spec( + spaces=self.get_spaces(), inference_only=True + ) + self.module = module_spec.build() + + self.host = "localhost" + self.port = int(self.config.env_config.get("port", 5555)) + self.worker_index + self.server_socket = None + self.client_socket = None + self.address = None + + self.metrics = MetricsLogger() + + self._episode_chunks_to_return: Optional[List[SingleAgentEpisode]] = None + self._done_episodes_for_metrics: List[SingleAgentEpisode] = [] + self._ongoing_episodes_for_metrics: DefaultDict[ + EpisodeID, List[SingleAgentEpisode] + ] = defaultdict(list) + + self._sample_lock = threading.Lock() + self._on_policy_lock = threading.Lock() + self._blocked_on_state = False + + # Start a background thread for client communication. + self.thread = threading.Thread( + target=self._client_message_listener, daemon=True + ) + self.thread.start() + + @override(EnvRunner) + def assert_healthy(self): + """Checks that the server socket is open and listening.""" + assert ( + self.server_socket is not None + ), "Server socket is None (not connected, not listening)." + + @override(EnvRunner) + def sample(self, **kwargs): + """Waits for the client to send episodes.""" + while True: + with self._sample_lock: + if self._episode_chunks_to_return is not None: + num_env_steps = 0 + num_episodes_completed = 0 + for eps in self._episode_chunks_to_return: + if eps.is_done: + self._done_episodes_for_metrics.append(eps) + num_episodes_completed += 1 + else: + self._ongoing_episodes_for_metrics[eps.id_].append(eps) + num_env_steps += len(eps) + + ret = self._episode_chunks_to_return + self._episode_chunks_to_return = None + + SingleAgentEnvRunner._increase_sampled_metrics( + self, num_env_steps, num_episodes_completed + ) + + return ret + time.sleep(0.01) + + @override(EnvRunner) + def get_metrics(self): + # TODO (sven): We should probably make this a utility function to be called + # from within Single/MultiAgentEnvRunner and other EnvRunner subclasses, as + # needed. + # Compute per-episode metrics (only on already completed episodes). + for eps in self._done_episodes_for_metrics: + assert eps.is_done + episode_length = len(eps) + episode_return = eps.get_return() + episode_duration_s = eps.get_duration_s() + # Don't forget about the already returned chunks of this episode. + if eps.id_ in self._ongoing_episodes_for_metrics: + for eps2 in self._ongoing_episodes_for_metrics[eps.id_]: + episode_length += len(eps2) + episode_return += eps2.get_return() + episode_duration_s += eps2.get_duration_s() + del self._ongoing_episodes_for_metrics[eps.id_] + + self._log_episode_metrics( + episode_length, episode_return, episode_duration_s + ) + + # Now that we have logged everything, clear cache of done episodes. + self._done_episodes_for_metrics.clear() + + # Return reduced metrics. + return self.metrics.reduce() + + def get_spaces(self): + return { + INPUT_ENV_SPACES: (self.config.observation_space, self.config.action_space), + DEFAULT_MODULE_ID: ( + self.config.observation_space, + self.config.action_space, + ), + } + + @override(EnvRunner) + def stop(self): + """Closes the client and server sockets.""" + self._close_sockets_if_necessary() + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + {"config": self.config}, # **kwargs + ) + + @override(Checkpointable) + def get_checkpointable_components(self): + return [ + (COMPONENT_RL_MODULE, self.module), + ] + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + return { + COMPONENT_RL_MODULE: self.module.get_state(), + WEIGHTS_SEQ_NO: self._weights_seq_no, + } + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + # Update the RLModule state. + if COMPONENT_RL_MODULE in state: + # A missing value for WEIGHTS_SEQ_NO or a value of 0 means: Force the + # update. + weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0) + + # Only update the weigths, if this is the first synchronization or + # if the weights of this `EnvRunner` lacks behind the actual ones. + if weights_seq_no == 0 or self._weights_seq_no < weights_seq_no: + rl_module_state = state[COMPONENT_RL_MODULE] + if ( + isinstance(rl_module_state, dict) + and DEFAULT_MODULE_ID in rl_module_state + ): + rl_module_state = rl_module_state[DEFAULT_MODULE_ID] + self.module.set_state(rl_module_state) + + # Update our weights_seq_no, if the new one is > 0. + if weights_seq_no > 0: + self._weights_seq_no = weights_seq_no + + if self._blocked_on_state is True: + self._send_set_state_message() + self._blocked_on_state = False + + def _client_message_listener(self): + """Entry point for the listener thread.""" + + # Set up the server socket and bind to the specified host and port. + self._recycle_sockets() + + # Enter an endless message receival- and processing loop. + while True: + # As long as we are blocked on a new state, sleep a bit and continue. + # Do NOT process any incoming messages (until we send out the new state + # back to the client). + if self._blocked_on_state is True: + time.sleep(0.01) + continue + + try: + # Blocking call to get next message. + msg_type, msg_body = get_rllink_message(self.client_socket) + + # Process the message received based on its type. + # Initial handshake. + if msg_type == RLlink.PING: + self._send_pong_message() + + # Episode data from the client. + elif msg_type in [ + RLlink.EPISODES, + RLlink.EPISODES_AND_GET_STATE, + ]: + self._process_episodes_message(msg_type, msg_body) + + # Client requests the state (model weights). + elif msg_type == RLlink.GET_STATE: + self._send_set_state_message() + + # Clients requests config information. + elif msg_type == RLlink.GET_CONFIG: + self._send_set_config_message() + + except ConnectionError as e: + print(f"Messaging/connection error {e}! Recycling sockets ...") + self._recycle_sockets(5.0) + continue + + def _recycle_sockets(self, sleep: float = 0.0): + # Close all old sockets, if they exist. + self._close_sockets_if_necessary() + + time.sleep(sleep) + + # Start listening on the configured port. + self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + # Allow reuse of the address. + self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.server_socket.bind((self.host, self.port)) + # Listen for a single connection. + self.server_socket.listen(1) + print(f"Waiting for client to connect to port {self.port}...") + + self.client_socket, self.address = self.server_socket.accept() + print(f"Connected to client at {self.address}") + + def _close_sockets_if_necessary(self): + if self.client_socket: + self.client_socket.close() + if self.server_socket: + self.server_socket.close() + + def _send_pong_message(self): + send_rllink_message(self.client_socket, {"type": RLlink.PONG.name}) + + def _process_episodes_message(self, msg_type, msg_body): + # On-policy training -> we have to block until we get a new `set_state` call + # (b/c the learning step is done and we can send new weights back to all + # clients). + if msg_type == RLlink.EPISODES_AND_GET_STATE: + self._blocked_on_state = True + + episodes = [] + for episode_state in msg_body["episodes"]: + episode = SingleAgentEpisode.from_state(episode_state) + episodes.append(episode.to_numpy()) + + # Push episodes into the to-be-returned list (for `sample()` requests). + with self._sample_lock: + if isinstance(self._episode_chunks_to_return, list): + self._episode_chunks_to_return.extend(episodes) + else: + self._episode_chunks_to_return = episodes + + def _send_set_state_message(self): + send_rllink_message( + self.client_socket, + { + "type": RLlink.SET_STATE.name, + "state": self.get_state(inference_only=True), + }, + ) + + def _send_set_config_message(self): + send_rllink_message( + self.client_socket, + { + "type": RLlink.SET_CONFIG.name, + # TODO (sven): We need AlgorithmConfig to be a `Checkpointable` with a + # msgpack'able state. + "config": pickle.dumps(self.config), + }, + ) + + def _log_episode_metrics(self, length, ret, sec): + # Log general episode metrics. + # To mimic the old API stack behavior, we'll use `window` here for + # these particular stats (instead of the default EMA). + win = self.config.metrics_num_episodes_for_smoothing + self.metrics.log_value(EPISODE_LEN_MEAN, length, window=win) + self.metrics.log_value(EPISODE_RETURN_MEAN, ret, window=win) + self.metrics.log_value(EPISODE_DURATION_SEC_MEAN, sec, window=win) + # Per-agent returns. + self.metrics.log_value( + ("agent_episode_returns_mean", DEFAULT_AGENT_ID), ret, window=win + ) + # Per-RLModule returns. + self.metrics.log_value( + ("module_episode_returns_mean", DEFAULT_MODULE_ID), ret, window=win + ) + + # For some metrics, log min/max as well. + self.metrics.log_value(EPISODE_LEN_MIN, length, reduce="min", window=win) + self.metrics.log_value(EPISODE_RETURN_MIN, ret, reduce="min", window=win) + self.metrics.log_value(EPISODE_LEN_MAX, length, reduce="max", window=win) + self.metrics.log_value(EPISODE_RETURN_MAX, ret, reduce="max", window=win) diff --git a/lib/python3.12/site-packages/ray/rllib/env/external/rllink.py b/lib/python3.12/site-packages/ray/rllib/env/external/rllink.py new file mode 100644 index 0000000000000000000000000000000000000000..12ec0c3c5f043fc481fc4ac1d1dc484e6c4eab00 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/external/rllink.py @@ -0,0 +1,109 @@ +from enum import Enum + +from packaging.version import Version + +from ray.rllib.utils.checkpoints import try_import_msgpack +from ray.util.annotations import DeveloperAPI + +msgpack = None + + +@DeveloperAPI +class RLlink(Enum): + PROTOCOL_VERSION = Version("0.0.1") + + # Requests: Client (external env) -> Server (RLlib). + # ---- + # Ping command (initial handshake). + PING = "PING" + # List of episodes (similar to what an EnvRunner.sample() call would return). + EPISODES = "EPISODES" + # Request state (e.g. model weights). + GET_STATE = "GET_STATE" + # Request Algorithm config. + GET_CONFIG = "GET_CONFIG" + # Send episodes and request the next state update right after that. + # Clients sending this message should wait for a SET_STATE message as an immediate + # response. Useful for external samplers that must collect on-policy data. + EPISODES_AND_GET_STATE = "EPISODES_AND_GET_STATE" + + # Responses: Server (RLlib) -> Client (external env). + # ---- + # Pong response (initial handshake). + PONG = "PONG" + # Set state (e.g. model weights). + SET_STATE = "SET_STATE" + # Set Algorithm config. + SET_CONFIG = "SET_CONFIG" + + # @OldAPIStack (to be deprecated soon). + ACTION_SPACE = "ACTION_SPACE" + OBSERVATION_SPACE = "OBSERVATION_SPACE" + GET_WORKER_ARGS = "GET_WORKER_ARGS" + GET_WEIGHTS = "GET_WEIGHTS" + REPORT_SAMPLES = "REPORT_SAMPLES" + START_EPISODE = "START_EPISODE" + GET_ACTION = "GET_ACTION" + LOG_ACTION = "LOG_ACTION" + LOG_RETURNS = "LOG_RETURNS" + END_EPISODE = "END_EPISODE" + + def __str__(self): + return self.name + + +@DeveloperAPI +def send_rllink_message(sock_, message: dict): + """Sends a message to the client with a length header.""" + global msgpack + if msgpack is None: + msgpack = try_import_msgpack(error=True) + + body = msgpack.packb(message, use_bin_type=True) # .encode("utf-8") + header = str(len(body)).zfill(8).encode("utf-8") + try: + sock_.sendall(header + body) + except Exception as e: + raise ConnectionError( + f"Error sending message {message} to server on socket {sock_}! " + f"Original error was: {e}" + ) + + +@DeveloperAPI +def get_rllink_message(sock_): + """Receives a message from the client following the length-header protocol.""" + global msgpack + if msgpack is None: + msgpack = try_import_msgpack(error=True) + + try: + # Read the length header (8 bytes) + header = _get_num_bytes(sock_, 8) + msg_length = int(header.decode("utf-8")) + # Read the message body + body = _get_num_bytes(sock_, msg_length) + # Decode JSON. + message = msgpack.unpackb(body, raw=False) # .loads(body.decode("utf-8")) + # Check for proper protocol. + if "type" not in message: + raise ConnectionError( + "Protocol Error! Message from peer does not contain `type` field." + ) + return RLlink(message.pop("type")), message + except Exception as e: + raise ConnectionError( + f"Error receiving message from peer on socket {sock_}! " + f"Original error was: {e}" + ) + + +def _get_num_bytes(sock_, num_bytes): + """Helper function to receive a specific number of bytes.""" + data = b"" + while len(data) < num_bytes: + packet = sock_.recv(num_bytes - len(data)) + if not packet: + raise ConnectionError(f"No data received from socket {sock_}!") + data += packet + return data diff --git a/lib/python3.12/site-packages/ray/rllib/env/external_env.py b/lib/python3.12/site-packages/ray/rllib/env/external_env.py new file mode 100644 index 0000000000000000000000000000000000000000..c9aae38f1852da855b6eb58eb0ce854ca442161a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/external_env.py @@ -0,0 +1,482 @@ +import queue +import threading +import uuid +from typing import TYPE_CHECKING, Callable, Optional, Tuple + +import gymnasium as gym + +from ray._common.deprecation import deprecation_warning +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.typing import ( + EnvActionType, + EnvInfoDict, + EnvObsType, + EnvType, + MultiEnvDict, +) + +if TYPE_CHECKING: + from ray.rllib.models.preprocessors import Preprocessor + + +@OldAPIStack +class ExternalEnv(threading.Thread): + """An environment that interfaces with external agents. + + Unlike simulator envs, control is inverted: The environment queries the + policy to obtain actions and in return logs observations and rewards for + training. This is in contrast to gym.Env, where the algorithm drives the + simulation through env.step() calls. + + You can use ExternalEnv as the backend for policy serving (by serving HTTP + requests in the run loop), for ingesting offline logs data (by reading + offline transitions in the run loop), or other custom use cases not easily + expressed through gym.Env. + + ExternalEnv supports both on-policy actions (through self.get_action()), + and off-policy actions (through self.log_action()). + + This env is thread-safe, but individual episodes must be executed serially. + + .. testcode:: + :skipif: True + + from ray.tune import register_env + from ray.rllib.algorithms.dqn import DQN + YourExternalEnv = ... + register_env("my_env", lambda config: YourExternalEnv(config)) + algo = DQN(env="my_env") + while True: + print(algo.train()) + """ + + def __init__( + self, + action_space: gym.Space, + observation_space: gym.Space, + max_concurrent: int = None, + ): + """Initializes an ExternalEnv instance. + + Args: + action_space: Action space of the env. + observation_space: Observation space of the env. + """ + + threading.Thread.__init__(self) + + self.daemon = True + self.action_space = action_space + self.observation_space = observation_space + self._episodes = {} + self._finished = set() + self._results_avail_condition = threading.Condition() + if max_concurrent is not None: + deprecation_warning( + "The `max_concurrent` argument has been deprecated. Please configure" + "the number of episodes using the `rollout_fragment_length` and" + "`batch_mode` arguments. Please raise an issue on the Ray Github if " + "these arguments do not support your expected use case for ExternalEnv", + error=True, + ) + + def run(self): + """Override this to implement the run loop. + + Your loop should continuously: + 1. Call self.start_episode(episode_id) + 2. Call self.[get|log]_action(episode_id, obs, [action]?) + 3. Call self.log_returns(episode_id, reward) + 4. Call self.end_episode(episode_id, obs) + 5. Wait if nothing to do. + + Multiple episodes may be started at the same time. + """ + raise NotImplementedError + + def start_episode( + self, episode_id: Optional[str] = None, training_enabled: bool = True + ) -> str: + """Record the start of an episode. + + Args: + episode_id: Unique string id for the episode or + None for it to be auto-assigned and returned. + training_enabled: Whether to use experiences for this + episode to improve the policy. + + Returns: + Unique string id for the episode. + """ + + if episode_id is None: + episode_id = uuid.uuid4().hex + + if episode_id in self._finished: + raise ValueError("Episode {} has already completed.".format(episode_id)) + + if episode_id in self._episodes: + raise ValueError("Episode {} is already started".format(episode_id)) + + self._episodes[episode_id] = _ExternalEnvEpisode( + episode_id, self._results_avail_condition, training_enabled + ) + + return episode_id + + def get_action(self, episode_id: str, observation: EnvObsType) -> EnvActionType: + """Record an observation and get the on-policy action. + + Args: + episode_id: Episode id returned from start_episode(). + observation: Current environment observation. + + Returns: + Action from the env action space. + """ + + episode = self._get(episode_id) + return episode.wait_for_action(observation) + + def log_action( + self, episode_id: str, observation: EnvObsType, action: EnvActionType + ) -> None: + """Record an observation and (off-policy) action taken. + + Args: + episode_id: Episode id returned from start_episode(). + observation: Current environment observation. + action: Action for the observation. + """ + + episode = self._get(episode_id) + episode.log_action(observation, action) + + def log_returns( + self, episode_id: str, reward: float, info: Optional[EnvInfoDict] = None + ) -> None: + """Records returns (rewards and infos) from the environment. + + The reward will be attributed to the previous action taken by the + episode. Rewards accumulate until the next action. If no reward is + logged before the next action, a reward of 0.0 is assumed. + + Args: + episode_id: Episode id returned from start_episode(). + reward: Reward from the environment. + info: Optional info dict. + """ + + episode = self._get(episode_id) + episode.cur_reward += reward + + if info: + episode.cur_info = info or {} + + def end_episode(self, episode_id: str, observation: EnvObsType) -> None: + """Records the end of an episode. + + Args: + episode_id: Episode id returned from start_episode(). + observation: Current environment observation. + """ + + episode = self._get(episode_id) + self._finished.add(episode.episode_id) + episode.done(observation) + + def _get(self, episode_id: str) -> "_ExternalEnvEpisode": + """Get a started episode by its ID or raise an error.""" + + if episode_id in self._finished: + raise ValueError("Episode {} has already completed.".format(episode_id)) + + if episode_id not in self._episodes: + raise ValueError("Episode {} not found.".format(episode_id)) + + return self._episodes[episode_id] + + def to_base_env( + self, + make_env: Optional[Callable[[int], EnvType]] = None, + num_envs: int = 1, + remote_envs: bool = False, + remote_env_batch_wait_ms: int = 0, + restart_failed_sub_environments: bool = False, + ) -> "BaseEnv": + """Converts an RLlib MultiAgentEnv into a BaseEnv object. + + The resulting BaseEnv is always vectorized (contains n + sub-environments) to support batched forward passes, where n may + also be 1. BaseEnv also supports async execution via the `poll` and + `send_actions` methods and thus supports external simulators. + + Args: + make_env: A callable taking an int as input (which indicates + the number of individual sub-environments within the final + vectorized BaseEnv) and returning one individual + sub-environment. + num_envs: The number of sub-environments to create in the + resulting (vectorized) BaseEnv. The already existing `env` + will be one of the `num_envs`. + remote_envs: Whether each sub-env should be a @ray.remote + actor. You can set this behavior in your config via the + `remote_worker_envs=True` option. + remote_env_batch_wait_ms: The wait time (in ms) to poll remote + sub-environments for, if applicable. Only used if + `remote_envs` is True. + + Returns: + The resulting BaseEnv object. + """ + if num_envs != 1: + raise ValueError( + "External(MultiAgent)Env does not currently support " + "num_envs > 1. One way of solving this would be to " + "treat your Env as a MultiAgentEnv hosting only one " + "type of agent but with several copies." + ) + env = ExternalEnvWrapper(self) + + return env + + +@OldAPIStack +class _ExternalEnvEpisode: + """Tracked state for each active episode.""" + + def __init__( + self, + episode_id: str, + results_avail_condition: threading.Condition, + training_enabled: bool, + multiagent: bool = False, + ): + self.episode_id = episode_id + self.results_avail_condition = results_avail_condition + self.training_enabled = training_enabled + self.multiagent = multiagent + self.data_queue = queue.Queue() + self.action_queue = queue.Queue() + if multiagent: + self.new_observation_dict = None + self.new_action_dict = None + self.cur_reward_dict = {} + self.cur_terminated_dict = {"__all__": False} + self.cur_truncated_dict = {"__all__": False} + self.cur_info_dict = {} + else: + self.new_observation = None + self.new_action = None + self.cur_reward = 0.0 + self.cur_terminated = False + self.cur_truncated = False + self.cur_info = {} + + def get_data(self): + if self.data_queue.empty(): + return None + return self.data_queue.get_nowait() + + def log_action(self, observation, action): + if self.multiagent: + self.new_observation_dict = observation + self.new_action_dict = action + else: + self.new_observation = observation + self.new_action = action + self._send() + self.action_queue.get(True, timeout=60.0) + + def wait_for_action(self, observation): + if self.multiagent: + self.new_observation_dict = observation + else: + self.new_observation = observation + self._send() + return self.action_queue.get(True, timeout=300.0) + + def done(self, observation): + if self.multiagent: + self.new_observation_dict = observation + self.cur_terminated_dict = {"__all__": True} + # TODO(sven): External env API does not currently support truncated, + # but we should deprecate external Env anyways in favor of a client-only + # approach. + self.cur_truncated_dict = {"__all__": False} + else: + self.new_observation = observation + self.cur_terminated = True + self.cur_truncated = False + self._send() + + def _send(self): + if self.multiagent: + if not self.training_enabled: + for agent_id in self.cur_info_dict: + self.cur_info_dict[agent_id]["training_enabled"] = False + item = { + "obs": self.new_observation_dict, + "reward": self.cur_reward_dict, + "terminated": self.cur_terminated_dict, + "truncated": self.cur_truncated_dict, + "info": self.cur_info_dict, + } + if self.new_action_dict is not None: + item["off_policy_action"] = self.new_action_dict + self.new_observation_dict = None + self.new_action_dict = None + self.cur_reward_dict = {} + else: + item = { + "obs": self.new_observation, + "reward": self.cur_reward, + "terminated": self.cur_terminated, + "truncated": self.cur_truncated, + "info": self.cur_info, + } + if self.new_action is not None: + item["off_policy_action"] = self.new_action + self.new_observation = None + self.new_action = None + self.cur_reward = 0.0 + if not self.training_enabled: + item["info"]["training_enabled"] = False + + with self.results_avail_condition: + self.data_queue.put_nowait(item) + self.results_avail_condition.notify() + + +@OldAPIStack +class ExternalEnvWrapper(BaseEnv): + """Internal adapter of ExternalEnv to BaseEnv.""" + + def __init__( + self, external_env: "ExternalEnv", preprocessor: "Preprocessor" = None + ): + from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv + + self.external_env = external_env + self.prep = preprocessor + self.multiagent = issubclass(type(external_env), ExternalMultiAgentEnv) + self._action_space = external_env.action_space + if preprocessor: + self._observation_space = preprocessor.observation_space + else: + self._observation_space = external_env.observation_space + external_env.start() + + @override(BaseEnv) + def poll( + self, + ) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, MultiEnvDict, MultiEnvDict]: + with self.external_env._results_avail_condition: + results = self._poll() + while len(results[0]) == 0: + self.external_env._results_avail_condition.wait() + results = self._poll() + if not self.external_env.is_alive(): + raise Exception("Serving thread has stopped.") + return results + + @override(BaseEnv) + def send_actions(self, action_dict: MultiEnvDict) -> None: + from ray.rllib.env.base_env import _DUMMY_AGENT_ID + + if self.multiagent: + for env_id, actions in action_dict.items(): + self.external_env._episodes[env_id].action_queue.put(actions) + else: + for env_id, action in action_dict.items(): + self.external_env._episodes[env_id].action_queue.put( + action[_DUMMY_AGENT_ID] + ) + + def _poll( + self, + ) -> Tuple[ + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + ]: + from ray.rllib.env.base_env import with_dummy_agent_id + + all_obs, all_rewards, all_terminateds, all_truncateds, all_infos = ( + {}, + {}, + {}, + {}, + {}, + ) + off_policy_actions = {} + for eid, episode in self.external_env._episodes.copy().items(): + data = episode.get_data() + cur_terminated = ( + episode.cur_terminated_dict["__all__"] + if self.multiagent + else episode.cur_terminated + ) + cur_truncated = ( + episode.cur_truncated_dict["__all__"] + if self.multiagent + else episode.cur_truncated + ) + if cur_terminated or cur_truncated: + del self.external_env._episodes[eid] + if data: + if self.prep: + all_obs[eid] = self.prep.transform(data["obs"]) + else: + all_obs[eid] = data["obs"] + all_rewards[eid] = data["reward"] + all_terminateds[eid] = data["terminated"] + all_truncateds[eid] = data["truncated"] + all_infos[eid] = data["info"] + if "off_policy_action" in data: + off_policy_actions[eid] = data["off_policy_action"] + if self.multiagent: + # Ensure a consistent set of keys + # rely on all_obs having all possible keys for now. + for eid, eid_dict in all_obs.items(): + for agent_id in eid_dict.keys(): + + def fix(d, zero_val): + if agent_id not in d[eid]: + d[eid][agent_id] = zero_val + + fix(all_rewards, 0.0) + fix(all_terminateds, False) + fix(all_truncateds, False) + fix(all_infos, {}) + return ( + all_obs, + all_rewards, + all_terminateds, + all_truncateds, + all_infos, + off_policy_actions, + ) + else: + return ( + with_dummy_agent_id(all_obs), + with_dummy_agent_id(all_rewards), + with_dummy_agent_id(all_terminateds, "__all__"), + with_dummy_agent_id(all_truncateds, "__all__"), + with_dummy_agent_id(all_infos), + with_dummy_agent_id(off_policy_actions), + ) + + @property + @override(BaseEnv) + def observation_space(self) -> gym.spaces.Dict: + return self._observation_space + + @property + @override(BaseEnv) + def action_space(self) -> gym.Space: + return self._action_space diff --git a/lib/python3.12/site-packages/ray/rllib/env/external_multi_agent_env.py b/lib/python3.12/site-packages/ray/rllib/env/external_multi_agent_env.py new file mode 100644 index 0000000000000000000000000000000000000000..6c7e333700f605f4e305c19e7c1f406ec8e607e2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/external_multi_agent_env.py @@ -0,0 +1,162 @@ +import uuid +from typing import Optional + +import gymnasium as gym + +from ray.rllib.env.external_env import ExternalEnv, _ExternalEnvEpisode +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.typing import MultiAgentDict + + +@OldAPIStack +class ExternalMultiAgentEnv(ExternalEnv): + """This is the multi-agent version of ExternalEnv.""" + + def __init__( + self, + action_space: gym.Space, + observation_space: gym.Space, + ): + """Initializes an ExternalMultiAgentEnv instance. + + Args: + action_space: Action space of the env. + observation_space: Observation space of the env. + """ + ExternalEnv.__init__(self, action_space, observation_space) + + # We require to know all agents' spaces. + if isinstance(self.action_space, dict) or isinstance( + self.observation_space, dict + ): + if not (self.action_space.keys() == self.observation_space.keys()): + raise ValueError( + "Agent ids disagree for action space and obs " + "space dict: {} {}".format( + self.action_space.keys(), self.observation_space.keys() + ) + ) + + def run(self): + """Override this to implement the multi-agent run loop. + + Your loop should continuously: + 1. Call self.start_episode(episode_id) + 2. Call self.get_action(episode_id, obs_dict) + -or- + self.log_action(episode_id, obs_dict, action_dict) + 3. Call self.log_returns(episode_id, reward_dict) + 4. Call self.end_episode(episode_id, obs_dict) + 5. Wait if nothing to do. + + Multiple episodes may be started at the same time. + """ + raise NotImplementedError + + @override(ExternalEnv) + def start_episode( + self, episode_id: Optional[str] = None, training_enabled: bool = True + ) -> str: + if episode_id is None: + episode_id = uuid.uuid4().hex + + if episode_id in self._finished: + raise ValueError("Episode {} has already completed.".format(episode_id)) + + if episode_id in self._episodes: + raise ValueError("Episode {} is already started".format(episode_id)) + + self._episodes[episode_id] = _ExternalEnvEpisode( + episode_id, self._results_avail_condition, training_enabled, multiagent=True + ) + + return episode_id + + @override(ExternalEnv) + def get_action( + self, episode_id: str, observation_dict: MultiAgentDict + ) -> MultiAgentDict: + """Record an observation and get the on-policy action. + + Thereby, observation_dict is expected to contain the observation + of all agents acting in this episode step. + + Args: + episode_id: Episode id returned from start_episode(). + observation_dict: Current environment observation. + + Returns: + action: Action from the env action space. + """ + + episode = self._get(episode_id) + return episode.wait_for_action(observation_dict) + + @override(ExternalEnv) + def log_action( + self, + episode_id: str, + observation_dict: MultiAgentDict, + action_dict: MultiAgentDict, + ) -> None: + """Record an observation and (off-policy) action taken. + + Args: + episode_id: Episode id returned from start_episode(). + observation_dict: Current environment observation. + action_dict: Action for the observation. + """ + + episode = self._get(episode_id) + episode.log_action(observation_dict, action_dict) + + @override(ExternalEnv) + def log_returns( + self, + episode_id: str, + reward_dict: MultiAgentDict, + info_dict: MultiAgentDict = None, + multiagent_done_dict: MultiAgentDict = None, + ) -> None: + """Record returns from the environment. + + The reward will be attributed to the previous action taken by the + episode. Rewards accumulate until the next action. If no reward is + logged before the next action, a reward of 0.0 is assumed. + + Args: + episode_id: Episode id returned from start_episode(). + reward_dict: Reward from the environment agents. + info_dict: Optional info dict. + multiagent_done_dict: Optional done dict for agents. + """ + + episode = self._get(episode_id) + + # Accumulate reward by agent. + # For existing agents, we want to add the reward up. + for agent, rew in reward_dict.items(): + if agent in episode.cur_reward_dict: + episode.cur_reward_dict[agent] += rew + else: + episode.cur_reward_dict[agent] = rew + + if multiagent_done_dict: + for agent, done in multiagent_done_dict.items(): + episode.cur_done_dict[agent] = done + + if info_dict: + episode.cur_info_dict = info_dict or {} + + @override(ExternalEnv) + def end_episode(self, episode_id: str, observation_dict: MultiAgentDict) -> None: + """Record the end of an episode. + + Args: + episode_id: Episode id returned from start_episode(). + observation_dict: Current environment observation. + """ + + episode = self._get(episode_id) + self._finished.add(episode.episode_id) + episode.done(observation_dict) diff --git a/lib/python3.12/site-packages/ray/rllib/env/multi_agent_env.py b/lib/python3.12/site-packages/ray/rllib/env/multi_agent_env.py new file mode 100644 index 0000000000000000000000000000000000000000..5e2a2de964356f26f00ec17e3fc5a13b8115f9eb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/multi_agent_env.py @@ -0,0 +1,807 @@ +import copy +import logging +from typing import Callable, Dict, List, Optional, Set, Tuple, Type, Union + +import gymnasium as gym +import numpy as np + +from ray._common.deprecation import Deprecated +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.env.env_context import EnvContext +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.typing import ( + AgentID, + EnvCreator, + EnvID, + EnvType, + MultiAgentDict, + MultiEnvDict, +) +from ray.util import log_once +from ray.util.annotations import DeveloperAPI, PublicAPI + +# If the obs space is Dict type, look for the global state under this key. +ENV_STATE = "state" + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="beta") +class MultiAgentEnv(gym.Env): + """An environment that hosts multiple independent agents. + + Agents are identified by AgentIDs (string). + """ + + # Optional mappings from AgentID to individual agents' spaces. + # Set this to an "exhaustive" dictionary, mapping all possible AgentIDs to + # individual agents' spaces. Alternatively, override + # `get_observation_space(agent_id=...)` and `get_action_space(agent_id=...)`, which + # is the API that RLlib uses to get individual spaces and whose default + # implementation is to simply look up `agent_id` in these dicts. + observation_spaces: Optional[Dict[AgentID, gym.Space]] = None + action_spaces: Optional[Dict[AgentID, gym.Space]] = None + + # All agents currently active in the environment. This attribute may change during + # the lifetime of the env or even during an individual episode. + agents: List[AgentID] = [] + # All agents that may appear in the environment, ever. + # This attribute should not be changed during the lifetime of this env. + possible_agents: List[AgentID] = [] + + # @OldAPIStack, use `observation_spaces` and `action_spaces`, instead. + observation_space: Optional[gym.Space] = None + action_space: Optional[gym.Space] = None + + def __init__(self): + super().__init__() + + # @OldAPIStack + if not hasattr(self, "_agent_ids"): + self._agent_ids = set() + + # If these important attributes are not set, try to infer them. + if not self.agents: + self.agents = list(self._agent_ids) + if not self.possible_agents: + self.possible_agents = self.agents.copy() + + def reset( + self, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Tuple[MultiAgentDict, MultiAgentDict]: # type: ignore + """Resets the env and returns observations from ready agents. + + Args: + seed: An optional seed to use for the new episode. + + Returns: + New observations for each ready agent. + + .. testcode:: + :skipif: True + + from ray.rllib.env.multi_agent_env import MultiAgentEnv + class MyMultiAgentEnv(MultiAgentEnv): + # Define your env here. + env = MyMultiAgentEnv() + obs, infos = env.reset(seed=42, options={}) + print(obs) + + .. testoutput:: + + { + "car_0": [2.4, 1.6], + "car_1": [3.4, -3.2], + "traffic_light_1": [0, 3, 5, 1], + } + """ + # Call super's `reset()` method to (maybe) set the given `seed`. + super().reset(seed=seed, options=options) + + def step( + self, action_dict: MultiAgentDict + ) -> Tuple[ + MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict + ]: + """Returns observations from ready agents. + + The returns are dicts mapping from agent_id strings to values. The + number of agents in the env can vary over time. + + Returns: + Tuple containing 1) new observations for + each ready agent, 2) reward values for each ready agent. If + the episode is just started, the value will be None. + 3) Terminated values for each ready agent. The special key + "__all__" (required) is used to indicate env termination. + 4) Truncated values for each ready agent. + 5) Info values for each agent id (may be empty dicts). + + .. testcode:: + :skipif: True + + env = ... + obs, rewards, terminateds, truncateds, infos = env.step(action_dict={ + "car_0": 1, "car_1": 0, "traffic_light_1": 2, + }) + print(rewards) + + print(terminateds) + + print(infos) + + .. testoutput:: + + { + "car_0": 3, + "car_1": -1, + "traffic_light_1": 0, + } + { + "car_0": False, # car_0 is still running + "car_1": True, # car_1 is terminated + "__all__": False, # the env is not terminated + } + { + "car_0": {}, # info for car_0 + "car_1": {}, # info for car_1 + } + + """ + raise NotImplementedError + + def render(self) -> None: + """Tries to render the environment.""" + + # By default, do nothing. + pass + + def get_observation_space(self, agent_id: AgentID) -> gym.Space: + if self.observation_spaces is not None: + return self.observation_spaces[agent_id] + + # @OldAPIStack behavior. + # `self.observation_space` is a `gym.spaces.Dict` AND contains `agent_id`. + if ( + isinstance(self.observation_space, gym.spaces.Dict) + and agent_id in self.observation_space.spaces + ): + return self.observation_space[agent_id] + # `self.observation_space` is not a `gym.spaces.Dict` OR doesn't contain + # `agent_id` -> The defined space is most likely meant to be the space + # for all agents. + else: + return self.observation_space + + def get_action_space(self, agent_id: AgentID) -> gym.Space: + if self.action_spaces is not None: + return self.action_spaces[agent_id] + + # @OldAPIStack behavior. + # `self.action_space` is a `gym.spaces.Dict` AND contains `agent_id`. + if ( + isinstance(self.action_space, gym.spaces.Dict) + and agent_id in self.action_space.spaces + ): + return self.action_space[agent_id] + # `self.action_space` is not a `gym.spaces.Dict` OR doesn't contain + # `agent_id` -> The defined space is most likely meant to be the space + # for all agents. + else: + return self.action_space + + @property + def num_agents(self) -> int: + return len(self.agents) + + @property + def max_num_agents(self) -> int: + return len(self.possible_agents) + + # fmt: off + # __grouping_doc_begin__ + def with_agent_groups( + self, + groups: Dict[str, List[AgentID]], + obs_space: gym.Space = None, + act_space: gym.Space = None, + ) -> "MultiAgentEnv": + """Convenience method for grouping together agents in this env. + + An agent group is a list of agent IDs that are mapped to a single + logical agent. All agents of the group must act at the same time in the + environment. The grouped agent exposes Tuple action and observation + spaces that are the concatenated action and obs spaces of the + individual agents. + + The rewards of all the agents in a group are summed. The individual + agent rewards are available under the "individual_rewards" key of the + group info return. + + Agent grouping is required to leverage algorithms such as Q-Mix. + + Args: + groups: Mapping from group id to a list of the agent ids + of group members. If an agent id is not present in any group + value, it will be left ungrouped. The group id becomes a new agent ID + in the final environment. + obs_space: Optional observation space for the grouped + env. Must be a tuple space. If not provided, will infer this to be a + Tuple of n individual agents spaces (n=num agents in a group). + act_space: Optional action space for the grouped env. + Must be a tuple space. If not provided, will infer this to be a Tuple + of n individual agents spaces (n=num agents in a group). + + .. testcode:: + :skipif: True + + from ray.rllib.env.multi_agent_env import MultiAgentEnv + class MyMultiAgentEnv(MultiAgentEnv): + # define your env here + ... + env = MyMultiAgentEnv(...) + grouped_env = env.with_agent_groups(env, { + "group1": ["agent1", "agent2", "agent3"], + "group2": ["agent4", "agent5"], + }) + + """ + + from ray.rllib.env.wrappers.group_agents_wrapper import GroupAgentsWrapper + return GroupAgentsWrapper(self, groups, obs_space, act_space) + + # __grouping_doc_end__ + # fmt: on + + @OldAPIStack + @Deprecated(new="MultiAgentEnv.possible_agents", error=False) + def get_agent_ids(self) -> Set[AgentID]: + if not hasattr(self, "_agent_ids"): + self._agent_ids = set() + if not isinstance(self._agent_ids, set): + self._agent_ids = set(self._agent_ids) + # Make this backward compatible as much as possible. + return self._agent_ids if self._agent_ids else set(self.agents) + + @OldAPIStack + def to_base_env( + self, + make_env: Optional[Callable[[int], EnvType]] = None, + num_envs: int = 1, + remote_envs: bool = False, + remote_env_batch_wait_ms: int = 0, + restart_failed_sub_environments: bool = False, + ) -> "BaseEnv": + """Converts an RLlib MultiAgentEnv into a BaseEnv object. + + The resulting BaseEnv is always vectorized (contains n + sub-environments) to support batched forward passes, where n may + also be 1. BaseEnv also supports async execution via the `poll` and + `send_actions` methods and thus supports external simulators. + + Args: + make_env: A callable taking an int as input (which indicates + the number of individual sub-environments within the final + vectorized BaseEnv) and returning one individual + sub-environment. + num_envs: The number of sub-environments to create in the + resulting (vectorized) BaseEnv. The already existing `env` + will be one of the `num_envs`. + remote_envs: Whether each sub-env should be a @ray.remote + actor. You can set this behavior in your config via the + `remote_worker_envs=True` option. + remote_env_batch_wait_ms: The wait time (in ms) to poll remote + sub-environments for, if applicable. Only used if + `remote_envs` is True. + restart_failed_sub_environments: If True and any sub-environment (within + a vectorized env) throws any error during env stepping, we will try to + restart the faulty sub-environment. This is done + without disturbing the other (still intact) sub-environments. + + Returns: + The resulting BaseEnv object. + """ + from ray.rllib.env.remote_base_env import RemoteBaseEnv + + if remote_envs: + env = RemoteBaseEnv( + make_env, + num_envs, + multiagent=True, + remote_env_batch_wait_ms=remote_env_batch_wait_ms, + restart_failed_sub_environments=restart_failed_sub_environments, + ) + # Sub-environments are not ray.remote actors. + else: + env = MultiAgentEnvWrapper( + make_env=make_env, + existing_envs=[self], + num_envs=num_envs, + restart_failed_sub_environments=restart_failed_sub_environments, + ) + + return env + + +@DeveloperAPI +def make_multi_agent( + env_name_or_creator: Union[str, EnvCreator], +) -> Type["MultiAgentEnv"]: + """Convenience wrapper for any single-agent env to be converted into MA. + + Allows you to convert a simple (single-agent) `gym.Env` class + into a `MultiAgentEnv` class. This function simply stacks n instances + of the given ```gym.Env``` class into one unified ``MultiAgentEnv`` class + and returns this class, thus pretending the agents act together in the + same environment, whereas - under the hood - they live separately from + each other in n parallel single-agent envs. + + Agent IDs in the resulting and are int numbers starting from 0 + (first agent). + + Args: + env_name_or_creator: String specifier or env_maker function taking + an EnvContext object as only arg and returning a gym.Env. + + Returns: + New MultiAgentEnv class to be used as env. + The constructor takes a config dict with `num_agents` key + (default=1). The rest of the config dict will be passed on to the + underlying single-agent env's constructor. + + .. testcode:: + :skipif: True + + from ray.rllib.env.multi_agent_env import make_multi_agent + # By gym string: + ma_cartpole_cls = make_multi_agent("CartPole-v1") + # Create a 2 agent multi-agent cartpole. + ma_cartpole = ma_cartpole_cls({"num_agents": 2}) + obs = ma_cartpole.reset() + print(obs) + + # By env-maker callable: + from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole + ma_stateless_cartpole_cls = make_multi_agent( + lambda config: StatelessCartPole(config)) + # Create a 3 agent multi-agent stateless cartpole. + ma_stateless_cartpole = ma_stateless_cartpole_cls( + {"num_agents": 3}) + print(obs) + + .. testoutput:: + + {0: [...], 1: [...]} + {0: [...], 1: [...], 2: [...]} + """ + + class MultiEnv(MultiAgentEnv): + def __init__(self, config: EnvContext = None): + super().__init__() + + # Note: Explicitly check for None here, because config + # can have an empty dict but meaningful data fields (worker_index, + # vector_index) etc. + # TODO (sven): Clean this up, so we are not mixing up dict fields + # with data fields. + if config is None: + config = {} + else: + # Note the deepcopy is needed b/c (a) we need to remove the + # `num_agents` keyword and (b) with `num_envs > 0` in the + # `VectorMultiAgentEnv` all following environment creations + # need the same config again. + config = copy.deepcopy(config) + num = config.pop("num_agents", 1) + if isinstance(env_name_or_creator, str): + self.envs = [gym.make(env_name_or_creator) for _ in range(num)] + else: + self.envs = [env_name_or_creator(config) for _ in range(num)] + self.terminateds = set() + self.truncateds = set() + self.observation_spaces = { + i: self.envs[i].observation_space for i in range(num) + } + self.action_spaces = {i: self.envs[i].action_space for i in range(num)} + self.agents = list(range(num)) + self.possible_agents = self.agents.copy() + + @override(MultiAgentEnv) + def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): + self.terminateds = set() + self.truncateds = set() + obs, infos = {}, {} + for i, env in enumerate(self.envs): + obs[i], infos[i] = env.reset(seed=seed, options=options) + if not self.observation_spaces[i].contains(obs[i]): + print("===> MultiEnv does not contain obs.") + + return obs, infos + + @override(MultiAgentEnv) + def step(self, action_dict): + obs, rew, terminated, truncated, info = {}, {}, {}, {}, {} + + # The environment is expecting an action for at least one agent. + if len(action_dict) == 0: + raise ValueError( + "The environment is expecting an action for at least one agent." + ) + + for i, action in action_dict.items(): + obs[i], rew[i], terminated[i], truncated[i], info[i] = self.envs[ + i + ].step(action) + if terminated[i]: + self.terminateds.add(i) + if truncated[i]: + self.truncateds.add(i) + # TODO: Flaw in our MultiAgentEnv API wrt. new gymnasium: Need to return + # an additional episode_done bool that covers cases where all agents are + # either terminated or truncated, but not all are truncated and not all are + # terminated. We can then get rid of the aweful `__all__` special keys! + terminated["__all__"] = len(self.terminateds | self.truncateds) == len( + self.envs + ) + truncated["__all__"] = len(self.truncateds) == len(self.envs) + return obs, rew, terminated, truncated, info + + @override(MultiAgentEnv) + def render(self): + # This render method simply renders all n underlying individual single-agent + # envs and concatenates their images (on top of each other if the returned + # images have dims where [width] > [height], otherwise next to each other). + render_images = [e.render() for e in self.envs] + if render_images[0].shape[1] > render_images[0].shape[0]: + concat_dim = 0 + else: + concat_dim = 1 + return np.concatenate(render_images, axis=concat_dim) + + return MultiEnv + + +@OldAPIStack +class MultiAgentEnvWrapper(BaseEnv): + """Internal adapter of MultiAgentEnv to BaseEnv. + + This also supports vectorization if num_envs > 1. + """ + + def __init__( + self, + make_env: Callable[[int], EnvType], + existing_envs: List["MultiAgentEnv"], + num_envs: int, + restart_failed_sub_environments: bool = False, + ): + """Wraps MultiAgentEnv(s) into the BaseEnv API. + + Args: + make_env: Factory that produces a new MultiAgentEnv instance taking the + vector index as only call argument. + Must be defined, if the number of existing envs is less than num_envs. + existing_envs: List of already existing multi-agent envs. + num_envs: Desired num multiagent envs to have at the end in + total. This will include the given (already created) + `existing_envs`. + restart_failed_sub_environments: If True and any sub-environment (within + this vectorized env) throws any error during env stepping, we will try + to restart the faulty sub-environment. This is done + without disturbing the other (still intact) sub-environments. + """ + self.make_env = make_env + self.envs = existing_envs + self.num_envs = num_envs + self.restart_failed_sub_environments = restart_failed_sub_environments + + self.terminateds = set() + self.truncateds = set() + while len(self.envs) < self.num_envs: + self.envs.append(self.make_env(len(self.envs))) + for env in self.envs: + assert isinstance(env, MultiAgentEnv) + self._init_env_state(idx=None) + self._unwrapped_env = self.envs[0].unwrapped + + @override(BaseEnv) + def poll( + self, + ) -> Tuple[ + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + ]: + obs, rewards, terminateds, truncateds, infos = {}, {}, {}, {}, {} + for i, env_state in enumerate(self.env_states): + ( + obs[i], + rewards[i], + terminateds[i], + truncateds[i], + infos[i], + ) = env_state.poll() + return obs, rewards, terminateds, truncateds, infos, {} + + @override(BaseEnv) + def send_actions(self, action_dict: MultiEnvDict) -> None: + for env_id, agent_dict in action_dict.items(): + if env_id in self.terminateds or env_id in self.truncateds: + raise ValueError( + f"Env {env_id} is already done and cannot accept new actions" + ) + env = self.envs[env_id] + try: + obs, rewards, terminateds, truncateds, infos = env.step(agent_dict) + except Exception as e: + if self.restart_failed_sub_environments: + logger.exception(e.args[0]) + self.try_restart(env_id=env_id) + obs = e + rewards = {} + terminateds = {"__all__": True} + truncateds = {"__all__": False} + infos = {} + else: + raise e + + assert isinstance( + obs, (dict, Exception) + ), "Not a multi-agent obs dict or an Exception!" + assert isinstance(rewards, dict), "Not a multi-agent reward dict!" + assert isinstance(terminateds, dict), "Not a multi-agent terminateds dict!" + assert isinstance(truncateds, dict), "Not a multi-agent truncateds dict!" + assert isinstance(infos, dict), "Not a multi-agent info dict!" + if isinstance(obs, dict): + info_diff = set(infos).difference(set(obs)) + if info_diff and info_diff != {"__common__"}: + raise ValueError( + "Key set for infos must be a subset of obs (plus optionally " + "the '__common__' key for infos concerning all/no agents): " + "{} vs {}".format(infos.keys(), obs.keys()) + ) + if "__all__" not in terminateds: + raise ValueError( + "In multi-agent environments, '__all__': True|False must " + "be included in the 'terminateds' dict: got {}.".format(terminateds) + ) + elif "__all__" not in truncateds: + raise ValueError( + "In multi-agent environments, '__all__': True|False must " + "be included in the 'truncateds' dict: got {}.".format(truncateds) + ) + + if terminateds["__all__"]: + self.terminateds.add(env_id) + if truncateds["__all__"]: + self.truncateds.add(env_id) + self.env_states[env_id].observe( + obs, rewards, terminateds, truncateds, infos + ) + + @override(BaseEnv) + def try_reset( + self, + env_id: Optional[EnvID] = None, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Optional[Tuple[MultiEnvDict, MultiEnvDict]]: + ret_obs = {} + ret_infos = {} + if isinstance(env_id, int): + env_id = [env_id] + if env_id is None: + env_id = list(range(len(self.envs))) + for idx in env_id: + obs, infos = self.env_states[idx].reset(seed=seed, options=options) + + if isinstance(obs, Exception): + if self.restart_failed_sub_environments: + self.env_states[idx].env = self.envs[idx] = self.make_env(idx) + else: + raise obs + else: + assert isinstance(obs, dict), "Not a multi-agent obs dict!" + if obs is not None: + if idx in self.terminateds: + self.terminateds.remove(idx) + if idx in self.truncateds: + self.truncateds.remove(idx) + ret_obs[idx] = obs + ret_infos[idx] = infos + return ret_obs, ret_infos + + @override(BaseEnv) + def try_restart(self, env_id: Optional[EnvID] = None) -> None: + if isinstance(env_id, int): + env_id = [env_id] + if env_id is None: + env_id = list(range(len(self.envs))) + for idx in env_id: + # Try closing down the old (possibly faulty) sub-env, but ignore errors. + try: + self.envs[idx].close() + except Exception as e: + if log_once("close_sub_env"): + logger.warning( + "Trying to close old and replaced sub-environment (at vector " + f"index={idx}), but closing resulted in error:\n{e}" + ) + # Try recreating the sub-env. + logger.warning(f"Trying to restart sub-environment at index {idx}.") + self.env_states[idx].env = self.envs[idx] = self.make_env(idx) + logger.warning(f"Sub-environment at index {idx} restarted successfully.") + + @override(BaseEnv) + def get_sub_environments( + self, as_dict: bool = False + ) -> Union[Dict[str, EnvType], List[EnvType]]: + if as_dict: + return {_id: env_state.env for _id, env_state in enumerate(self.env_states)} + return [state.env for state in self.env_states] + + @override(BaseEnv) + def try_render(self, env_id: Optional[EnvID] = None) -> None: + if env_id is None: + env_id = 0 + assert isinstance(env_id, int) + return self.envs[env_id].render() + + @property + @override(BaseEnv) + def observation_space(self) -> gym.spaces.Dict: + return self.envs[0].observation_space + + @property + @override(BaseEnv) + def action_space(self) -> gym.Space: + return self.envs[0].action_space + + @override(BaseEnv) + def get_agent_ids(self) -> Set[AgentID]: + return self.envs[0].get_agent_ids() + + def _init_env_state(self, idx: Optional[int] = None) -> None: + """Resets all or one particular sub-environment's state (by index). + + Args: + idx: The index to reset at. If None, reset all the sub-environments' states. + """ + # If index is None, reset all sub-envs' states: + if idx is None: + self.env_states = [ + _MultiAgentEnvState(env, self.restart_failed_sub_environments) + for env in self.envs + ] + # Index provided, reset only the sub-env's state at the given index. + else: + assert isinstance(idx, int) + self.env_states[idx] = _MultiAgentEnvState( + self.envs[idx], self.restart_failed_sub_environments + ) + + +@OldAPIStack +class _MultiAgentEnvState: + def __init__(self, env: MultiAgentEnv, return_error_as_obs: bool = False): + assert isinstance(env, MultiAgentEnv) + self.env = env + self.return_error_as_obs = return_error_as_obs + + self.initialized = False + self.last_obs = {} + self.last_rewards = {} + self.last_terminateds = {"__all__": False} + self.last_truncateds = {"__all__": False} + self.last_infos = {} + + def poll( + self, + ) -> Tuple[ + MultiAgentDict, + MultiAgentDict, + MultiAgentDict, + MultiAgentDict, + MultiAgentDict, + ]: + if not self.initialized: + # TODO(sven): Should we make it possible to pass in a seed here? + self.reset() + self.initialized = True + + observations = self.last_obs + rewards = {} + terminateds = {"__all__": self.last_terminateds["__all__"]} + truncateds = {"__all__": self.last_truncateds["__all__"]} + infos = self.last_infos + + # If episode is done or we have an error, release everything we have. + if ( + terminateds["__all__"] + or truncateds["__all__"] + or isinstance(observations, Exception) + ): + rewards = self.last_rewards + self.last_rewards = {} + terminateds = self.last_terminateds + if isinstance(observations, Exception): + terminateds["__all__"] = True + truncateds["__all__"] = False + self.last_terminateds = {} + truncateds = self.last_truncateds + self.last_truncateds = {} + self.last_obs = {} + infos = self.last_infos + self.last_infos = {} + # Only release those agents' rewards/terminateds/truncateds/infos, whose + # observations we have. + else: + for ag in observations.keys(): + if ag in self.last_rewards: + rewards[ag] = self.last_rewards[ag] + del self.last_rewards[ag] + if ag in self.last_terminateds: + terminateds[ag] = self.last_terminateds[ag] + del self.last_terminateds[ag] + if ag in self.last_truncateds: + truncateds[ag] = self.last_truncateds[ag] + del self.last_truncateds[ag] + + self.last_terminateds["__all__"] = False + self.last_truncateds["__all__"] = False + return observations, rewards, terminateds, truncateds, infos + + def observe( + self, + obs: MultiAgentDict, + rewards: MultiAgentDict, + terminateds: MultiAgentDict, + truncateds: MultiAgentDict, + infos: MultiAgentDict, + ): + self.last_obs = obs + for ag, r in rewards.items(): + if ag in self.last_rewards: + self.last_rewards[ag] += r + else: + self.last_rewards[ag] = r + for ag, d in terminateds.items(): + if ag in self.last_terminateds: + self.last_terminateds[ag] = self.last_terminateds[ag] or d + else: + self.last_terminateds[ag] = d + for ag, t in truncateds.items(): + if ag in self.last_truncateds: + self.last_truncateds[ag] = self.last_truncateds[ag] or t + else: + self.last_truncateds[ag] = t + self.last_infos = infos + + def reset( + self, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Tuple[MultiAgentDict, MultiAgentDict]: + try: + obs_and_infos = self.env.reset(seed=seed, options=options) + except Exception as e: + if self.return_error_as_obs: + logger.exception(e.args[0]) + obs_and_infos = e, e + else: + raise e + + self.last_obs, self.last_infos = obs_and_infos + self.last_rewards = {} + self.last_terminateds = {"__all__": False} + self.last_truncateds = {"__all__": False} + + return self.last_obs, self.last_infos diff --git a/lib/python3.12/site-packages/ray/rllib/env/multi_agent_env_runner.py b/lib/python3.12/site-packages/ray/rllib/env/multi_agent_env_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..f83733471fb116f1b55bcc8a252f620dacc9358b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/multi_agent_env_runner.py @@ -0,0 +1,1102 @@ +import logging +import math +import time +from collections import defaultdict +from functools import partial +from typing import Collection, DefaultDict, Dict, List, Optional, Union + +import gymnasium as gym + +import ray +from ray._common.deprecation import Deprecated +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.callbacks.utils import make_callback +from ray.rllib.core import ( + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + COMPONENT_RL_MODULE, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule, MultiRLModuleSpec +from ray.rllib.env import INPUT_ENV_SINGLE_SPACES, INPUT_ENV_SPACES +from ray.rllib.env.env_context import EnvContext +from ray.rllib.env.env_runner import ENV_STEP_FAILURE, EnvRunner +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.env.utils import _gym_env_creator +from ray.rllib.env.vector.registration import make_vec +from ray.rllib.env.vector.vector_multi_agent_env import VectorMultiAgentEnv +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.framework import get_device, try_import_torch +from ray.rllib.utils.metrics import ( + ENV_TO_MODULE_CONNECTOR, + EPISODE_DURATION_SEC_MEAN, + EPISODE_LEN_MAX, + EPISODE_LEN_MEAN, + EPISODE_LEN_MIN, + EPISODE_RETURN_MAX, + EPISODE_RETURN_MEAN, + EPISODE_RETURN_MIN, + MODULE_TO_ENV_CONNECTOR, + NUM_AGENT_STEPS_SAMPLED, + NUM_AGENT_STEPS_SAMPLED_LIFETIME, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_EPISODES, + NUM_EPISODES_LIFETIME, + NUM_MODULE_STEPS_SAMPLED, + NUM_MODULE_STEPS_SAMPLED_LIFETIME, + RLMODULE_INFERENCE_TIMER, + SAMPLE_TIMER, + TIME_BETWEEN_SAMPLING, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.pre_checks.env import check_multiagent_environments +from ray.rllib.utils.typing import EpisodeID, ModelWeights, ResultDict, StateDict +from ray.tune.registry import ENV_CREATOR, _global_registry +from ray.util.annotations import PublicAPI + +torch, _ = try_import_torch() +logger = logging.getLogger("ray.rllib") + + +# TODO (sven): As soon as RolloutWorker is no longer supported, make `EnvRunner` itself +# a Checkpointable. Currently, only some of its subclasses are Checkpointables. +@PublicAPI(stability="alpha") +class MultiAgentEnvRunner(EnvRunner, Checkpointable): + """The genetic environment runner for the multi-agent case.""" + + @override(EnvRunner) + def __init__(self, config: AlgorithmConfig, **kwargs): + """Initializes a MultiAgentEnvRunner instance. + + Args: + config: An `AlgorithmConfig` object containing all settings needed to + build this `EnvRunner` class. + """ + super().__init__(config=config, **kwargs) + + # Raise an Error, if the provided config is not a multi-agent one. + if not self.config.is_multi_agent: + raise ValueError( + f"Cannot use this EnvRunner class ({type(self).__name__}), if your " + "setup is not multi-agent! Try adding multi-agent information to your " + "AlgorithmConfig via calling the `config.multi_agent(policies=..., " + "policy_mapping_fn=...)`." + ) + + self.spaces = kwargs.get("spaces", {}) + + self._setup_metrics() + + # Create our callbacks object. + self._callbacks = [cls() for cls in force_list(self.config.callbacks_class)] + + # Set device. + self._device = get_device( + self.config, + 0 if not self.worker_index else self.config.num_gpus_per_env_runner, + ) + + # Create the vectorized gymnasium env. + self.env: Optional[VectorMultiAgentEnv] = None + self.num_envs: int = 0 + if ( + self.worker_index is None + or self.worker_index > 0 + or self.config.create_env_on_local_worker + or self.config.num_env_runners == 0 + ): + self.make_env() + + # Create the env-to-module connector pipeline. + self._env_to_module = self.config.build_env_to_module_connector( + env=self.env, spaces=self.spaces, device=self._device + ) + # Cached env-to-module results taken at the end of a `_sample_timesteps()` + # call to make sure the final observation (before an episode cut) gets properly + # processed (and maybe postprocessed and re-stored into the episode). + # For example, if we had a connector that normalizes observations and directly + # re-inserts these new obs back into the episode, the last observation in each + # sample call would NOT be processed, which could be very harmful in cases, + # in which value function bootstrapping of those (truncation) observations is + # required in the learning step. + self._cached_to_module = None + + # Construct the MultiRLModule. + self.module: Optional[MultiRLModule] = None + self.make_module() + + # Create the module-to-env connector pipeline. + self._module_to_env = self.config.build_module_to_env_connector( + env=self.env.unwrapped if self.env else None, spaces=self.spaces + ) + + self._needs_initial_reset: bool = True + self._episode: Optional[MultiAgentEpisode] = None + self._shared_data = None + + self._weights_seq_no: int = 0 + + # Measures the time passed between returning from `sample()` + # and receiving the next `sample()` request from the user. + self._time_after_sampling = None + + @override(EnvRunner) + def sample( + self, + *, + num_timesteps: int = None, + num_episodes: int = None, + explore: bool = None, + random_actions: bool = False, + force_reset: bool = False, + ) -> List[MultiAgentEpisode]: + """Runs and returns a sample (n timesteps or m episodes) on the env(s). + + Args: + num_timesteps: The number of timesteps to sample during this call. + Note that only one of `num_timesteps` or `num_episodes` may be provided. + num_episodes: The number of episodes to sample during this call. + Note that only one of `num_timesteps` or `num_episodes` may be provided. + explore: If True, will use the RLModule's `forward_exploration()` + method to compute actions. If False, will use the RLModule's + `forward_inference()` method. If None (default), will use the `explore` + boolean setting from `self.config` passed into this EnvRunner's + constructor. You can change this setting in your config via + `config.env_runners(explore=True|False)`. + random_actions: If True, actions will be sampled randomly (from the action + space of the environment). If False (default), actions or action + distribution parameters are computed by the RLModule. + force_reset: Whether to force-reset all (vector) environments before + sampling. Useful if you would like to collect a clean slate of new + episodes via this call. Note that when sampling n episodes + (`num_episodes != None`), this is fixed to True. + + Returns: + A list of `MultiAgentEpisode` instances, carrying the sampled data. + """ + if self.env is None: + raise ValueError( + f"{self} doesn't have an env! Can't call `sample()` on it." + ) + + assert not (num_timesteps is not None and num_episodes is not None), ( + "Provide " + "either `num_timesteps` or `num_episodes`. Both provided here:" + f"{num_timesteps=}, {num_episodes=}" + ) + # Log time between `sample()` requests. + if self._time_after_sampling is not None: + self.metrics.log_value( + key=TIME_BETWEEN_SAMPLING, + value=time.perf_counter() - self._time_after_sampling, + ) + + # Log current weight seq no. + self.metrics.log_value( + key=WEIGHTS_SEQ_NO, + value=self._weights_seq_no, + window=1, + ) + + with self.metrics.log_time(SAMPLE_TIMER): + # If no execution details are provided, use the config to try to infer the + # desired timesteps/episodes to sample and exploration behavior. + if explore is None: + explore = self.config.explore + if ( + num_timesteps is None + and num_episodes is None + and self.config.batch_mode == "truncate_episodes" + ): + num_timesteps = ( + self.config.get_rollout_fragment_length(self.worker_index) + * self.num_envs + ) + + # Sample "num_timesteps" timesteps. + if num_timesteps is not None: + samples = self._sample( + num_timesteps=num_timesteps, + explore=explore, + random_actions=random_actions, + force_reset=force_reset, + ) + # Sample "num_episodes" episodes. + elif num_episodes is not None: + samples = self._sample( + num_episodes=num_episodes, + explore=explore, + random_actions=random_actions, + ) + # For batch_mode="complete_episodes" (env_runners configuration), continue sampling as long as the number of timesteps done is smaller than the `train_batch_size`. + else: + samples = self._sample( + num_episodes=self.num_envs, + explore=explore, + random_actions=random_actions, + ) + + # Make the `on_sample_end` callback. + make_callback( + "on_sample_end", + callbacks_objects=self._callbacks, + callbacks_functions=self.config.callbacks_on_sample_end, + kwargs=dict( + env_runner=self, + metrics_logger=self.metrics, + samples=samples, + ), + ) + + self._time_after_sampling = time.perf_counter() + + return samples + + def _sample( + self, + *, + num_timesteps: Optional[int] = None, + num_episodes: Optional[int] = None, + explore: bool, + random_actions: bool = False, + force_reset: bool = False, + ) -> List[MultiAgentEpisode]: + + done_episodes_to_return: List[MultiAgentEpisode] = [] + + # Have to reset the env (on all vector sub_envs). + if force_reset or num_episodes is not None or self._needs_initial_reset: + episodes = self._episodes = [None for _ in range(self.num_envs)] + + shared_data = self._shared_data = {} + self._reset_envs(episodes, shared_data, explore) + + # We just reset the env. Don't have to force this again in the next + # call to `self._sample_timesteps()`. + self._needs_initial_reset = False + else: + episodes = self._episodes + shared_data = self._shared_data + + if num_episodes is not None: + self._needs_initial_reset = True + + # Loop through `num_timesteps` timesteps or `num_episodes` episodes. + env_ts = 0 + agent_ts = 0 + eps = 0 + while ( + (eps < num_episodes) + if num_timesteps is None + else ( + env_ts < num_timesteps + if self.config.count_steps_by == "env_steps" + else agent_ts < num_timesteps + ) + ): + # Act randomly. + if random_actions: + to_env = { + Columns.ACTIONS: [ + { + # Only act (randomly) for those agents that had an + # observation. + aid: self.env.envs[i] + .unwrapped.get_action_space(aid) + .sample() + for aid in episodes[i].get_agents_to_act() + } + for i in range(self.num_envs) + ] + } + # Compute an action using the RLModule. + else: + # Env-to-module connector (already cached). + to_module = self._cached_to_module + assert to_module is not None + self._cached_to_module = None + + if to_module: + # MultiRLModule forward pass: Explore or not. + if explore: + # Global env steps sampled are (roughly) this EnvRunner's lifetime + # count times the number of env runners in the algo. + global_env_steps_lifetime = ( + self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME, default=0) + + env_ts + ) * (self.config.num_env_runners or 1) + with self.metrics.log_time(RLMODULE_INFERENCE_TIMER): + to_env = self.module.forward_exploration( + to_module, t=global_env_steps_lifetime + ) + else: + with self.metrics.log_time(RLMODULE_INFERENCE_TIMER): + to_env = self.module.forward_inference(to_module) + + # Module-to-env connector. + to_env = self._module_to_env( + rl_module=self.module, + batch=to_env, + episodes=episodes, + explore=explore, + shared_data=shared_data, + metrics=self.metrics, + metrics_prefix_key=(MODULE_TO_ENV_CONNECTOR,), + ) + # In case all environments had been terminated `to_module` will be + # empty and no actions are needed b/c we reset all environments. + else: + to_env = {} + shared_data["vector_env_episodes_map"] = {} + + # Extract the (vectorized) actions (to be sent to the env) from the + # module/connector output. Note that these actions are fully ready (e.g. + # already unsquashed/clipped) to be sent to the environment and might not + # be identical to the actions produced by the RLModule/distribution, which + # are the ones stored permanently in the episode objects. + actions = to_env.pop(Columns.ACTIONS, [{} for _ in episodes]) + actions_for_env = to_env.pop(Columns.ACTIONS_FOR_ENV, actions) + # Try stepping the environment. + results = self._try_env_step(actions_for_env) + if results == ENV_STEP_FAILURE: + logging.warning( + f"RLlib {self.__class__.__name__}: Environment step failed. Will force reset env(s) in this EnvRunner." + ) + return self._sample( + num_timesteps=num_timesteps, + num_episodes=num_episodes, + explore=explore, + random_actions=random_actions, + force_reset=True, + ) + observations, rewards, terminateds, truncateds, infos = results + + call_on_episode_start = set() + # Store the data from the last environment step into the + # episodes for all sub-environments. + for env_index in range(self.num_envs): + extra_model_outputs = defaultdict(dict) + # `to_env` returns a dictionary with column keys and + # (AgentID, value) tuple values. + for col, ma_dict_list in to_env.items(): + ma_dict = ma_dict_list[env_index] + for agent_id, val in ma_dict.items(): + extra_model_outputs[agent_id][col] = val + extra_model_outputs[agent_id][ + WEIGHTS_SEQ_NO + ] = self._weights_seq_no + extra_model_outputs = dict(extra_model_outputs) + + # Episode has no data in it yet -> Was just reset and needs to be called + # with its `add_env_reset()` method. + if not self._episodes[env_index].is_reset: + # Add the reset step data to the episode. + episodes[env_index].add_env_reset( + observations=observations[env_index], + infos=infos[env_index], + ) + # Call the callback on episode start so users can hook in. + call_on_episode_start.add(env_index) + + # Call `add_env_step()` method on episode. + else: + episodes[env_index].add_env_step( + observations=observations[env_index], + actions=actions[env_index], + rewards=rewards[env_index], + infos=infos[env_index], + terminateds=terminateds[env_index], + truncateds=truncateds[env_index], + extra_model_outputs=extra_model_outputs, + ) + + # Ray metrics + self._log_env_steps( + metric=self._metrics_num_env_steps_sampled, num_steps=1 + ) + + # Only increase ts when we actually stepped (not reset'd as a reset + # does not count as a timestep). + env_ts += self._increase_sampled_metrics( + 1, observations[env_index], episodes[env_index] + ) + agent_ts += len(observations[env_index]) + + done_episodes_to_run_env_to_module = [] + for env_index in range(self.num_envs): + # Call `on_episode_start()` callback (always after reset). + if env_index in call_on_episode_start: + self._make_on_episode_callback( + "on_episode_start", env_index, episodes + ) + # Make the `on_episode_step` callbacks. + else: + self._make_on_episode_callback( + "on_episode_step", env_index, episodes + ) + + # Episode is done. + if episodes[env_index].is_done: + eps += 1 + + # Make the `on_episode_end` callbacks (before finalizing the episode + # object). + self._make_on_episode_callback( + "on_episode_end", env_index, episodes + ) + + # TODO (simon): Check, if needed. I guess not b/c the complete episode is done. + # This needs to be executed here to remove the `SingleAgentEpisode`s + # that are done. + self._prune_zero_len_sa_episodes(episodes[env_index]) + + done_episodes_to_return.append(episodes[env_index]) + # Run a last time the `env_to_module` pipeline for these episodes + # to postprocess artifacts (e.g. observations to one-hot). + done_episodes_to_run_env_to_module.append(episodes[env_index]) + # Also early-out if we reach the number of episodes within this + # for-loop. + if eps == num_episodes: + break + + old_episode_id = episodes[env_index].id_ + # Create a new episode object with no data in it and execute + # `on_episode_created` callback (before the `env.reset()` call). + self._new_episode(env_index, episodes) + # Register the mapping of new episode ID to old episode ID. + shared_data["vector_env_episodes_map"].update( + {old_episode_id: episodes[env_index].id_} + ) + + # Env-to-module connector pass (cache results as we will do the RLModule + # forward pass only in the next `while`-iteration). + # Note, running the pipeline here ensures that we are not executing a + # pipeline run for agents that have died. This increases performance in + # case of environments in which 1000's of agents exist. Because the + # `VectorMultiAgentEnv` calls `reset()` for any terminated/truncated sub- + # environment `observations` and `infos` will always be returned for the + # `MultiAgentEpisode.add_reset_step`. + if self.module is not None: + if done_episodes_to_run_env_to_module: + # Run the env-to-module connector pipeline for all done episodes. + # Note, this is needed to postprocess last-step data, e.g. if the + # user uses a connector that one-hot encodes observations. + # Note, this pipeline run is not timed as the number of episodes + # can differ from `num_envs_per_env_runner` and would bias time + # measurements. + self._env_to_module( + batch={}, + episodes=done_episodes_to_run_env_to_module, + explore=explore, + rl_module=self.module, + shared_data=shared_data, + metrics=None, + ) + self._cached_to_module = self._env_to_module( + batch={}, + episodes=episodes, + explore=explore, + rl_module=self.module, + shared_data=shared_data, + metrics=self.metrics, + metrics_prefix_key=(ENV_TO_MODULE_CONNECTOR,), + ) + + # Numpy'ize the done episodes after running the connector pipeline. Note, + # that we need simple `list` objects in the + # `AddObservationsFromEpisodesToBatch` connector. Furthermore, we spare + # multiple `if` calls. + if self.config.episodes_to_numpy: + for episode in done_episodes_to_return: + # Any possibly compress observations. + episode.to_numpy() + + # Return done episodes ... + self._done_episodes_for_metrics.extend(done_episodes_to_return) + # ... and all ongoing episode chunks. + + # Also, make sure we start new episode chunks (continuing the ongoing episodes + # from the to-be-returned chunks). + ongoing_episodes_to_return = [] + # Only if we are doing individual timesteps: We have to maybe cut an ongoing + # episode and continue building it on the next call to `sample()`. + if num_timesteps is not None: + ongoing_episodes_continuations = [ + eps.cut(len_lookback_buffer=self.config.episode_lookback_horizon) + for eps in self._episodes + ] + + for eps in self._episodes: + # Just started Episodes do not have to be returned. There is no data + # in them anyway. + if eps.env_t == 0: + continue + eps.validate() + self._ongoing_episodes_for_metrics[eps.id_].append(eps) + + self._prune_zero_len_sa_episodes(eps) + + # Numpy'ize the episode. + if self.config.episodes_to_numpy: + # Any possibly compress observations. + ongoing_episodes_to_return.append(eps.to_numpy()) + # Leave episode as lists of individual (obs, action, etc..) items. + else: + ongoing_episodes_to_return.append(eps) + + # Continue collecting into the cut Episode chunks. + self._episodes = ongoing_episodes_continuations + + # Return collected episode data. + return done_episodes_to_return + ongoing_episodes_to_return + + def _reset_envs(self, episodes, shared_data, explore): + for env_index in range(self.num_envs): + self._new_episode(env_index, episodes) + + # Erase all cached ongoing episodes. This EnvRunner never completes or returns + # these from `get_metrics`, causing a memory leak. + self._ongoing_episodes_for_metrics.clear() + + # Try resetting the environment. + observations, infos = self._try_env_reset( + # Only seed (if seed provided) upon initial reset. + seed=self._seed if self._needs_initial_reset else None, + # TODO (sven): Support options? + options=None, + ) + + # Set the initial obs and infos in the episodes. + for env_index in range(self.num_envs): + episodes[env_index].add_env_reset( + observations=observations[env_index], + infos=infos[env_index], + ) + + # Run the env-to-module connector to make sure the reset-obs/infos have + # properly been processed (if applicable). + self._cached_to_module = None + if self.module: + self._cached_to_module = self._env_to_module( + rl_module=self.module, + episodes=episodes, + explore=explore, + shared_data=shared_data, + metrics=self.metrics, + metrics_key_prefix=(ENV_TO_MODULE_CONNECTOR,), + ) + + # Call `on_episode_start()` callbacks (always after reset). + for env_index in range(self.num_envs): + self._make_on_episode_callback("on_episode_start", env_index, episodes) + + @override(EnvRunner) + def get_spaces(self): + if self.env is None: + return self.spaces + # Return the already agent-to-module translated spaces from our connector + # pipeline. + return { + INPUT_ENV_SPACES: (self.env.observation_space, self.env.action_space), + INPUT_ENV_SINGLE_SPACES: ( + self.env.single_observation_space, + self.env.single_action_space, + ), + **{ + mid: (o, self._env_to_module.action_space[mid]) + for mid, o in self._env_to_module.observation_space.spaces.items() + }, + } + + @override(EnvRunner) + def get_metrics(self) -> ResultDict: + # Compute per-episode metrics (only on already completed episodes). + for eps in self._done_episodes_for_metrics: + assert eps.is_done + episode_length = len(eps) + agent_steps = defaultdict( + int, + {str(aid): len(sa_eps) for aid, sa_eps in eps.agent_episodes.items()}, + ) + episode_return = eps.get_return() + episode_duration_s = eps.get_duration_s() + + agent_episode_returns = defaultdict( + float, + { + str(sa_eps.agent_id): sa_eps.get_return() + for sa_eps in eps.agent_episodes.values() + }, + ) + module_episode_returns = defaultdict( + float, + { + sa_eps.module_id: sa_eps.get_return() + for sa_eps in eps.agent_episodes.values() + }, + ) + + # Don't forget about the already returned chunks of this episode. + if eps.id_ in self._ongoing_episodes_for_metrics: + for eps2 in self._ongoing_episodes_for_metrics[eps.id_]: + return_eps2 = eps2.get_return() + episode_length += len(eps2) + episode_return += return_eps2 + episode_duration_s += eps2.get_duration_s() + + for sa_eps in eps2.agent_episodes.values(): + return_sa = sa_eps.get_return() + agent_steps[str(sa_eps.agent_id)] += len(sa_eps) + agent_episode_returns[str(sa_eps.agent_id)] += return_sa + module_episode_returns[sa_eps.module_id] += return_sa + + del self._ongoing_episodes_for_metrics[eps.id_] + + self._log_episode_metrics( + episode_length, + episode_return, + episode_duration_s, + agent_episode_returns, + module_episode_returns, + dict(agent_steps), + ) + + # Now that we have logged everything, clear cache of done episodes. + self._done_episodes_for_metrics.clear() + + # Return reduced metrics. + return self.metrics.reduce() + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + # Basic state dict. + state = { + NUM_ENV_STEPS_SAMPLED_LIFETIME: ( + self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME, default=0) + ), + } + + # RLModule (MultiRLModule) component. + if self._check_component(COMPONENT_RL_MODULE, components, not_components): + state[COMPONENT_RL_MODULE] = self.module.get_state( + components=self._get_subcomponents(COMPONENT_RL_MODULE, components), + not_components=self._get_subcomponents( + COMPONENT_RL_MODULE, not_components + ), + **kwargs, + ) + state[WEIGHTS_SEQ_NO] = self._weights_seq_no + + # Env-to-module connector. + if self._check_component( + COMPONENT_ENV_TO_MODULE_CONNECTOR, components, not_components + ): + state[COMPONENT_ENV_TO_MODULE_CONNECTOR] = self._env_to_module.get_state() + # Module-to-env connector. + if self._check_component( + COMPONENT_MODULE_TO_ENV_CONNECTOR, components, not_components + ): + state[COMPONENT_MODULE_TO_ENV_CONNECTOR] = self._module_to_env.get_state() + + return state + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + if COMPONENT_ENV_TO_MODULE_CONNECTOR in state: + self._env_to_module.set_state(state[COMPONENT_ENV_TO_MODULE_CONNECTOR]) + if COMPONENT_MODULE_TO_ENV_CONNECTOR in state: + self._module_to_env.set_state(state[COMPONENT_MODULE_TO_ENV_CONNECTOR]) + + # Update RLModule state. + if COMPONENT_RL_MODULE in state: + # A missing value for WEIGHTS_SEQ_NO or a value of 0 means: Force the + # update. + weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0) + + # Only update the weights, if this is the first synchronization or + # if the weights of this `EnvRunner` lacks behind the actual ones. + if weights_seq_no == 0 or self._weights_seq_no < weights_seq_no: + rl_module_state = state[COMPONENT_RL_MODULE] + if isinstance(rl_module_state, ray.ObjectRef): + rl_module_state = ray.get(rl_module_state) + self.module.set_state(rl_module_state) + + # Update weights_seq_no, if the new one is > 0. + if weights_seq_no > 0: + self._weights_seq_no = weights_seq_no + + # Update lifetime counters. + if NUM_ENV_STEPS_SAMPLED_LIFETIME in state: + self.metrics.set_value( + key=NUM_ENV_STEPS_SAMPLED_LIFETIME, + value=state[NUM_ENV_STEPS_SAMPLED_LIFETIME], + reduce="sum", + with_throughput=True, + ) + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + {"config": self.config}, # **kwargs + ) + + @override(Checkpointable) + def get_metadata(self): + metadata = Checkpointable.get_metadata(self) + metadata.update( + { + # TODO (sven): Maybe add serialized (JSON-writable) config here? + } + ) + return metadata + + @override(Checkpointable) + def get_checkpointable_components(self): + return [ + (COMPONENT_RL_MODULE, self.module), + (COMPONENT_ENV_TO_MODULE_CONNECTOR, self._env_to_module), + (COMPONENT_MODULE_TO_ENV_CONNECTOR, self._module_to_env), + ] + + @override(EnvRunner) + def assert_healthy(self): + """Checks that self.__init__() has been completed properly. + + Ensures that the instances has a `MultiRLModule` and an + environment defined. + + Raises: + AssertionError: If the EnvRunner Actor has NOT been properly initialized. + """ + # Make sure, we have built our gym.vector.Env and RLModule properly. + assert self.env and self.module + + @override(EnvRunner) + def make_env(self): + # If an env already exists, try closing it first (to allow it to properly + # cleanup). + if self.env is not None: + try: + self.env.close() + except Exception as e: + logger.warning( + "Tried closing the existing env (multi-agent), but failed with " + f"error: {e.args[0]}" + ) + del self.env + + env_ctx = self.config.env_config + if not isinstance(env_ctx, EnvContext): + env_ctx = EnvContext( + env_ctx, + worker_index=self.worker_index, + num_workers=self.config.num_env_runners, + remote=self.config.remote_worker_envs, + ) + + # No env provided -> Error. + if not self.config.env: + raise ValueError( + "`config.env` is not provided! You should provide a valid environment " + "to your config through `config.environment([env descriptor e.g. " + "'CartPole-v1'])`." + ) + # Register env for the local context. + # Note, `gym.register` has to be called on each worker. + elif isinstance(self.config.env, str) and _global_registry.contains( + ENV_CREATOR, self.config.env + ): + entry_point = partial( + _global_registry.get(ENV_CREATOR, self.config.env), + env_ctx, + ) + else: + entry_point = partial( + _gym_env_creator, + env_descriptor=self.config.env, + env_context=env_ctx, + ) + gym.register( + "rllib-multi-agent-env-v0", + entry_point=entry_point, + disable_env_checker=True, + ) + vectorize_mode = self.config.gym_env_vectorize_mode + + # Perform actual gym.make call. + self.env = make_vec( + "rllib-multi-agent-env-v0", + num_envs=self.config.num_envs_per_env_runner, + vectorization_mode=( + vectorize_mode + if isinstance(vectorize_mode, gym.envs.registration.VectorizeMode) + else gym.envs.registration.VectorizeMode(vectorize_mode.lower()) + ), + ) + self.num_envs: int = self.env.num_envs + assert self.num_envs == self.config.num_envs_per_env_runner + # self.env: MultiAgentEnv = gym.make("rllib-multi-agent-env-v0") + # self.num_envs = 1 + # If required, check the created MultiAgentEnv instances. + if not self.config.disable_env_checking: + try: + for env in self.env.envs: + check_multiagent_environments(env.unwrapped) + except Exception as e: + logger.exception(e.args[0]) + # If not required, still check the type (must be `VectorMultiAgentEnv``). + else: + try: + assert isinstance(self.env, VectorMultiAgentEnv) + assert isinstance(self.env.envs[0].unwrapped, MultiAgentEnv) + except AssertionError: + logger.exception( + "When using the `MultiAgentEnvRunner`, the environment must " + f"inherit from `ray.rllib.env.vector.vector_multi_agent_env." + f"VectorMultiAgentEnv` (but yours is {self.env}) and the individual" + " envs must inherit from `MultiAgentEnv` (but yours is " + f"{self.env.envs[0].unwrapped})!" + ) + except TypeError: + logger.exception( + "When using the `MultiAgentEnvRunner`, the env must " + "have a subscriptable `self.envs` attribute!" + ) + + # Set the flag to reset all envs upon the next `sample()` call. + self._needs_initial_reset = True + + # Call the `on_environment_created` callback. + make_callback( + "on_environment_created", + callbacks_objects=self._callbacks, + callbacks_functions=self.config.callbacks_on_environment_created, + kwargs=dict( + env_runner=self, + metrics_logger=self.metrics, + env=self.env.unwrapped, + env_context=env_ctx, + ), + ) + + @override(EnvRunner) + def make_module(self): + env = self.env.unwrapped if self.env is not None else None + try: + module_spec: MultiRLModuleSpec = self.config.get_multi_rl_module_spec( + env=env, spaces=self.get_spaces(), inference_only=True + ) + # Build the module from its spec. + self.module = module_spec.build() + # Move the RLModule to our device. + # TODO (sven): In order to make this framework-agnostic, we should maybe + # make the MultiRLModule.build() method accept a device OR create an + # additional `(Multi)RLModule.to()` override. + if torch: + self.module.foreach_module( + lambda mid, mod: ( + mod.to(self._device) + if isinstance(mod, torch.nn.Module) + else mod + ) + ) + + # If `AlgorithmConfig.get_rl_module_spec()` is not implemented, this env runner + # will not have an RLModule, but might still be usable with random actions. + except NotImplementedError: + self.module = None + + @override(EnvRunner) + def stop(self): + # Note, `MultiAgentEnv` inherits `close()`-method from `gym.Env`. + if self.env is not None: + self.env.close() + + def _setup_metrics(self): + self._done_episodes_for_metrics: List[MultiAgentEpisode] = [] + self._ongoing_episodes_for_metrics: DefaultDict[ + EpisodeID, List[MultiAgentEpisode] + ] = defaultdict(list) + + def _new_episode(self, env_index, episodes=None): + episodes = episodes if episodes is not None else self._episodes + episodes[env_index] = MultiAgentEpisode( + observation_space={ + aid: self.env.envs[env_index].unwrapped.get_observation_space(aid) + for aid in self.env.envs[env_index].unwrapped.possible_agents + }, + action_space={ + aid: self.env.envs[env_index].unwrapped.get_action_space(aid) + for aid in self.env.envs[env_index].unwrapped.possible_agents + }, + agent_to_module_mapping_fn=self.config.policy_mapping_fn, + ) + self._make_on_episode_callback("on_episode_created", env_index, episodes) + + def _make_on_episode_callback( + self, which: str, idx: int, episodes: List[MultiAgentEpisode] + ): + kwargs = dict( + episode=episodes[idx], + env_runner=self, + metrics_logger=self.metrics, + env=self.env.unwrapped, + rl_module=self.module, + env_index=idx, + ) + if which == "on_episode_end": + kwargs["prev_episode_chunks"] = self._ongoing_episodes_for_metrics[ + episodes[idx].id_ + ] + + make_callback( + which, + callbacks_objects=self._callbacks, + callbacks_functions=getattr(self.config, f"callbacks_{which}"), + kwargs=kwargs, + ) + + def _increase_sampled_metrics(self, num_steps, next_obs, episode): + # Env steps. + self.metrics.log_value( + NUM_ENV_STEPS_SAMPLED, num_steps, reduce="sum", clear_on_reduce=True + ) + self.metrics.log_value( + NUM_ENV_STEPS_SAMPLED_LIFETIME, + num_steps, + reduce="sum", + with_throughput=True, + ) + # Completed episodes. + if episode.is_done: + self.metrics.log_value(NUM_EPISODES, 1, reduce="sum", clear_on_reduce=True) + self.metrics.log_value(NUM_EPISODES_LIFETIME, 1, reduce="sum") + + # Record agent and module metrics. + for aid in next_obs: + self.metrics.log_value( + (NUM_AGENT_STEPS_SAMPLED, str(aid)), + 1, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_SAMPLED_LIFETIME, str(aid)), + 1, + reduce="sum", + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_SAMPLED, episode.module_for(aid)), + 1, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_SAMPLED_LIFETIME, episode.module_for(aid)), + 1, + reduce="sum", + ) + return num_steps + + def _log_episode_metrics( + self, + length, + ret, + sec, + agents=None, + modules=None, + agent_steps=None, + ): + # Log general episode metrics. + + # Use the configured window, but factor in the parallelism of the EnvRunners. + # As a result, we only log the last `window / num_env_runners` steps here, + # b/c everything gets parallel-merged in the Algorithm process. + win = max( + 1, + int( + math.ceil( + self.config.metrics_num_episodes_for_smoothing + / (self.config.num_env_runners or 1) + ) + ), + ) + + self.metrics.log_dict( + { + EPISODE_LEN_MEAN: length, + EPISODE_RETURN_MEAN: ret, + EPISODE_DURATION_SEC_MEAN: sec, + **( + { + # Per-agent returns. + "agent_episode_returns_mean": agents, + # Per-RLModule returns. + "module_episode_returns_mean": modules, + "agent_steps": agent_steps, + } + if agents is not None + else {} + ), + }, + window=win, + ) + # For some metrics, log min/max as well. + self.metrics.log_dict( + { + EPISODE_LEN_MIN: length, + EPISODE_RETURN_MIN: ret, + }, + reduce="min", + window=win, + ) + self.metrics.log_dict( + { + EPISODE_LEN_MAX: length, + EPISODE_RETURN_MAX: ret, + }, + reduce="max", + window=win, + ) + + @staticmethod + def _prune_zero_len_sa_episodes(episode: MultiAgentEpisode): + for agent_id, agent_eps in episode.agent_episodes.copy().items(): + if not agent_eps: + del episode.agent_episodes[agent_id] + + @Deprecated( + new="MultiAgentEnvRunner.get_state(components='rl_module')", + error=False, + ) + def get_weights(self, modules=None): + rl_module_state = self.get_state(components=COMPONENT_RL_MODULE)[ + COMPONENT_RL_MODULE + ] + return rl_module_state + + @Deprecated(new="MultiAgentEnvRunner.set_state()", error=False) + def set_weights( + self, + weights: ModelWeights, + global_vars: Optional[Dict] = None, + weights_seq_no: int = 0, + ) -> None: + assert global_vars is None + return self.set_state( + { + COMPONENT_RL_MODULE: weights, + WEIGHTS_SEQ_NO: weights_seq_no, + } + ) diff --git a/lib/python3.12/site-packages/ray/rllib/env/multi_agent_episode.py b/lib/python3.12/site-packages/ray/rllib/env/multi_agent_episode.py new file mode 100644 index 0000000000000000000000000000000000000000..35327fca5b927014150a5a68c1c1427083932e66 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/multi_agent_episode.py @@ -0,0 +1,2795 @@ +import copy +import time +import uuid +from collections import defaultdict +from typing import ( + Any, + Callable, + Collection, + DefaultDict, + Dict, + List, + Optional, + Set, + Union, +) + +import gymnasium as gym + +from ray._common.deprecation import Deprecated +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.env.utils.infinite_lookback_buffer import InfiniteLookbackBuffer +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils import force_list +from ray.rllib.utils.error import MultiAgentEnvError +from ray.rllib.utils.spaces.space_utils import batch +from ray.rllib.utils.typing import AgentID, ModuleID, MultiAgentDict +from ray.util.annotations import PublicAPI + + +# TODO (simon): Include cases in which the number of agents in an +# episode are shrinking or growing during the episode itself. +@PublicAPI(stability="alpha") +class MultiAgentEpisode: + """Stores multi-agent episode data. + + The central attribute of the class is the timestep mapping + `self.env_t_to_agent_t` that maps AgentIDs to their specific environment steps to + the agent's own scale/timesteps. + + Each AgentID in the `MultiAgentEpisode` has its own `SingleAgentEpisode` object + in which this agent's data is stored. Together with the env_t_to_agent_t mapping, + we can extract information either on any individual agent's time scale or from + the (global) multi-agent environment time scale. + + Extraction of data from a MultiAgentEpisode happens via the getter APIs, e.g. + `get_observations()`, which work analogous to the ones implemented in the + `SingleAgentEpisode` class. + + Note that recorded `terminateds`/`truncateds` come as simple + `MultiAgentDict`s mapping AgentID to bools and thus have no assignment to a + certain timestep (analogous to a SingleAgentEpisode's single `terminated/truncated` + boolean flag). Instead we assign it to the last observation recorded. + Theoretically, there could occur edge cases in some environments + where an agent receives partial rewards and then terminates without + a last observation. In these cases, we duplicate the last observation. + + Also, if no initial observation has been received yet for an agent, but + some rewards for this same agent already occurred, we delete the agent's data + up to here, b/c there is nothing to learn from these "premature" rewards. + """ + + __slots__ = ( + "id_", + "agent_to_module_mapping_fn", + "_agent_to_module_mapping", + "observation_space", + "action_space", + "env_t_started", + "env_t", + "agent_t_started", + "env_t_to_agent_t", + "_hanging_actions_end", + "_hanging_extra_model_outputs_end", + "_hanging_rewards_end", + "_hanging_rewards_begin", + "is_terminated", + "is_truncated", + "agent_episodes", + "_last_step_time", + "_len_lookback_buffers", + "_start_time", + "_custom_data", + ) + + SKIP_ENV_TS_TAG = "S" + + def __init__( + self, + id_: Optional[str] = None, + *, + observations: Optional[List[MultiAgentDict]] = None, + observation_space: Optional[gym.Space] = None, + infos: Optional[List[MultiAgentDict]] = None, + actions: Optional[List[MultiAgentDict]] = None, + action_space: Optional[gym.Space] = None, + rewards: Optional[List[MultiAgentDict]] = None, + terminateds: Union[MultiAgentDict, bool] = False, + truncateds: Union[MultiAgentDict, bool] = False, + extra_model_outputs: Optional[List[MultiAgentDict]] = None, + env_t_started: Optional[int] = None, + agent_t_started: Optional[Dict[AgentID, int]] = None, + len_lookback_buffer: Union[int, str] = "auto", + agent_episode_ids: Optional[Dict[AgentID, str]] = None, + agent_module_ids: Optional[Dict[AgentID, ModuleID]] = None, + agent_to_module_mapping_fn: Optional[ + Callable[[AgentID, "MultiAgentEpisode"], ModuleID] + ] = None, + ): + """Initializes a `MultiAgentEpisode`. + + Args: + id_: Optional. Either a string to identify an episode or None. + If None, a hexadecimal id is created. In case of providing + a string, make sure that it is unique, as episodes get + concatenated via this string. + observations: A list of dictionaries mapping agent IDs to observations. + Can be None. If provided, should match all other episode data + (actions, rewards, etc.) in terms of list lengths and agent IDs. + observation_space: An optional gym.spaces.Dict mapping agent IDs to + individual agents' spaces, which all (individual agents') observations + should abide to. If not None and this MultiAgentEpisode is numpy'ized + (via the `self.to_numpy()` method), and data is appended or set, the new + data will be checked for correctness. + infos: A list of dictionaries mapping agent IDs to info dicts. + Can be None. If provided, should match all other episode data + (observations, rewards, etc.) in terms of list lengths and agent IDs. + actions: A list of dictionaries mapping agent IDs to actions. + Can be None. If provided, should match all other episode data + (observations, rewards, etc.) in terms of list lengths and agent IDs. + action_space: An optional gym.spaces.Dict mapping agent IDs to + individual agents' spaces, which all (individual agents') actions + should abide to. If not None and this MultiAgentEpisode is numpy'ized + (via the `self.to_numpy()` method), and data is appended or set, the new + data will be checked for correctness. + rewards: A list of dictionaries mapping agent IDs to rewards. + Can be None. If provided, should match all other episode data + (actions, rewards, etc.) in terms of list lengths and agent IDs. + terminateds: A boolean defining if an environment has + terminated OR a MultiAgentDict mapping individual agent ids + to boolean flags indicating whether individual agents have terminated. + A special __all__ key in these dicts indicates, whether the episode + is terminated for all agents. + The default is `False`, i.e. the episode has not been terminated. + truncateds: A boolean defining if the environment has been + truncated OR a MultiAgentDict mapping individual agent ids + to boolean flags indicating whether individual agents have been + truncated. A special __all__ key in these dicts indicates, whether the + episode is truncated for all agents. + The default is `False`, i.e. the episode has not been truncated. + extra_model_outputs: A list of dictionaries mapping agent IDs to their + corresponding extra model outputs. Each of these "outputs" is a dict + mapping keys (str) to model output values, for example for + `key=STATE_OUT`, the values would be the internal state outputs for + that agent. + env_t_started: The env timestep (int) that defines the starting point + of the episode. This is only larger zero, if an already ongoing episode + chunk is being created, for example by slicing an ongoing episode or + by calling the `cut()` method on an ongoing episode. + agent_t_started: A dict mapping AgentIDs to the respective agent's (local) + timestep at which its SingleAgentEpisode chunk started. + len_lookback_buffer: The size of the lookback buffers to keep in + front of this Episode for each type of data (observations, actions, + etc..). If larger 0, will interpret the first `len_lookback_buffer` + items in each type of data as NOT part of this actual + episode chunk, but instead serve as "historical" record that may be + viewed and used to derive new data from. For example, it might be + necessary to have a lookback buffer of four if you would like to do + observation frame stacking and your episode has been cut and you are now + operating on a new chunk (continuing from the cut one). Then, for the + first 3 items, you would have to be able to look back into the old + chunk's data. + If `len_lookback_buffer` is "auto" (default), will interpret all + provided data in the constructor as part of the lookback buffers. + agent_episode_ids: An optional dict mapping AgentIDs + to their corresponding `SingleAgentEpisode`. If None, each + `SingleAgentEpisode` in `MultiAgentEpisode.agent_episodes` + will generate a hexadecimal code. If a dictionary is provided, + make sure that IDs are unique, because the agents' `SingleAgentEpisode` + instances are concatenated or recreated by it. + agent_module_ids: An optional dict mapping AgentIDs to their respective + ModuleIDs (these mapping are always valid for an entire episode and + thus won't change during the course of this episode). If a mapping from + agent to module has already been provided via this dict, the (optional) + `agent_to_module_mapping_fn` will NOT be used again to map the same + agent (agents do not change their assigned module in the course of + one episode). + agent_to_module_mapping_fn: A callable taking an AgentID and a + MultiAgentEpisode as args and returning a ModuleID. Used to map agents + that have not been mapped yet (because they just entered this episode) + to a ModuleID. The resulting ModuleID is only stored inside the agent's + SingleAgentEpisode object. + """ + self.id_: str = id_ or uuid.uuid4().hex + if agent_to_module_mapping_fn is None: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + agent_to_module_mapping_fn = ( + AlgorithmConfig.DEFAULT_AGENT_TO_MODULE_MAPPING_FN + ) + self.agent_to_module_mapping_fn = agent_to_module_mapping_fn + # In case a user - e.g. via callbacks - already forces a mapping to happen + # via the `module_for()` API even before the agent has entered the episode + # (and has its SingleAgentEpisode created), we store all aldeary done mappings + # in this dict here. + self._agent_to_module_mapping: Dict[AgentID, ModuleID] = agent_module_ids or {} + + # Lookback buffer length is not provided. Interpret all provided data as + # lookback buffer. + if len_lookback_buffer == "auto": + len_lookback_buffer = len(rewards or []) + self._len_lookback_buffers = len_lookback_buffer + + self.observation_space = observation_space or {} + self.action_space = action_space or {} + + terminateds = terminateds or {} + truncateds = truncateds or {} + + # The global last timestep of the episode and the timesteps when this chunk + # started (excluding a possible lookback buffer). + self.env_t_started = env_t_started or 0 + self.env_t = ( + (len(rewards) if rewards is not None else 0) + - self._len_lookback_buffers + + self.env_t_started + ) + self.agent_t_started = defaultdict(int, agent_t_started or {}) + + # Keeps track of the correspondence between agent steps and environment steps. + # Under each AgentID as key is a InfiniteLookbackBuffer with the following + # data in it: + # The indices of the items in the data represent environment timesteps, + # starting from index=0 for the `env.reset()` and with each `env.step()` call + # increase by 1. + # The values behind these (env timestep) indices represent the agent timesteps + # happening at these env timesteps and the special value of + # `self.SKIP_ENV_TS_TAG` means that the agent did NOT step at the given env + # timestep. + # Thus, agents that are part of the reset obs, will start their mapping data + # with a [0 ...], all other agents will start their mapping data with: + # [self.SKIP_ENV_TS_TAG, ...]. + self.env_t_to_agent_t: DefaultDict[ + AgentID, InfiniteLookbackBuffer + ] = defaultdict(InfiniteLookbackBuffer) + + # Create caches for hanging actions/rewards/extra_model_outputs. + # When an agent gets an observation (and then sends an action), but does not + # receive immediately a next observation, we store the "hanging" action (and + # related rewards and extra model outputs) in the caches postfixed w/ `_end` + # until the next observation is received. + self._hanging_actions_end = {} + self._hanging_extra_model_outputs_end = defaultdict(dict) + self._hanging_rewards_end = defaultdict(float) + + # In case of a `cut()` or `slice()`, we also need to store the hanging actions, + # rewards, and extra model outputs that were already "hanging" in preceeding + # episode slice. + self._hanging_rewards_begin = defaultdict(float) + + # If this is an ongoing episode than the last `__all__` should be `False` + self.is_terminated: bool = ( + terminateds + if isinstance(terminateds, bool) + else terminateds.get("__all__", False) + ) + + # If this is an ongoing episode than the last `__all__` should be `False` + self.is_truncated: bool = ( + truncateds + if isinstance(truncateds, bool) + else truncateds.get("__all__", False) + ) + + # The individual agent SingleAgentEpisode objects. + self.agent_episodes: Dict[AgentID, SingleAgentEpisode] = {} + self._init_single_agent_episodes( + agent_module_ids=agent_module_ids, + agent_episode_ids=agent_episode_ids, + observations=observations, + infos=infos, + actions=actions, + rewards=rewards, + terminateds=terminateds, + truncateds=truncateds, + extra_model_outputs=extra_model_outputs, + ) + + # Cache for custom data. May be used to store custom metrics from within a + # callback for the ongoing episode (e.g. render images). + self._custom_data = {} + + # Keep timer stats on deltas between steps. + self._start_time = None + self._last_step_time = None + + # Validate ourselves. + self.validate() + + def add_env_reset( + self, + *, + observations: MultiAgentDict, + infos: Optional[MultiAgentDict] = None, + ) -> None: + """Stores initial observation. + + Args: + observations: A dictionary mapping agent IDs to initial observations. + Note that some agents may not have an initial observation. + infos: A dictionary mapping agent IDs to initial info dicts. + Note that some agents may not have an initial info dict. If not None, + the agent IDs in `infos` must be a subset of those in `observations` + meaning it would not be allowed to have an agent with an info dict, + but not with an observation. + """ + assert not self.is_done + # Assume that this episode is completely empty and has not stepped yet. + # Leave self.env_t (and self.env_t_started) at 0. + assert self.env_t == self.env_t_started == 0 + infos = infos or {} + + # Note, all agents will have an initial observation, some may have an initial + # info dict as well. + for agent_id, agent_obs in observations.items(): + # Update env_t_to_agent_t mapping (all agents that are part of the reset + # obs have their first mapping 0 (env_t) -> 0 (agent_t)). + self.env_t_to_agent_t[agent_id].append(0) + # Create SingleAgentEpisode, if necessary. + if agent_id not in self.agent_episodes: + self.agent_episodes[agent_id] = SingleAgentEpisode( + agent_id=agent_id, + module_id=self.module_for(agent_id), + multi_agent_episode_id=self.id_, + observation_space=self.observation_space.get(agent_id), + action_space=self.action_space.get(agent_id), + ) + # Add initial observations (and infos) to the agent's episode. + self.agent_episodes[agent_id].add_env_reset( + observation=agent_obs, + infos=infos.get(agent_id), + ) + + # Validate our data. + self.validate() + + # Start the timer for this episode. + self._start_time = time.perf_counter() + + def add_env_step( + self, + observations: MultiAgentDict, + actions: MultiAgentDict, + rewards: MultiAgentDict, + infos: Optional[MultiAgentDict] = None, + *, + terminateds: Optional[MultiAgentDict] = None, + truncateds: Optional[MultiAgentDict] = None, + extra_model_outputs: Optional[MultiAgentDict] = None, + ) -> None: + """Adds a timestep to the episode. + + Args: + observations: A dictionary mapping agent IDs to their corresponding + next observations. Note that some agents may not have stepped at this + timestep. + actions: Mandatory. A dictionary mapping agent IDs to their + corresponding actions. Note that some agents may not have stepped at + this timestep. + rewards: Mandatory. A dictionary mapping agent IDs to their + corresponding observations. Note that some agents may not have stepped + at this timestep. + infos: A dictionary mapping agent IDs to their + corresponding info. Note that some agents may not have stepped at this + timestep. + terminateds: A dictionary mapping agent IDs to their `terminated` flags, + indicating, whether the environment has been terminated for them. + A special `__all__` key indicates that the episode is terminated for + all agent IDs. + terminateds: A dictionary mapping agent IDs to their `truncated` flags, + indicating, whether the environment has been truncated for them. + A special `__all__` key indicates that the episode is `truncated` for + all agent IDs. + extra_model_outputs: A dictionary mapping agent IDs to their + corresponding specific model outputs (also in a dictionary; e.g. + `vf_preds` for PPO). + """ + # Cannot add data to an already done episode. + if self.is_done: + raise MultiAgentEnvError( + "Cannot call `add_env_step` on a MultiAgentEpisode that is already " + "done!" + ) + + infos = infos or {} + terminateds = terminateds or {} + truncateds = truncateds or {} + extra_model_outputs = extra_model_outputs or {} + + # Increase (global) env step by one. + self.env_t += 1 + + # Find out, whether this episode is terminated/truncated (for all agents). + # Case 1: all agents are terminated or all are truncated. + self.is_terminated = terminateds.get("__all__", False) + self.is_truncated = truncateds.get("__all__", False) + # Find all agents that were done at prior timesteps and add the agents that are + # done at the present timestep. + agents_done = set( + [aid for aid, sa_eps in self.agent_episodes.items() if sa_eps.is_done] + + [aid for aid in terminateds if terminateds[aid]] + + [aid for aid in truncateds if truncateds[aid]] + ) + # Case 2: Some agents are truncated and the others are terminated -> Declare + # this episode as terminated. + if all(aid in set(agents_done) for aid in self.agent_ids): + self.is_terminated = True + + # For all agents that are not stepping in this env step, but that are not done + # yet -> Add a skip tag to their env- to agent-step mappings. + stepped_agent_ids = set(observations.keys()) + for agent_id, env_t_to_agent_t in self.env_t_to_agent_t.items(): + if agent_id not in stepped_agent_ids: + env_t_to_agent_t.append(self.SKIP_ENV_TS_TAG) + + # Loop through all agent IDs that we received data for in this step: + # Those found in observations, actions, and rewards. + agent_ids_with_data = ( + set(observations.keys()) + | set(actions.keys()) + | set(rewards.keys()) + | set(terminateds.keys()) + | set(truncateds.keys()) + | set( + self.agent_episodes.keys() + if terminateds.get("__all__") or truncateds.get("__all__") + else set() + ) + ) - {"__all__"} + for agent_id in agent_ids_with_data: + if agent_id not in self.agent_episodes: + sa_episode = SingleAgentEpisode( + agent_id=agent_id, + module_id=self.module_for(agent_id), + multi_agent_episode_id=self.id_, + observation_space=self.observation_space.get(agent_id), + action_space=self.action_space.get(agent_id), + ) + else: + sa_episode = self.agent_episodes[agent_id] + + # Collect value to be passed (at end of for-loop) into `add_env_step()` + # call. + _observation = observations.get(agent_id) + _action = actions.get(agent_id) + _reward = rewards.get(agent_id) + _infos = infos.get(agent_id) + _terminated = terminateds.get(agent_id, False) or self.is_terminated + _truncated = truncateds.get(agent_id, False) or self.is_truncated + _extra_model_outputs = extra_model_outputs.get(agent_id) + + # The value to place into the env- to agent-step map for this agent ID. + # _agent_step = self.SKIP_ENV_TS_TAG + + # Agents, whose SingleAgentEpisode had already been done before this + # step should NOT have received any data in this step. + if sa_episode.is_done and any( + v is not None + for v in [_observation, _action, _reward, _infos, _extra_model_outputs] + ): + raise MultiAgentEnvError( + f"Agent {agent_id} already had its `SingleAgentEpisode.is_done` " + f"set to True, but still received data in a following step! " + f"obs={_observation} act={_action} rew={_reward} info={_infos} " + f"extra_model_outputs={_extra_model_outputs}." + ) + _reward = _reward or 0.0 + + # CASE 1: A complete agent step is available (in one env step). + # ------------------------------------------------------------- + # We have an observation and an action for this agent -> + # Add the agent step to the single agent episode. + # ... action -> next obs + reward ... + if _observation is not None and _action is not None: + if agent_id not in rewards: + raise MultiAgentEnvError( + f"Agent {agent_id} acted (and received next obs), but did NOT " + f"receive any reward from the env!" + ) + + # CASE 2: Step gets completed with a hanging action OR first observation. + # ------------------------------------------------------------------------ + # We have an observation, but no action -> + # a) Action (and extra model outputs) must be hanging already. Also use + # collected hanging rewards and extra_model_outputs. + # b) The observation is the first observation for this agent ID. + elif _observation is not None and _action is None: + _action = self._hanging_actions_end.pop(agent_id, None) + + # We have a hanging action (the agent had acted after the previous + # observation, but the env had not responded - until now - with another + # observation). + # ...[hanging action] ... ... -> next obs + (reward)? ... + if _action is not None: + # Get the extra model output if available. + _extra_model_outputs = self._hanging_extra_model_outputs_end.pop( + agent_id, None + ) + _reward = self._hanging_rewards_end.pop(agent_id, 0.0) + _reward + # First observation for this agent, we have no hanging action. + # ... [done]? ... -> [1st obs for agent ID] + else: + # The agent is already done -> The agent thus has never stepped once + # and we do not have to create a SingleAgentEpisode for it. + if _terminated or _truncated: + self._del_hanging(agent_id) + continue + # This must be the agent's initial observation. + else: + # Prepend n skip tags to this agent's mapping + the initial [0]. + assert agent_id not in self.env_t_to_agent_t + self.env_t_to_agent_t[agent_id].extend( + [self.SKIP_ENV_TS_TAG] * self.env_t + [0] + ) + self.env_t_to_agent_t[ + agent_id + ].lookback = self._len_lookback_buffers + # Make `add_env_reset` call and continue with next agent. + sa_episode.add_env_reset(observation=_observation, infos=_infos) + # Add possible reward to begin cache. + self._hanging_rewards_begin[agent_id] += _reward + # Now that the SAEps is valid, add it to our dict. + self.agent_episodes[agent_id] = sa_episode + continue + + # CASE 3: Step is started (by an action), but not completed (no next obs). + # ------------------------------------------------------------------------ + # We have no observation, but we have a hanging action (used when we receive + # the next obs for this agent in the future). + elif agent_id not in observations and agent_id in actions: + # Agent got truncated -> Error b/c we would need a last (truncation) + # observation for this (otherwise, e.g. bootstrapping would not work). + # [previous obs] [action] (hanging) ... ... [truncated] + if _truncated: + raise MultiAgentEnvError( + f"Agent {agent_id} acted and then got truncated, but did NOT " + "receive a last (truncation) observation, required for e.g. " + "value function bootstrapping!" + ) + # Agent got terminated. + # [previous obs] [action] (hanging) ... ... [terminated] + elif _terminated: + # If the agent was terminated and no observation is provided, + # duplicate the previous one (this is a technical "fix" to properly + # complete the single agent episode; this last observation is never + # used for learning anyway). + _observation = sa_episode._last_added_observation + _infos = sa_episode._last_added_infos + # Agent is still alive. + # [previous obs] [action] (hanging) ... + else: + # Hanging action, reward, and extra_model_outputs. + assert agent_id not in self._hanging_actions_end + self._hanging_actions_end[agent_id] = _action + self._hanging_rewards_end[agent_id] = _reward + self._hanging_extra_model_outputs_end[ + agent_id + ] = _extra_model_outputs + + # CASE 4: Step has started in the past and is still ongoing (no observation, + # no action). + # -------------------------------------------------------------------------- + # Record reward and terminated/truncated flags. + else: + _action = self._hanging_actions_end.get(agent_id) + + # Agent is done. + if _terminated or _truncated: + # If the agent has NOT stepped, we treat it as not being + # part of this episode. + # ... ... [other agents doing stuff] ... ... [agent done] + if _action is None: + self._del_hanging(agent_id) + continue + + # Agent got truncated -> Error b/c we would need a last (truncation) + # observation for this (otherwise, e.g. bootstrapping would not + # work). + if _truncated: + raise MultiAgentEnvError( + f"Agent {agent_id} acted and then got truncated, but did " + "NOT receive a last (truncation) observation, required " + "for e.g. value function bootstrapping!" + ) + + # [obs] ... ... [hanging action] ... ... [done] + # If the agent was terminated and no observation is provided, + # duplicate the previous one (this is a technical "fix" to properly + # complete the single agent episode; this last observation is never + # used for learning anyway). + _observation = sa_episode._last_added_observation + _infos = sa_episode._last_added_infos + # `_action` is already `get` above. We don't need to pop out from + # the cache as it gets wiped out anyway below b/c the agent is + # done. + _extra_model_outputs = self._hanging_extra_model_outputs_end.pop( + agent_id, None + ) + _reward = self._hanging_rewards_end.pop(agent_id, 0.0) + _reward + # The agent is still alive, just add current reward to cache. + else: + # But has never stepped in this episode -> add to begin cache. + if agent_id not in self.agent_episodes: + self._hanging_rewards_begin[agent_id] += _reward + # Otherwise, add to end cache. + else: + self._hanging_rewards_end[agent_id] += _reward + + # If agent is stepping, add timestep to `SingleAgentEpisode`. + if _observation is not None: + sa_episode.add_env_step( + observation=_observation, + action=_action, + reward=_reward, + infos=_infos, + terminated=_terminated, + truncated=_truncated, + extra_model_outputs=_extra_model_outputs, + ) + # Update the env- to agent-step mapping. + self.env_t_to_agent_t[agent_id].append( + len(sa_episode) + sa_episode.observations.lookback + ) + + # Agent is also done. -> Erase all hanging values for this agent + # (they should be empty at this point anyways). + if _terminated or _truncated: + self._del_hanging(agent_id) + + # Validate our data. + self.validate() + + # Step time stats. + self._last_step_time = time.perf_counter() + if self._start_time is None: + self._start_time = self._last_step_time + + def validate(self) -> None: + """Validates the episode's data. + + This function ensures that the data stored to a `MultiAgentEpisode` is + in order (e.g. that the correct number of observations, actions, rewards + are there). + """ + for eps in self.agent_episodes.values(): + eps.validate() + + # TODO (sven): Validate MultiAgentEpisode specifics, like the timestep mappings, + # action/reward caches, etc.. + + @property + def custom_data(self): + return self._custom_data + + @property + def is_reset(self) -> bool: + """Returns True if `self.add_env_reset()` has already been called.""" + return any( + len(sa_episode.observations) > 0 + for sa_episode in self.agent_episodes.values() + ) + + @property + def is_numpy(self) -> bool: + """True, if the data in this episode is already stored as numpy arrays.""" + is_numpy = next(iter(self.agent_episodes.values())).is_numpy + # Make sure that all single agent's episodes' `is_numpy` flags are the same. + if not all(eps.is_numpy is is_numpy for eps in self.agent_episodes.values()): + raise RuntimeError( + f"Only some SingleAgentEpisode objects in {self} are converted to " + f"numpy, others are not!" + ) + return is_numpy + + @property + def is_done(self): + """Whether the episode is actually done (terminated or truncated). + + A done episode cannot be continued via `self.add_env_step()` or being + concatenated on its right-side with another episode chunk or being + succeeded via `self.cut()`. + + Note that in a multi-agent environment this does not necessarily + correspond to single agents having terminated or being truncated. + + `self.is_terminated` should be `True`, if all agents are terminated and + `self.is_truncated` should be `True`, if all agents are truncated. If + only one or more (but not all!) agents are `terminated/truncated the + `MultiAgentEpisode.is_terminated/is_truncated` should be `False`. This + information about single agent's terminated/truncated states can always + be retrieved from the `SingleAgentEpisode`s inside the 'MultiAgentEpisode` + one. + + If all agents are either terminated or truncated, but in a mixed fashion, + i.e. some are terminated and others are truncated: This is currently + undefined and could potentially be a problem (if a user really implemented + such a multi-agent env that behaves this way). + + Returns: + Boolean defining if an episode has either terminated or truncated. + """ + return self.is_terminated or self.is_truncated + + def to_numpy(self) -> "MultiAgentEpisode": + """Converts this Episode's list attributes to numpy arrays. + + This means in particular that this episodes' lists (per single agent) of + (possibly complex) data (e.g. an agent having a dict obs space) will be + converted to (possibly complex) structs, whose leafs are now numpy arrays. + Each of these leaf numpy arrays will have the same length (batch dimension) + as the length of the original lists. + + Note that Columns.INFOS are NEVER numpy'ized and will remain a list + (normally, a list of the original, env-returned dicts). This is due to the + heterogeneous nature of INFOS returned by envs, which would make it unwieldy to + convert this information to numpy arrays. + + After calling this method, no further data may be added to this episode via + the `self.add_env_step()` method. + + Examples: + + .. testcode:: + + import numpy as np + + from ray.rllib.env.multi_agent_episode import MultiAgentEpisode + from ray.rllib.env.tests.test_multi_agent_episode import ( + TestMultiAgentEpisode + ) + + # Create some multi-agent episode data. + ( + observations, + actions, + rewards, + terminateds, + truncateds, + infos, + ) = TestMultiAgentEpisode._mock_multi_agent_records() + # Define the agent ids. + agent_ids = ["agent_1", "agent_2", "agent_3", "agent_4", "agent_5"] + + episode = MultiAgentEpisode( + observations=observations, + infos=infos, + actions=actions, + rewards=rewards, + # Note: terminated/truncated have nothing to do with an episode + # being converted `to_numpy` or not (via the `self.to_numpy()` method)! + terminateds=terminateds, + truncateds=truncateds, + len_lookback_buffer=0, # no lookback; all data is actually "in" episode + ) + + # Episode has not been numpy'ized yet. + assert not episode.is_numpy + # We are still operating on lists. + assert ( + episode.get_observations( + indices=[1], + agent_ids="agent_1", + ) == {"agent_1": [1]} + ) + + # Numpy'ized the episode. + episode.to_numpy() + assert episode.is_numpy + + # Everything is now numpy arrays (with 0-axis of size + # B=[len of requested slice]). + assert ( + isinstance(episode.get_observations( + indices=[1], + agent_ids="agent_1", + )["agent_1"], np.ndarray) + ) + + Returns: + This `MultiAgentEpisode` object with the converted numpy data. + """ + + for agent_id, agent_eps in self.agent_episodes.copy().items(): + agent_eps.to_numpy() + + return self + + def concat_episode(self, other: "MultiAgentEpisode") -> None: + """Adds the given `other` MultiAgentEpisode to the right side of `self`. + + In order for this to work, both chunks (`self` and `other`) must fit + together. This is checked by the IDs (must be identical), the time step counters + (`self.env_t` must be the same as `episode_chunk.env_t_started`), as well as the + observations/infos of the individual agents at the concatenation boundaries. + Also, `self.is_done` must not be True, meaning `self.is_terminated` and + `self.is_truncated` are both False. + + Args: + other: The other `MultiAgentEpisode` to be concatenated to this one. + + Returns: + A `MultiAgentEpisode` instance containing the concatenated data + from both episodes (`self` and `other`). + """ + # Make sure the IDs match. + assert other.id_ == self.id_ + # NOTE (sven): This is what we agreed on. As the replay buffers must be + # able to concatenate. + assert not self.is_done + # Make sure the timesteps match. + assert self.env_t == other.env_t_started + # Validate `other`. + other.validate() + + # Concatenate the individual SingleAgentEpisodes from both chunks. + all_agent_ids = set(self.agent_ids) | set(other.agent_ids) + for agent_id in all_agent_ids: + sa_episode = self.agent_episodes.get(agent_id) + + # If agent is only in the new episode chunk -> Store all the data of `other` + # wrt agent in `self`. + if sa_episode is None: + self.agent_episodes[agent_id] = other.agent_episodes[agent_id] + self.env_t_to_agent_t[agent_id] = other.env_t_to_agent_t[agent_id] + self.agent_t_started[agent_id] = other.agent_t_started[agent_id] + self._copy_hanging(agent_id, other) + + # If the agent was done in `self`, ignore and continue. There should not be + # any data of that agent in `other`. + elif sa_episode.is_done: + continue + + # If the agent has data in both chunks, concatenate on the single-agent + # level, thereby making sure the hanging values (begin and end) match. + elif agent_id in other.agent_episodes: + # If `other` has hanging (end) values -> Add these to `self`'s agent + # SingleAgentEpisode (as a new timestep) and only then concatenate. + # Otherwise, the concatentaion would fail b/c of missing data. + if agent_id in self._hanging_actions_end: + assert agent_id in self._hanging_extra_model_outputs_end + sa_episode.add_env_step( + observation=other.agent_episodes[agent_id].get_observations(0), + infos=other.agent_episodes[agent_id].get_infos(0), + action=self._hanging_actions_end[agent_id], + reward=( + self._hanging_rewards_end[agent_id] + + other._hanging_rewards_begin[agent_id] + ), + extra_model_outputs=( + self._hanging_extra_model_outputs_end[agent_id] + ), + ) + sa_episode.concat_episode(other.agent_episodes[agent_id]) + # Override `self`'s hanging (end) values with `other`'s hanging (end). + if agent_id in other._hanging_actions_end: + self._hanging_actions_end[agent_id] = copy.deepcopy( + other._hanging_actions_end[agent_id] + ) + self._hanging_rewards_end[agent_id] = other._hanging_rewards_end[ + agent_id + ] + self._hanging_extra_model_outputs_end[agent_id] = copy.deepcopy( + other._hanging_extra_model_outputs_end[agent_id] + ) + + # Concatenate the env- to agent-timestep mappings. + j = self.env_t + for i, val in enumerate(other.env_t_to_agent_t[agent_id][1:]): + if val == self.SKIP_ENV_TS_TAG: + self.env_t_to_agent_t[agent_id].append(self.SKIP_ENV_TS_TAG) + else: + self.env_t_to_agent_t[agent_id].append(i + 1 + j) + + # Otherwise, the agent is only in `self` and not done. All data is stored + # already -> skip + # else: pass + + # Update all timestep counters. + self.env_t = other.env_t + # Check, if the episode is terminated or truncated. + if other.is_terminated: + self.is_terminated = True + elif other.is_truncated: + self.is_truncated = True + + # Merge with `other`'s custom_data, but give `other` priority b/c we assume + # that as a follow-up chunk of `self` other has a more complete version of + # `custom_data`. + self.custom_data.update(other.custom_data) + + # Validate. + self.validate() + + def cut(self, len_lookback_buffer: int = 0) -> "MultiAgentEpisode": + """Returns a successor episode chunk (of len=0) continuing from this Episode. + + The successor will have the same ID as `self`. + If no lookback buffer is requested (len_lookback_buffer=0), the successor's + observations will be the last observation(s) of `self` and its length will + therefore be 0 (no further steps taken yet). If `len_lookback_buffer` > 0, + the returned successor will have `len_lookback_buffer` observations (and + actions, rewards, etc..) taken from the right side (end) of `self`. For example + if `len_lookback_buffer=2`, the returned successor's lookback buffer actions + will be identical to the results of `self.get_actions([-2, -1])`. + + This method is useful if you would like to discontinue building an episode + chunk (b/c you have to return it from somewhere), but would like to have a new + episode instance to continue building the actual gym.Env episode at a later + time. Vie the `len_lookback_buffer` argument, the continuing chunk (successor) + will still be able to "look back" into this predecessor episode's data (at + least to some extend, depending on the value of `len_lookback_buffer`). + + Args: + len_lookback_buffer: The number of environment timesteps to take along into + the new chunk as "lookback buffer". A lookback buffer is additional data + on the left side of the actual episode data for visibility purposes + (but without actually being part of the new chunk). For example, if + `self` ends in actions: agent_1=5,6,7 and agent_2=6,7, and we call + `self.cut(len_lookback_buffer=2)`, the returned chunk will have + actions 6 and 7 for both agents already in it, but still + `t_started`==t==8 (not 7!) and a length of 0. If there is not enough + data in `self` yet to fulfil the `len_lookback_buffer` request, the + value of `len_lookback_buffer` is automatically adjusted (lowered). + + Returns: + The successor Episode chunk of this one with the same ID and state and the + only observation being the last observation in self. + """ + assert len_lookback_buffer >= 0 + if self.is_done: + raise RuntimeError( + "Can't call `MultiAgentEpisode.cut()` when the episode is already done!" + ) + + # If there is hanging data (e.g. actions) in the agents' caches, we might have + # to re-adjust the lookback len further into the past to make sure that these + # agents have at least one observation to look back to. Otherwise, the timestep + # that got cut into will be "lost" for learning from it. + orig_len_lb = len_lookback_buffer + for agent_id, agent_actions in self._hanging_actions_end.items(): + assert self.env_t_to_agent_t[agent_id].get(-1) == self.SKIP_ENV_TS_TAG + for i in range(orig_len_lb, len(self.env_t_to_agent_t[agent_id].data) + 1): + if self.env_t_to_agent_t[agent_id].get(-i) != self.SKIP_ENV_TS_TAG: + len_lookback_buffer = max(len_lookback_buffer, i - 1) + break + + # Initialize this episode chunk with the most recent observations + # and infos (even if lookback is zero). Similar to an initial `env.reset()` + indices_obs_and_infos = slice(-len_lookback_buffer - 1, None) + indices_rest = ( + slice(-len_lookback_buffer, None) + if len_lookback_buffer > 0 + else slice(None, 0) # -> empty slice + ) + + observations = self.get_observations( + indices=indices_obs_and_infos, return_list=True + ) + infos = self.get_infos(indices=indices_obs_and_infos, return_list=True) + actions = self.get_actions(indices=indices_rest, return_list=True) + rewards = self.get_rewards(indices=indices_rest, return_list=True) + extra_model_outputs = self.get_extra_model_outputs( + key=None, # all keys + indices=indices_rest, + return_list=True, + ) + + successor = MultiAgentEpisode( + # Same ID. + id_=self.id_, + observations=observations, + observation_space=self.observation_space, + infos=infos, + actions=actions, + action_space=self.action_space, + rewards=rewards, + # List of MADicts, mapping agent IDs to their respective extra model output + # dicts. + extra_model_outputs=extra_model_outputs, + terminateds=self.get_terminateds(), + truncateds=self.get_truncateds(), + # Continue with `self`'s current timesteps. + env_t_started=self.env_t, + agent_t_started={ + aid: self.agent_episodes[aid].t + for aid in self.agent_ids + if not self.agent_episodes[aid].is_done + }, + # Same AgentIDs and SingleAgentEpisode IDs. + agent_episode_ids=self.agent_episode_ids, + agent_module_ids={ + aid: self.agent_episodes[aid].module_id for aid in self.agent_ids + }, + agent_to_module_mapping_fn=self.agent_to_module_mapping_fn, + # All data we provided to the c'tor goes into the lookback buffer. + len_lookback_buffer="auto", + ) + + # Copy over the hanging (end) values into the hanging (begin) caches of the + # successor. + successor._hanging_rewards_begin = self._hanging_rewards_end.copy() + + # Deepcopy all custom data in `self` to be continued in the cut episode. + successor._custom_data = copy.deepcopy(self.custom_data) + + return successor + + @property + def agent_ids(self) -> Set[AgentID]: + """Returns the agent ids.""" + return set(self.agent_episodes.keys()) + + @property + def agent_episode_ids(self) -> MultiAgentDict: + """Returns ids from each agent's `SingleAgentEpisode`.""" + + return { + agent_id: agent_eps.id_ + for agent_id, agent_eps in self.agent_episodes.items() + } + + def module_for(self, agent_id: AgentID) -> Optional[ModuleID]: + """Returns the ModuleID for a given AgentID. + + Forces the agent-to-module mapping to be performed (via + `self.agent_to_module_mapping_fn`), if this has not been done yet. + Note that all such mappings are stored in the `self._agent_to_module_mapping` + property. + + Args: + agent_id: The AgentID to get a mapped ModuleID for. + + Returns: + The ModuleID mapped to from the given `agent_id`. + """ + if agent_id not in self._agent_to_module_mapping: + module_id = self._agent_to_module_mapping[ + agent_id + ] = self.agent_to_module_mapping_fn(agent_id, self) + return module_id + else: + return self._agent_to_module_mapping[agent_id] + + def get_observations( + self, + indices: Optional[Union[int, List[int], slice]] = None, + agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None, + *, + env_steps: bool = True, + # global_indices: bool = False, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + one_hot_discrete: bool = False, + return_list: bool = False, + ) -> Union[MultiAgentDict, List[MultiAgentDict]]: + """Returns agents' observations or batched ranges thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual observation stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual observations in a batch of size len(indices). + A slice object is interpreted as a range of observations to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all observations (from ts=0 to the end). + agent_ids: An optional collection of AgentIDs or a single AgentID to get + observations for. If None, will return observations for all agents in + this episode. + env_steps: Whether `indices` should be interpreted as environment time steps + (True) or per-agent timesteps (False). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with agent A's observations + [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range + (ts=0 item is 7), will respond to `get_observations(-1, agent_ids=[A], + neg_index_as_lookback=True)` with {A: `6`} and to + `get_observations(slice(-2, 1), agent_ids=[A], + neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with agent A' observations [10, 11, 12, 13, 14] + and lookback buffer size of 2 (meaning observations `10` and `11` are + part of the lookback buffer) will respond to + `get_observations(slice(-7, -2), agent_ids=[A], fill=0.0)` with + `{A: [0.0, 0.0, 10, 11, 12]}`. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) observation + space that are Discrete or MultiDiscrete. Note that if `fill=0` and the + requested `indices` are out of the range of our data, the returned + one-hot vectors will actually be zero-hot (all slots zero). + return_list: Whether to return a list of multi-agent dicts (instead of + a single multi-agent dict of lists/structs). False by default. This + option can only be used when `env_steps` is True due to the fact the + such a list can only be interpreted as one env step per list item + (would not work with agent steps). + + Returns: + A dictionary mapping agent IDs to observations (at the given + `indices`). If `env_steps` is True, only agents that have stepped + (were ready) at the given env step `indices` are returned (i.e. not all + agent IDs are necessarily in the keys). + If `return_list` is True, returns a list of MultiAgentDicts (mapping agent + IDs to observations) instead. + """ + return self._get( + what="observations", + indices=indices, + agent_ids=agent_ids, + env_steps=env_steps, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + one_hot_discrete=one_hot_discrete, + return_list=return_list, + ) + + def get_infos( + self, + indices: Optional[Union[int, List[int], slice]] = None, + agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None, + *, + env_steps: bool = True, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + return_list: bool = False, + ) -> Union[MultiAgentDict, List[MultiAgentDict]]: + """Returns agents' info dicts or list (ranges) thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual info dict stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual info dicts in a list of size len(indices). + A slice object is interpreted as a range of info dicts to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all infos (from ts=0 to the end). + agent_ids: An optional collection of AgentIDs or a single AgentID to get + info dicts for. If None, will return info dicts for all agents in + this episode. + env_steps: Whether `indices` should be interpreted as environment time steps + (True) or per-agent timesteps (False). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with agent A's info dicts + [{"l":4}, {"l":5}, {"l":6}, {"a":7}, {"b":8}, {"c":9}], where the + first 3 items are the lookback buffer (ts=0 item is {"a": 7}), will + respond to `get_infos(-1, agent_ids=A, neg_index_as_lookback=True)` + with `{A: {"l":6}}` and to + `get_infos(slice(-2, 1), agent_ids=A, neg_index_as_lookback=True)` + with `{A: [{"l":5}, {"l":6}, {"a":7}]}`. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to + auto-fill. For example, an episode with agent A's infos being + [{"l":10}, {"l":11}, {"a":12}, {"b":13}, {"c":14}] and lookback buffer + size of 2 (meaning infos {"l":10}, {"l":11} are part of the lookback + buffer) will respond to `get_infos(slice(-7, -2), agent_ids=A, + fill={"o": 0.0})` with + `{A: [{"o":0.0}, {"o":0.0}, {"l":10}, {"l":11}, {"a":12}]}`. + return_list: Whether to return a list of multi-agent dicts (instead of + a single multi-agent dict of lists/structs). False by default. This + option can only be used when `env_steps` is True due to the fact the + such a list can only be interpreted as one env step per list item + (would not work with agent steps). + + Returns: + A dictionary mapping agent IDs to observations (at the given + `indices`). If `env_steps` is True, only agents that have stepped + (were ready) at the given env step `indices` are returned (i.e. not all + agent IDs are necessarily in the keys). + If `return_list` is True, returns a list of MultiAgentDicts (mapping agent + IDs to infos) instead. + """ + return self._get( + what="infos", + indices=indices, + agent_ids=agent_ids, + env_steps=env_steps, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + return_list=return_list, + ) + + def get_actions( + self, + indices: Optional[Union[int, List[int], slice]] = None, + agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None, + *, + env_steps: bool = True, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + one_hot_discrete: bool = False, + return_list: bool = False, + ) -> Union[MultiAgentDict, List[MultiAgentDict]]: + """Returns agents' actions or batched ranges thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual actions stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual actions in a batch of size len(indices). + A slice object is interpreted as a range of actions to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all actions (from ts=0 to the end). + agent_ids: An optional collection of AgentIDs or a single AgentID to get + actions for. If None, will return actions for all agents in + this episode. + env_steps: Whether `indices` should be interpreted as environment time steps + (True) or per-agent timesteps (False). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with agent A's actions + [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range + (ts=0 item is 7), will respond to `get_actions(-1, agent_ids=[A], + neg_index_as_lookback=True)` with {A: `6`} and to + `get_actions(slice(-2, 1), agent_ids=[A], + neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with agent A' actions [10, 11, 12, 13, 14] + and lookback buffer size of 2 (meaning actions `10` and `11` are + part of the lookback buffer) will respond to + `get_actions(slice(-7, -2), agent_ids=[A], fill=0.0)` with + `{A: [0.0, 0.0, 10, 11, 12]}`. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) observation + space that are Discrete or MultiDiscrete. Note that if `fill=0` and the + requested `indices` are out of the range of our data, the returned + one-hot vectors will actually be zero-hot (all slots zero). + return_list: Whether to return a list of multi-agent dicts (instead of + a single multi-agent dict of lists/structs). False by default. This + option can only be used when `env_steps` is True due to the fact the + such a list can only be interpreted as one env step per list item + (would not work with agent steps). + + Returns: + A dictionary mapping agent IDs to actions (at the given + `indices`). If `env_steps` is True, only agents that have stepped + (were ready) at the given env step `indices` are returned (i.e. not all + agent IDs are necessarily in the keys). + If `return_list` is True, returns a list of MultiAgentDicts (mapping agent + IDs to actions) instead. + """ + return self._get( + what="actions", + indices=indices, + agent_ids=agent_ids, + env_steps=env_steps, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + one_hot_discrete=one_hot_discrete, + return_list=return_list, + ) + + def get_rewards( + self, + indices: Optional[Union[int, List[int], slice]] = None, + agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None, + *, + env_steps: bool = True, + neg_index_as_lookback: bool = False, + fill: Optional[float] = None, + return_list: bool = False, + ) -> Union[MultiAgentDict, List[MultiAgentDict]]: + """Returns agents' rewards or batched ranges thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual rewards stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual rewards in a batch of size len(indices). + A slice object is interpreted as a range of rewards to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all rewards (from ts=0 to the end). + agent_ids: An optional collection of AgentIDs or a single AgentID to get + rewards for. If None, will return rewards for all agents in + this episode. + env_steps: Whether `indices` should be interpreted as environment time steps + (True) or per-agent timesteps (False). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with agent A's rewards + [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range + (ts=0 item is 7), will respond to `get_rewards(-1, agent_ids=[A], + neg_index_as_lookback=True)` with {A: `6`} and to + `get_rewards(slice(-2, 1), agent_ids=[A], + neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}. + fill: An optional float value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with agent A' rewards [10, 11, 12, 13, 14] + and lookback buffer size of 2 (meaning rewards `10` and `11` are + part of the lookback buffer) will respond to + `get_rewards(slice(-7, -2), agent_ids=[A], fill=0.0)` with + `{A: [0.0, 0.0, 10, 11, 12]}`. + return_list: Whether to return a list of multi-agent dicts (instead of + a single multi-agent dict of lists/structs). False by default. This + option can only be used when `env_steps` is True due to the fact the + such a list can only be interpreted as one env step per list item + (would not work with agent steps). + + Returns: + A dictionary mapping agent IDs to rewards (at the given + `indices`). If `env_steps` is True, only agents that have stepped + (were ready) at the given env step `indices` are returned (i.e. not all + agent IDs are necessarily in the keys). + If `return_list` is True, returns a list of MultiAgentDicts (mapping agent + IDs to rewards) instead. + """ + return self._get( + what="rewards", + indices=indices, + agent_ids=agent_ids, + env_steps=env_steps, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + return_list=return_list, + ) + + def get_extra_model_outputs( + self, + key: Optional[str] = None, + indices: Optional[Union[int, List[int], slice]] = None, + agent_ids: Optional[Union[Collection[AgentID], AgentID]] = None, + *, + env_steps: bool = True, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + return_list: bool = False, + ) -> Union[MultiAgentDict, List[MultiAgentDict]]: + """Returns agents' actions or batched ranges thereof from this episode. + + Args: + key: The `key` within each agents' extra_model_outputs dict to extract + data for. If None, return data of all extra model output keys. + indices: A single int is interpreted as an index, from which to return the + individual extra model outputs stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual extra model outputs in a batch of size len(indices). + A slice object is interpreted as a range of extra model outputs to be + returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all extra model outputs (from ts=0 to the end). + agent_ids: An optional collection of AgentIDs or a single AgentID to get + extra model outputs for. If None, will return extra model outputs for + all agents in this episode. + env_steps: Whether `indices` should be interpreted as environment time steps + (True) or per-agent timesteps (False). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with agent A's actions + [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the lookback buffer range + (ts=0 item is 7), will respond to `get_actions(-1, agent_ids=[A], + neg_index_as_lookback=True)` with {A: `6`} and to + `get_actions(slice(-2, 1), agent_ids=[A], + neg_index_as_lookback=True)` with {A: `[5, 6, 7]`}. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with agent A' actions [10, 11, 12, 13, 14] + and lookback buffer size of 2 (meaning actions `10` and `11` are + part of the lookback buffer) will respond to + `get_actions(slice(-7, -2), agent_ids=[A], fill=0.0)` with + `{A: [0.0, 0.0, 10, 11, 12]}`. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) observation + space that are Discrete or MultiDiscrete. Note that if `fill=0` and the + requested `indices` are out of the range of our data, the returned + one-hot vectors will actually be zero-hot (all slots zero). + return_list: Whether to return a list of multi-agent dicts (instead of + a single multi-agent dict of lists/structs). False by default. This + option can only be used when `env_steps` is True due to the fact the + such a list can only be interpreted as one env step per list item + (would not work with agent steps). + + Returns: + A dictionary mapping agent IDs to actions (at the given + `indices`). If `env_steps` is True, only agents that have stepped + (were ready) at the given env step `indices` are returned (i.e. not all + agent IDs are necessarily in the keys). + If `return_list` is True, returns a list of MultiAgentDicts (mapping agent + IDs to extra_model_outputs) instead. + """ + return self._get( + what="extra_model_outputs", + extra_model_outputs_key=key, + indices=indices, + agent_ids=agent_ids, + env_steps=env_steps, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + return_list=return_list, + ) + + def get_terminateds(self) -> MultiAgentDict: + """Gets the terminateds at given indices.""" + terminateds = { + agent_id: self.agent_episodes[agent_id].is_terminated + for agent_id in self.agent_ids + } + terminateds.update({"__all__": self.is_terminated}) + return terminateds + + def get_truncateds(self) -> MultiAgentDict: + truncateds = { + agent_id: self.agent_episodes[agent_id].is_truncated + for agent_id in self.agent_ids + } + truncateds.update({"__all__": self.is_terminated}) + return truncateds + + def slice( + self, + slice_: slice, + *, + len_lookback_buffer: Optional[int] = None, + ) -> "MultiAgentEpisode": + """Returns a slice of this episode with the given slice object. + + Works analogous to + :py:meth:`~ray.rllib.env.single_agent_episode.SingleAgentEpisode.slice` + + However, the important differences are: + - `slice_` is provided in (global) env steps, not agent steps. + - In case `slice_` ends - for a certain agent - in an env step, where that + particular agent does not have an observation, the previous observation will + be included, but the next action and sum of rewards until this point will + be stored in the agent's hanging values caches for the returned + MultiAgentEpisode slice. + + .. testcode:: + + from ray.rllib.env.multi_agent_episode import MultiAgentEpisode + from ray.rllib.utils.test_utils import check + + # Generate a simple multi-agent episode. + observations = [ + {"a0": 0, "a1": 0}, # 0 + { "a1": 1}, # 1 + { "a1": 2}, # 2 + {"a0": 3, "a1": 3}, # 3 + {"a0": 4}, # 4 + ] + # Actions are the same as observations (except for last obs, which doesn't + # have an action). + actions = observations[:-1] + # Make up a reward for each action. + rewards = [ + {aid: r / 10 + 0.1 for aid, r in o.items()} + for o in observations + ] + episode = MultiAgentEpisode( + observations=observations, + actions=actions, + rewards=rewards, + len_lookback_buffer=0, + ) + + # Slice the episode and check results. + slice = episode[1:3] + a0 = slice.agent_episodes["a0"] + a1 = slice.agent_episodes["a1"] + check((a0.observations, a1.observations), ([3], [1, 2, 3])) + check((a0.actions, a1.actions), ([], [1, 2])) + check((a0.rewards, a1.rewards), ([], [0.2, 0.3])) + check((a0.is_done, a1.is_done), (False, False)) + + # If a slice ends in a "gap" for an agent, expect actions and rewards to be + # cached for this agent. + slice = episode[:2] + a0 = slice.agent_episodes["a0"] + check(a0.observations, [0]) + check(a0.actions, []) + check(a0.rewards, []) + check(slice._hanging_actions_end["a0"], 0) + check(slice._hanging_rewards_end["a0"], 0.1) + + Args: + slice_: The slice object to use for slicing. This should exclude the + lookback buffer, which will be prepended automatically to the returned + slice. + len_lookback_buffer: If not None, forces the returned slice to try to have + this number of timesteps in its lookback buffer (if available). If None + (default), tries to make the returned slice's lookback as large as the + current lookback buffer of this episode (`self`). + + Returns: + The new MultiAgentEpisode representing the requested slice. + """ + if slice_.step not in [1, None]: + raise NotImplementedError( + "Slicing MultiAgentEnv with a step other than 1 (you used" + f" {slice_.step}) is not supported!" + ) + + # Translate `slice_` into one that only contains 0-or-positive ints and will + # NOT contain any None. + start = slice_.start + stop = slice_.stop + + # Start is None -> 0. + if start is None: + start = 0 + # Start is negative -> Interpret index as counting "from end". + elif start < 0: + start = max(len(self) + start, 0) + # Start is larger than len(self) -> Clip to len(self). + elif start > len(self): + start = len(self) + + # Stop is None -> Set stop to our len (one ts past last valid index). + if stop is None: + stop = len(self) + # Stop is negative -> Interpret index as counting "from end". + elif stop < 0: + stop = max(len(self) + stop, 0) + # Stop is larger than len(self) -> Clip to len(self). + elif stop > len(self): + stop = len(self) + + ref_lookback = None + try: + for aid, sa_episode in self.agent_episodes.items(): + if ref_lookback is None: + ref_lookback = sa_episode.observations.lookback + assert sa_episode.observations.lookback == ref_lookback + assert sa_episode.actions.lookback == ref_lookback + assert sa_episode.rewards.lookback == ref_lookback + assert all( + ilb.lookback == ref_lookback + for ilb in sa_episode.extra_model_outputs.values() + ) + except AssertionError: + raise ValueError( + "Can only slice a MultiAgentEpisode if all lookback buffers in this " + "episode have the exact same size!" + ) + + # Determine terminateds/truncateds and when (in agent timesteps) the + # single-agent episode slices start. + terminateds = {} + truncateds = {} + agent_t_started = {} + for aid, sa_episode in self.agent_episodes.items(): + mapping = self.env_t_to_agent_t[aid] + # If the (agent) timestep directly at the slice stop boundary is equal to + # the length of the single-agent episode of this agent -> Use the + # single-agent episode's terminated/truncated flags. + # If `stop` is already beyond this agent's single-agent episode, then we + # don't have to keep track of this: The MultiAgentEpisode initializer will + # automatically determine that this agent must be done (b/c it has no action + # following its final observation). + if ( + stop < len(mapping) + and mapping[stop] != self.SKIP_ENV_TS_TAG + and len(sa_episode) == mapping[stop] + ): + terminateds[aid] = sa_episode.is_terminated + truncateds[aid] = sa_episode.is_truncated + # Determine this agent's t_started. + if start < len(mapping): + for i in range(start, len(mapping)): + if mapping[i] != self.SKIP_ENV_TS_TAG: + agent_t_started[aid] = sa_episode.t_started + mapping[i] + break + terminateds["__all__"] = all( + terminateds.get(aid) for aid in self.agent_episodes + ) + truncateds["__all__"] = all(truncateds.get(aid) for aid in self.agent_episodes) + + # Determine all other slice contents. + _lb = len_lookback_buffer if len_lookback_buffer is not None else ref_lookback + if start - _lb < 0 and ref_lookback < (_lb - start): + _lb = ref_lookback + start + observations = self.get_observations( + slice(start - _lb, stop + 1), + neg_index_as_lookback=True, + return_list=True, + ) + actions = self.get_actions( + slice(start - _lb, stop), + neg_index_as_lookback=True, + return_list=True, + ) + rewards = self.get_rewards( + slice(start - _lb, stop), + neg_index_as_lookback=True, + return_list=True, + ) + extra_model_outputs = self.get_extra_model_outputs( + indices=slice(start - _lb, stop), + neg_index_as_lookback=True, + return_list=True, + ) + + # Create the actual slice to be returned. + ma_episode = MultiAgentEpisode( + id_=self.id_, + # In the following, offset `start`s automatically by lookbacks. + observations=observations, + observation_space=self.observation_space, + actions=actions, + action_space=self.action_space, + rewards=rewards, + extra_model_outputs=extra_model_outputs, + terminateds=terminateds, + truncateds=truncateds, + len_lookback_buffer=_lb, + env_t_started=self.env_t_started + start, + agent_episode_ids={ + aid: eid.id_ for aid, eid in self.agent_episodes.items() + }, + agent_t_started=agent_t_started, + agent_module_ids=self._agent_to_module_mapping, + agent_to_module_mapping_fn=self.agent_to_module_mapping_fn, + ) + + # Numpy'ize slice if `self` is also finalized. + if self.is_numpy: + ma_episode.to_numpy() + + return ma_episode + + def __len__(self): + """Returns the length of an `MultiAgentEpisode`. + + Note that the length of an episode is defined by the difference + between its actual timestep and the starting point. + + Returns: An integer defining the length of the episode or an + error if the episode has not yet started. + """ + return self.env_t - self.env_t_started + + def __repr__(self): + sa_eps_returns = { + aid: sa_eps.get_return() for aid, sa_eps in self.agent_episodes.items() + } + return ( + f"MAEps(len={len(self)} done={self.is_done} " + f"Rs={sa_eps_returns} id_={self.id_})" + ) + + def print(self) -> None: + """Prints this MultiAgentEpisode as a table of observations for the agents.""" + + # Find the maximum timestep across all agents to determine the grid width. + max_ts = max(ts.len_incl_lookback() for ts in self.env_t_to_agent_t.values()) + lookback = next(iter(self.env_t_to_agent_t.values())).lookback + longest_agent = max(len(aid) for aid in self.agent_ids) + # Construct the header. + header = ( + "ts" + + (" " * longest_agent) + + " ".join(str(i) for i in range(-lookback, max_ts - lookback)) + + "\n" + ) + # Construct each agent's row. + rows = [] + for agent, inf_buffer in self.env_t_to_agent_t.items(): + row = f"{agent} " + (" " * (longest_agent - len(agent))) + for t in inf_buffer.data: + # Two spaces for alignment. + if t == "S": + row += " " + # Mark the step with an x. + else: + row += " x " + # Remove trailing space for alignment. + rows.append(row.rstrip()) + + # Join all components into a final string + print(header + "\n".join(rows)) + + def get_state(self) -> Dict[str, Any]: + """Returns the state of a multi-agent episode. + + Note that from an episode's state the episode itself can + be recreated. + + Returns: A dicitonary containing pickable data for a + `MultiAgentEpisode`. + """ + return { + "id_": self.id_, + "agent_to_module_mapping_fn": self.agent_to_module_mapping_fn, + "_agent_to_module_mapping": self._agent_to_module_mapping, + "observation_space": self.observation_space, + "action_space": self.action_space, + "env_t_started": self.env_t_started, + "env_t": self.env_t, + "agent_t_started": self.agent_t_started, + # TODO (simon): Check, if we can store the `InfiniteLookbackBuffer` + "env_t_to_agent_t": self.env_t_to_agent_t, + "_hanging_actions_end": self._hanging_actions_end, + "_hanging_extra_model_outputs_end": self._hanging_extra_model_outputs_end, + "_hanging_rewards_end": self._hanging_rewards_end, + "_hanging_rewards_begin": self._hanging_rewards_begin, + "is_terminated": self.is_terminated, + "is_truncated": self.is_truncated, + "agent_episodes": list( + { + agent_id: agent_eps.get_state() + for agent_id, agent_eps in self.agent_episodes.items() + }.items() + ), + "_start_time": self._start_time, + "_last_step_time": self._last_step_time, + "custom_data": self.custom_data, + } + + @staticmethod + def from_state(state: Dict[str, Any]) -> "MultiAgentEpisode": + """Creates a multi-agent episode from a state dictionary. + + See `MultiAgentEpisode.get_state()` for creating a state for + a `MultiAgentEpisode` pickable state. For recreating a + `MultiAgentEpisode` from a state, this state has to be complete, + i.e. all data must have been stored in the state. + + Args: + state: A dict containing all data required to recreate a MultiAgentEpisode`. + See `MultiAgentEpisode.get_state()`. + + Returns: + A `MultiAgentEpisode` instance created from the state data. + """ + # Create an empty `MultiAgentEpisode` instance. + episode = MultiAgentEpisode(id_=state["id_"]) + # Fill the instance with the state data. + episode.agent_to_module_mapping_fn = state["agent_to_module_mapping_fn"] + episode._agent_to_module_mapping = state["_agent_to_module_mapping"] + episode.observation_space = state["observation_space"] + episode.action_space = state["action_space"] + episode.env_t_started = state["env_t_started"] + episode.env_t = state["env_t"] + episode.agent_t_started = state["agent_t_started"] + episode.env_t_to_agent_t = state["env_t_to_agent_t"] + episode._hanging_actions_end = state["_hanging_actions_end"] + episode._hanging_extra_model_outputs_end = state[ + "_hanging_extra_model_outputs_end" + ] + episode._hanging_rewards_end = state["_hanging_rewards_end"] + episode._hanging_rewards_begin = state["_hanging_rewards_begin"] + episode.is_terminated = state["is_terminated"] + episode.is_truncated = state["is_truncated"] + episode.agent_episodes = { + agent_id: SingleAgentEpisode.from_state(agent_state) + for agent_id, agent_state in state["agent_episodes"] + } + episode._start_time = state["_start_time"] + episode._last_step_time = state["_last_step_time"] + episode._custom_data = state.get("custom_data", {}) + + # Validate the episode. + episode.validate() + + return episode + + def get_sample_batch(self) -> MultiAgentBatch: + """Converts this `MultiAgentEpisode` into a `MultiAgentBatch`. + + Each `SingleAgentEpisode` instances in `MultiAgentEpisode.agent_epiosdes` + will be converted into a `SampleBatch` and the environment timestep will be + passed as the returned MultiAgentBatch's `env_steps`. + + Returns: + A MultiAgentBatch containing all of this episode's data. + """ + # TODO (simon): Check, if timesteps should be converted into global + # timesteps instead of agent steps. + # Note, only agents that have stepped are included into the batch. + return MultiAgentBatch( + policy_batches={ + agent_id: agent_eps.get_sample_batch() + for agent_id, agent_eps in self.agent_episodes.items() + if agent_eps.t - agent_eps.t_started > 0 + }, + env_steps=self.env_t - self.env_t_started, + ) + + def get_return( + self, + include_hanging_rewards: bool = False, + ) -> float: + """Returns all-agent return. + + Args: + include_hanging_rewards: Whether we should also consider + hanging rewards wehn calculating the overall return. Agents might + have received partial rewards, i.e. rewards without an + observation. These are stored in the "hanging" caches (begin and end) + for each agent and added up until the next observation is received by + that agent. + + Returns: + The sum of all single-agents' returns (maybe including the hanging + rewards per agent). + """ + env_return = sum( + agent_eps.get_return() for agent_eps in self.agent_episodes.values() + ) + if include_hanging_rewards: + for hanging_r in self._hanging_rewards_begin.values(): + env_return += hanging_r + for hanging_r in self._hanging_rewards_end.values(): + env_return += hanging_r + + return env_return + + def get_agents_to_act(self) -> Set[AgentID]: + """Returns a set of agent IDs required to send an action to `env.step()` next. + + Those are generally the agents that received an observation in the most recent + `env.step()` call. + + Returns: + A set of AgentIDs that are supposed to send actions to the next `env.step()` + call. + """ + return { + aid + for aid in self.get_observations(-1).keys() + if not self.agent_episodes[aid].is_done + } + + def get_agents_that_stepped(self) -> Set[AgentID]: + """Returns a set of agent IDs of those agents that just finished stepping. + + These are all the agents that have an observation logged at the last env + timestep, which may include agents, whose single agent episode just terminated + or truncated. + + Returns: + A set of AgentIDs of those agents that just finished stepping (that have a + most recent observation on the env timestep scale), regardless of whether + their single agent episodes are done or not. + """ + return set(self.get_observations(-1).keys()) + + def get_duration_s(self) -> float: + """Returns the duration of this Episode (chunk) in seconds.""" + if self._last_step_time is None: + return 0.0 + return self._last_step_time - self._start_time + + def set_observations( + self, + *, + new_data: MultiAgentDict, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some single-agent Episode's observations with the provided data. + + This is a helper method to batch `SingleAgentEpisode.set_observations`. + For more detail, see `SingleAgentEpisode.set_observations`. + + Args: + new_data: A dict mapping agent IDs to new observation data. + Each value in the dict is the new observation data to overwrite existing data with. + This may be a list of individual observation(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be (possibly complex) struct matching the + observation space and with a batch size of its leafs exactly the size + of the to-be-overwritten slice or segment (provided by `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single observation). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + observations = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_observations(individual_observation, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + observations buffer with the provided "individual_observation". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + for agent_id, new_agent_data in new_data.items(): + if agent_id not in self.agent_episodes: + raise KeyError(f"AgentID '{agent_id}' not found in this episode.") + self.agent_episodes[agent_id].set_observations( + new_data=new_agent_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def set_actions( + self, + *, + new_data: MultiAgentDict, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of this Episode's actions with the provided data. + + This is a helper method to batch `SingleAgentEpisode.set_actions`. + For more detail, see `SingleAgentEpisode.set_actions`. + + Args: + new_data: A dict mapping agent IDs to new action data. + Each value in the dict is the new action data to overwrite existing data with. + This may be a list of individual action(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be (possibly complex) struct matching the + action space and with a batch size of its leafs exactly the size + of the to-be-overwritten slice or segment (provided by `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single observation). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + actions = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_actions(individual_action, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + actions buffer with the provided "individual_action". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + for agent_id, new_agent_data in new_data.items(): + if agent_id not in self.agent_episodes: + raise KeyError(f"AgentID '{agent_id}' not found in this episode.") + self.agent_episodes[agent_id].set_actions( + new_data=new_agent_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def set_rewards( + self, + *, + new_data: MultiAgentDict, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of this Episode's rewards with the provided data. + + This is a helper method to batch `SingleAgentEpisode.set_rewards`. + For more detail, see `SingleAgentEpisode.set_rewards`. + + Args: + new_data: A dict mapping agent IDs to new reward data. + Each value in the dict is the new reward data to overwrite existing data with. + This may be a list of individual reward(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be a np.ndarray with a length exactly + the size of the to-be-overwritten slice or segment (provided by + `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single reward). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + rewards = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_rewards(individual_reward, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + rewards buffer with the provided "individual_reward". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + for agent_id, new_agent_data in new_data.items(): + if agent_id not in self.agent_episodes: + raise KeyError(f"AgentID '{agent_id}' not found in this episode.") + self.agent_episodes[agent_id].set_rewards( + new_data=new_agent_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def set_extra_model_outputs( + self, + *, + key, + new_data: MultiAgentDict, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of this Episode's extra model outputs with `new_data`. + + This is a helper method to batch `SingleAgentEpisode.set_extra_model_outputs`. + For more detail, see `SingleAgentEpisode.set_extra_model_outputs`. + + Args: + key: The `key` within `self.extra_model_outputs` to override data on or + to insert as a new key into `self.extra_model_outputs`. + new_data: A dict mapping agent IDs to new extra model outputs data. + Each value in the dict is the new extra model outputs data to overwrite existing data with. + This may be a list of individual reward(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be a np.ndarray with a length exactly + the size of the to-be-overwritten slice or segment (provided by + `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single extra model output). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + extra_model_outputs[key][agent_id] = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_extra_model_outputs(key, individual_output, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + extra_model_outputs[key][agent_id] buffer with the provided "individual_output". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + for agent_id, new_agent_data in new_data.items(): + if agent_id not in self.agent_episodes: + raise KeyError(f"AgentID '{agent_id}' not found in this episode.") + self.agent_episodes[agent_id].set_extra_model_outputs( + key=key, + new_data=new_agent_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def env_steps(self) -> int: + """Returns the number of environment steps. + + Note, this episode instance could be a chunk of an actual episode. + + Returns: + An integer that counts the number of environment steps this episode instance + has seen. + """ + return len(self) + + def agent_steps(self) -> int: + """Number of agent steps. + + Note, there are >= 1 agent steps per environment step. + + Returns: + An integer counting the number of agent steps executed during the time this + episode instance records. + """ + return sum(len(eps) for eps in self.agent_episodes.values()) + + def __getitem__(self, item: slice) -> "MultiAgentEpisode": + """Enable squared bracket indexing- and slicing syntax, e.g. episode[-4:].""" + if isinstance(item, slice): + return self.slice(slice_=item) + else: + raise NotImplementedError( + f"MultiAgentEpisode does not support getting item '{item}'! " + "Only slice objects allowed with the syntax: `episode[a:b]`." + ) + + def _init_single_agent_episodes( + self, + *, + agent_module_ids: Optional[Dict[AgentID, ModuleID]] = None, + agent_episode_ids: Optional[Dict[AgentID, str]] = None, + observations: Optional[List[MultiAgentDict]] = None, + actions: Optional[List[MultiAgentDict]] = None, + rewards: Optional[List[MultiAgentDict]] = None, + infos: Optional[List[MultiAgentDict]] = None, + terminateds: Union[MultiAgentDict, bool] = False, + truncateds: Union[MultiAgentDict, bool] = False, + extra_model_outputs: Optional[List[MultiAgentDict]] = None, + ): + if observations is None: + return + if actions is None: + assert not rewards + assert not extra_model_outputs + actions = [] + rewards = [] + extra_model_outputs = [] + + # Infos and `extra_model_outputs` are allowed to be None -> Fill them with + # proper dummy values, if so. + if infos is None: + infos = [{} for _ in range(len(observations))] + if extra_model_outputs is None: + extra_model_outputs = [{} for _ in range(len(actions))] + + observations_per_agent = defaultdict(list) + infos_per_agent = defaultdict(list) + actions_per_agent = defaultdict(list) + rewards_per_agent = defaultdict(list) + extra_model_outputs_per_agent = defaultdict(list) + done_per_agent = defaultdict(bool) + len_lookback_buffer_per_agent = defaultdict(lambda: self._len_lookback_buffers) + + all_agent_ids = set( + agent_episode_ids.keys() if agent_episode_ids is not None else [] + ) + agent_module_ids = agent_module_ids or {} + + # Step through all observations and interpret these as the (global) env steps. + for data_idx, (obs, inf) in enumerate(zip(observations, infos)): + # If we do have actions/extra outs/rewards for this timestep, use the data. + # It may be that these lists have the same length as the observations list, + # in which case the data will be cached (agent did step/send an action, + # but the step has not been concluded yet by the env). + act = actions[data_idx] if len(actions) > data_idx else {} + extra_outs = ( + extra_model_outputs[data_idx] + if len(extra_model_outputs) > data_idx + else {} + ) + rew = rewards[data_idx] if len(rewards) > data_idx else {} + + for agent_id, agent_obs in obs.items(): + all_agent_ids.add(agent_id) + + observations_per_agent[agent_id].append(agent_obs) + infos_per_agent[agent_id].append(inf.get(agent_id, {})) + + # Pull out hanging action (if not first obs for this agent) and + # complete step for agent. + if len(observations_per_agent[agent_id]) > 1: + actions_per_agent[agent_id].append( + self._hanging_actions_end.pop(agent_id) + ) + extra_model_outputs_per_agent[agent_id].append( + self._hanging_extra_model_outputs_end.pop(agent_id) + ) + rewards_per_agent[agent_id].append( + self._hanging_rewards_end.pop(agent_id) + ) + # First obs for this agent. Make sure the agent's mapping is + # appropriately prepended with self.SKIP_ENV_TS_TAG tags. + else: + if agent_id not in self.env_t_to_agent_t: + self.env_t_to_agent_t[agent_id].extend( + [self.SKIP_ENV_TS_TAG] * data_idx + ) + len_lookback_buffer_per_agent[agent_id] -= data_idx + + # Agent is still continuing (has an action for the next step). + if agent_id in act: + # Always push actions/extra outputs into cache, then remove them + # from there, once the next observation comes in. Same for rewards. + self._hanging_actions_end[agent_id] = act[agent_id] + self._hanging_extra_model_outputs_end[agent_id] = extra_outs.get( + agent_id, {} + ) + self._hanging_rewards_end[agent_id] += rew.get(agent_id, 0.0) + # Agent is done (has no action for the next step). + elif terminateds.get(agent_id) or truncateds.get(agent_id): + done_per_agent[agent_id] = True + # There is more (global) action/reward data. This agent must therefore + # be done. Automatically add it to `done_per_agent` and `terminateds`. + elif data_idx < len(observations) - 1: + done_per_agent[agent_id] = terminateds[agent_id] = True + + # Update env_t_to_agent_t mapping. + self.env_t_to_agent_t[agent_id].append( + len(observations_per_agent[agent_id]) - 1 + ) + + # Those agents that did NOT step: + # - Get self.SKIP_ENV_TS_TAG added to their env_t_to_agent_t mapping. + # - Get their reward (if any) added up. + for agent_id in all_agent_ids: + if agent_id not in obs and agent_id not in done_per_agent: + self.env_t_to_agent_t[agent_id].append(self.SKIP_ENV_TS_TAG) + # If we are still in the global lookback buffer segment, deduct 1 + # from this agents' lookback buffer, b/c we don't want the agent + # to use this (missing) obs/data in its single-agent lookback. + if ( + len(self.env_t_to_agent_t[agent_id]) + - self._len_lookback_buffers + <= 0 + ): + len_lookback_buffer_per_agent[agent_id] -= 1 + self._hanging_rewards_end[agent_id] += rew.get(agent_id, 0.0) + + # - Validate per-agent data. + # - Fix lookback buffers of env_t_to_agent_t mappings. + for agent_id in list(self.env_t_to_agent_t.keys()): + # Skip agent if it doesn't seem to have any data. + if agent_id not in observations_per_agent: + del self.env_t_to_agent_t[agent_id] + continue + assert ( + len(observations_per_agent[agent_id]) + == len(infos_per_agent[agent_id]) + == len(actions_per_agent[agent_id]) + 1 + == len(extra_model_outputs_per_agent[agent_id]) + 1 + == len(rewards_per_agent[agent_id]) + 1 + ) + self.env_t_to_agent_t[agent_id].lookback = self._len_lookback_buffers + + # Now create the individual episodes from the collected per-agent data. + for agent_id, agent_obs in observations_per_agent.items(): + # If agent only has a single obs AND is already done, remove all its traces + # from this MultiAgentEpisode. + if len(agent_obs) == 1 and done_per_agent.get(agent_id): + self._del_agent(agent_id) + continue + + # Try to figure out the module ID for this agent. + # If not provided explicitly by the user that initializes this episode + # object, try our mapping function. + module_id = agent_module_ids.get( + agent_id, self.agent_to_module_mapping_fn(agent_id, self) + ) + # Create this agent's SingleAgentEpisode. + sa_episode = SingleAgentEpisode( + id_=( + agent_episode_ids.get(agent_id) + if agent_episode_ids is not None + else None + ), + agent_id=agent_id, + module_id=module_id, + multi_agent_episode_id=self.id_, + observations=agent_obs, + observation_space=self.observation_space.get(agent_id), + infos=infos_per_agent[agent_id], + actions=actions_per_agent[agent_id], + action_space=self.action_space.get(agent_id), + rewards=rewards_per_agent[agent_id], + extra_model_outputs=( + { + k: [i[k] for i in extra_model_outputs_per_agent[agent_id]] + for k in extra_model_outputs_per_agent[agent_id][0].keys() + } + if extra_model_outputs_per_agent[agent_id] + else None + ), + terminated=terminateds.get(agent_id, False), + truncated=truncateds.get(agent_id, False), + t_started=self.agent_t_started[agent_id], + len_lookback_buffer=max(len_lookback_buffer_per_agent[agent_id], 0), + ) + # .. and store it. + self.agent_episodes[agent_id] = sa_episode + + def _get( + self, + *, + what, + indices, + agent_ids=None, + env_steps=True, + neg_index_as_lookback=False, + fill=None, + one_hot_discrete=False, + return_list=False, + extra_model_outputs_key=None, + ): + agent_ids = set(force_list(agent_ids)) or self.agent_ids + + kwargs = dict( + what=what, + indices=indices, + agent_ids=agent_ids, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + # Rewards and infos do not support one_hot_discrete option. + one_hot_discrete=dict( + {} if not one_hot_discrete else {"one_hot_discrete": one_hot_discrete} + ), + extra_model_outputs_key=extra_model_outputs_key, + ) + + # User specified agent timesteps (indices) -> Simply delegate everything + # to the individual agents' SingleAgentEpisodes. + if env_steps is False: + if return_list: + raise ValueError( + f"`MultiAgentEpisode.get_{what}()` can't be called with both " + "`env_steps=False` and `return_list=True`!" + ) + return self._get_data_by_agent_steps(**kwargs) + # User specified env timesteps (indices) -> We need to translate them for each + # agent into agent-timesteps. + # Return a list of individual per-env-timestep multi-agent dicts. + elif return_list: + return self._get_data_by_env_steps_as_list(**kwargs) + # Return a single multi-agent dict with lists/arrays as leafs. + else: + return self._get_data_by_env_steps(**kwargs) + + def _get_data_by_agent_steps( + self, + *, + what, + indices, + agent_ids, + neg_index_as_lookback, + fill, + one_hot_discrete, + extra_model_outputs_key, + ): + # Return requested data by agent-steps. + ret = {} + # For each agent, we retrieve the data through passing the given indices into + # the SingleAgentEpisode of that agent. + for agent_id, sa_episode in self.agent_episodes.items(): + if agent_id not in agent_ids: + continue + inf_lookback_buffer = getattr(sa_episode, what) + hanging_val = self._get_hanging_value(what, agent_id) + # User wants a specific `extra_model_outputs` key. + if extra_model_outputs_key is not None: + inf_lookback_buffer = inf_lookback_buffer[extra_model_outputs_key] + hanging_val = hanging_val[extra_model_outputs_key] + agent_value = inf_lookback_buffer.get( + indices=indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + _add_last_ts_value=hanging_val, + **one_hot_discrete, + ) + if agent_value is None or agent_value == []: + continue + ret[agent_id] = agent_value + return ret + + def _get_data_by_env_steps_as_list( + self, + *, + what: str, + indices: Union[int, slice, List[int]], + agent_ids: Collection[AgentID], + neg_index_as_lookback: bool, + fill: Any, + one_hot_discrete, + extra_model_outputs_key: str, + ) -> List[MultiAgentDict]: + # Collect indices for each agent first, so we can construct the list in + # the next step. + agent_indices = {} + for agent_id in self.agent_episodes.keys(): + if agent_id not in agent_ids: + continue + agent_indices[agent_id] = self.env_t_to_agent_t[agent_id].get( + indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=self.SKIP_ENV_TS_TAG, + # For those records where there is no "hanging" last timestep (all + # other than obs and infos), we have to ignore the last entry in + # the env_t_to_agent_t mappings. + _ignore_last_ts=what not in ["observations", "infos"], + ) + if not agent_indices: + return [] + ret = [] + for i in range(len(next(iter(agent_indices.values())))): + ret2 = {} + for agent_id, idxes in agent_indices.items(): + hanging_val = self._get_hanging_value(what, agent_id) + ( + inf_lookback_buffer, + indices_to_use, + ) = self._get_inf_lookback_buffer_or_dict( + agent_id, + what, + extra_model_outputs_key, + hanging_val, + filter_for_skip_indices=idxes[i], + ) + if ( + what == "extra_model_outputs" + and not inf_lookback_buffer + and not hanging_val + ): + continue + agent_value = self._get_single_agent_data_by_index( + what=what, + inf_lookback_buffer=inf_lookback_buffer, + agent_id=agent_id, + index_incl_lookback=indices_to_use, + fill=fill, + one_hot_discrete=one_hot_discrete, + extra_model_outputs_key=extra_model_outputs_key, + hanging_val=hanging_val, + ) + if agent_value is not None: + ret2[agent_id] = agent_value + ret.append(ret2) + return ret + + def _get_data_by_env_steps( + self, + *, + what: str, + indices: Union[int, slice, List[int]], + agent_ids: Collection[AgentID], + neg_index_as_lookback: bool, + fill: Any, + one_hot_discrete: bool, + extra_model_outputs_key: str, + ) -> MultiAgentDict: + ignore_last_ts = what not in ["observations", "infos"] + ret = {} + for agent_id, sa_episode in self.agent_episodes.items(): + if agent_id not in agent_ids: + continue + hanging_val = self._get_hanging_value(what, agent_id) + agent_indices = self.env_t_to_agent_t[agent_id].get( + indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=self.SKIP_ENV_TS_TAG if fill is not None else None, + # For those records where there is no "hanging" last timestep (all + # other than obs and infos), we have to ignore the last entry in + # the env_t_to_agent_t mappings. + _ignore_last_ts=ignore_last_ts, + ) + inf_lookback_buffer, agent_indices = self._get_inf_lookback_buffer_or_dict( + agent_id, + what, + extra_model_outputs_key, + hanging_val, + filter_for_skip_indices=agent_indices, + ) + if isinstance(agent_indices, list): + agent_values = self._get_single_agent_data_by_env_step_indices( + what=what, + agent_id=agent_id, + indices_incl_lookback=agent_indices, + fill=fill, + one_hot_discrete=one_hot_discrete, + hanging_val=hanging_val, + extra_model_outputs_key=extra_model_outputs_key, + ) + if len(agent_values) > 0: + ret[agent_id] = agent_values + else: + agent_values = self._get_single_agent_data_by_index( + what=what, + inf_lookback_buffer=inf_lookback_buffer, + agent_id=agent_id, + index_incl_lookback=agent_indices, + fill=fill, + one_hot_discrete=one_hot_discrete, + extra_model_outputs_key=extra_model_outputs_key, + hanging_val=hanging_val, + ) + if agent_values is not None: + ret[agent_id] = agent_values + return ret + + def _get_single_agent_data_by_index( + self, + *, + what: str, + inf_lookback_buffer: InfiniteLookbackBuffer, + agent_id: AgentID, + index_incl_lookback: Union[int, str], + fill: Any, + one_hot_discrete: dict, + extra_model_outputs_key: str, + hanging_val: Any, + ) -> Any: + sa_episode = self.agent_episodes[agent_id] + + if index_incl_lookback == self.SKIP_ENV_TS_TAG: + # We don't want to fill -> Skip this agent. + if fill is None: + return + # Provide filled value for this agent. + return getattr(sa_episode, f"get_{what}")( + indices=1000000000000, + neg_index_as_lookback=False, + fill=fill, + **dict( + {} + if extra_model_outputs_key is None + else {"key": extra_model_outputs_key} + ), + **one_hot_discrete, + ) + + # No skip timestep -> Provide value at given index for this agent. + + # Special case: extra_model_outputs and key=None (return all keys as + # a dict). Note that `inf_lookback_buffer` is NOT an infinite lookback + # buffer, but a dict mapping keys to individual infinite lookback + # buffers. + elif what == "extra_model_outputs" and extra_model_outputs_key is None: + assert hanging_val is None or isinstance(hanging_val, dict) + ret = {} + if inf_lookback_buffer: + for key, sub_buffer in inf_lookback_buffer.items(): + ret[key] = sub_buffer.get( + indices=index_incl_lookback - sub_buffer.lookback, + neg_index_as_lookback=True, + fill=fill, + _add_last_ts_value=( + None if hanging_val is None else hanging_val[key] + ), + **one_hot_discrete, + ) + else: + for key in hanging_val.keys(): + ret[key] = InfiniteLookbackBuffer().get( + indices=index_incl_lookback, + neg_index_as_lookback=True, + fill=fill, + _add_last_ts_value=hanging_val[key], + **one_hot_discrete, + ) + return ret + + # Extract data directly from the infinite lookback buffer object. + else: + return inf_lookback_buffer.get( + indices=index_incl_lookback - inf_lookback_buffer.lookback, + neg_index_as_lookback=True, + fill=fill, + _add_last_ts_value=hanging_val, + **one_hot_discrete, + ) + + def _get_single_agent_data_by_env_step_indices( + self, + *, + what: str, + agent_id: AgentID, + indices_incl_lookback: Union[int, str], + fill: Optional[Any] = None, + one_hot_discrete: bool = False, + extra_model_outputs_key: Optional[str] = None, + hanging_val: Optional[Any] = None, + ) -> Any: + """Returns single data item from the episode based on given (env step) indices. + + The returned data item will have a batch size that matches the env timesteps + defined via `indices_incl_lookback`. + + Args: + what: A (str) descriptor of what data to collect. Must be one of + "observations", "infos", "actions", "rewards", or "extra_model_outputs". + indices_incl_lookback: A list of ints specifying, which indices + to pull from the InfiniteLookbackBuffer defined by `agent_id` and `what` + (and maybe `extra_model_outputs_key`). Note that these indices + disregard the special logic of the lookback buffer. Meaning if one + index in `indices_incl_lookback` is 0, then the first value in the + lookback buffer should be returned, not the first value after the + lookback buffer (which would be normal behavior for pulling items from + an `InfiniteLookbackBuffer` object). + agent_id: The individual agent ID to pull data for. Used to lookup the + `SingleAgentEpisode` object for this agent in `self`. + fill: An optional float value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the buffer's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, a buffer with data [10, 11, 12, 13, 14] and lookback + buffer size of 2 (meaning `10` and `11` are part of the lookback buffer) + will respond to `indices_incl_lookback=[-1, -2, 0]` and `fill=0.0` + with `[0.0, 0.0, 10]`. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) space + that are Discrete or MultiDiscrete. Note that if `fill=0` and the + requested `indices_incl_lookback` are out of the range of our data, the + returned one-hot vectors will actually be zero-hot (all slots zero). + extra_model_outputs_key: Only if what is "extra_model_outputs", this + specifies the sub-key (str) inside the extra_model_outputs dict, e.g. + STATE_OUT or ACTION_DIST_INPUTS. + hanging_val: In case we are pulling actions, rewards, or extra_model_outputs + data, there might be information "hanging" (cached). For example, + if an agent receives an observation o0 and then immediately sends an + action a0 back, but then does NOT immediately reveive a next + observation, a0 is now cached (not fully logged yet with this + episode). The currently cached value must be provided here to be able + to return it in case the index is -1 (most recent timestep). + + Returns: + A data item corresponding to the provided args. + """ + sa_episode = self.agent_episodes[agent_id] + + inf_lookback_buffer = getattr(sa_episode, what) + if extra_model_outputs_key is not None: + inf_lookback_buffer = inf_lookback_buffer[extra_model_outputs_key] + + # If there are self.SKIP_ENV_TS_TAG items in `indices_incl_lookback` and user + # wants to fill these (together with outside-episode-bounds indices) -> + # Provide these skipped timesteps as filled values. + if self.SKIP_ENV_TS_TAG in indices_incl_lookback and fill is not None: + single_fill_value = inf_lookback_buffer.get( + indices=1000000000000, + neg_index_as_lookback=False, + fill=fill, + **one_hot_discrete, + ) + ret = [] + for i in indices_incl_lookback: + if i == self.SKIP_ENV_TS_TAG: + ret.append(single_fill_value) + else: + ret.append( + inf_lookback_buffer.get( + indices=i - getattr(sa_episode, what).lookback, + neg_index_as_lookback=True, + fill=fill, + _add_last_ts_value=hanging_val, + **one_hot_discrete, + ) + ) + if self.is_numpy: + ret = batch(ret) + else: + # Filter these indices out up front. + indices = [ + i - inf_lookback_buffer.lookback + for i in indices_incl_lookback + if i != self.SKIP_ENV_TS_TAG + ] + ret = inf_lookback_buffer.get( + indices=indices, + neg_index_as_lookback=True, + fill=fill, + _add_last_ts_value=hanging_val, + **one_hot_discrete, + ) + return ret + + def _get_hanging_value(self, what: str, agent_id: AgentID) -> Any: + """Returns the hanging action/reward/extra_model_outputs for given agent.""" + if what == "actions": + return self._hanging_actions_end.get(agent_id) + elif what == "extra_model_outputs": + return self._hanging_extra_model_outputs_end.get(agent_id) + elif what == "rewards": + return self._hanging_rewards_end.get(agent_id) + + def _copy_hanging(self, agent_id: AgentID, other: "MultiAgentEpisode") -> None: + """Copies hanging action, reward, extra_model_outputs from `other` to `self.""" + if agent_id in other._hanging_rewards_begin: + self._hanging_rewards_begin[agent_id] = other._hanging_rewards_begin[ + agent_id + ] + if agent_id in other._hanging_rewards_end: + self._hanging_actions_end[agent_id] = copy.deepcopy( + other._hanging_actions_end[agent_id] + ) + self._hanging_rewards_end[agent_id] = other._hanging_rewards_end[agent_id] + self._hanging_extra_model_outputs_end[agent_id] = copy.deepcopy( + other._hanging_extra_model_outputs_end[agent_id] + ) + + def _del_hanging(self, agent_id: AgentID) -> None: + """Deletes all hanging action, reward, extra_model_outputs of given agent.""" + self._hanging_rewards_begin.pop(agent_id, None) + + self._hanging_actions_end.pop(agent_id, None) + self._hanging_extra_model_outputs_end.pop(agent_id, None) + self._hanging_rewards_end.pop(agent_id, None) + + def _del_agent(self, agent_id: AgentID) -> None: + """Deletes all data of given agent from this episode.""" + self._del_hanging(agent_id) + self.agent_episodes.pop(agent_id, None) + self.agent_ids.discard(agent_id) + self.env_t_to_agent_t.pop(agent_id, None) + self._agent_to_module_mapping.pop(agent_id, None) + self.agent_t_started.pop(agent_id, None) + + def _get_inf_lookback_buffer_or_dict( + self, + agent_id: AgentID, + what: str, + extra_model_outputs_key: Optional[str] = None, + hanging_val: Optional[Any] = None, + filter_for_skip_indices=None, + ): + """Returns a single InfiniteLookbackBuffer or a dict of such. + + In case `what` is "extra_model_outputs" AND `extra_model_outputs_key` is None, + a dict is returned. In all other cases, a single InfiniteLookbackBuffer is + returned. + """ + inf_lookback_buffer_or_dict = inf_lookback_buffer = getattr( + self.agent_episodes[agent_id], what + ) + if what == "extra_model_outputs": + if extra_model_outputs_key is not None: + inf_lookback_buffer = inf_lookback_buffer_or_dict[ + extra_model_outputs_key + ] + elif inf_lookback_buffer_or_dict: + inf_lookback_buffer = next(iter(inf_lookback_buffer_or_dict.values())) + elif filter_for_skip_indices is not None: + return inf_lookback_buffer_or_dict, filter_for_skip_indices + else: + return inf_lookback_buffer_or_dict + + if filter_for_skip_indices is not None: + inf_lookback_buffer_len = ( + len(inf_lookback_buffer) + + inf_lookback_buffer.lookback + + (hanging_val is not None) + ) + ignore_last_ts = what not in ["observations", "infos"] + if isinstance(filter_for_skip_indices, list): + filter_for_skip_indices = [ + "S" if ignore_last_ts and i == inf_lookback_buffer_len else i + for i in filter_for_skip_indices + ] + elif ignore_last_ts and filter_for_skip_indices == inf_lookback_buffer_len: + filter_for_skip_indices = "S" + return inf_lookback_buffer_or_dict, filter_for_skip_indices + else: + return inf_lookback_buffer_or_dict + + @Deprecated(new="MultiAgentEpisode.custom_data[some-key] = ...", error=True) + def add_temporary_timestep_data(self): + pass + + @Deprecated(new="MultiAgentEpisode.custom_data[some-key]", error=True) + def get_temporary_timestep_data(self): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/env/policy_client.py b/lib/python3.12/site-packages/ray/rllib/env/policy_client.py new file mode 100644 index 0000000000000000000000000000000000000000..9ae02d5714b7a8b9e8532d86bf0d8da5ba6c7fc5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/policy_client.py @@ -0,0 +1,321 @@ +import logging +import threading +import time +from typing import Optional, Union + +import ray.cloudpickle as pickle + +# Backward compatibility. +from ray.rllib.env.external.rllink import RLlink as Commands +from ray.rllib.env.external_env import ExternalEnv +from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import ( + EnvActionType, + EnvInfoDict, + EnvObsType, + MultiAgentDict, +) + +logger = logging.getLogger(__name__) + +try: + import requests # `requests` is not part of stdlib. +except ImportError: + requests = None + logger.warning( + "Couldn't import `requests` library. Be sure to install it on" + " the client side." + ) + + +@OldAPIStack +class PolicyClient: + """REST client to interact with an RLlib policy server.""" + + def __init__( + self, + address: str, + inference_mode: str = "local", + update_interval: float = 10.0, + session: Optional[requests.Session] = None, + ): + self.address = address + self.session = session + self.env: ExternalEnv = None + if inference_mode == "local": + self.local = True + self._setup_local_rollout_worker(update_interval) + elif inference_mode == "remote": + self.local = False + else: + raise ValueError("inference_mode must be either 'local' or 'remote'") + + def start_episode( + self, episode_id: Optional[str] = None, training_enabled: bool = True + ) -> str: + if self.local: + self._update_local_policy() + return self.env.start_episode(episode_id, training_enabled) + + return self._send( + { + "episode_id": episode_id, + "command": Commands.START_EPISODE, + "training_enabled": training_enabled, + } + )["episode_id"] + + def get_action( + self, episode_id: str, observation: Union[EnvObsType, MultiAgentDict] + ) -> Union[EnvActionType, MultiAgentDict]: + if self.local: + self._update_local_policy() + if isinstance(episode_id, (list, tuple)): + actions = { + eid: self.env.get_action(eid, observation[eid]) + for eid in episode_id + } + return actions + else: + return self.env.get_action(episode_id, observation) + else: + return self._send( + { + "command": Commands.GET_ACTION, + "observation": observation, + "episode_id": episode_id, + } + )["action"] + + def log_action( + self, + episode_id: str, + observation: Union[EnvObsType, MultiAgentDict], + action: Union[EnvActionType, MultiAgentDict], + ) -> None: + if self.local: + self._update_local_policy() + return self.env.log_action(episode_id, observation, action) + + self._send( + { + "command": Commands.LOG_ACTION, + "observation": observation, + "action": action, + "episode_id": episode_id, + } + ) + + def log_returns( + self, + episode_id: str, + reward: float, + info: Union[EnvInfoDict, MultiAgentDict] = None, + multiagent_done_dict: Optional[MultiAgentDict] = None, + ) -> None: + if self.local: + self._update_local_policy() + if multiagent_done_dict is not None: + assert isinstance(reward, dict) + return self.env.log_returns( + episode_id, reward, info, multiagent_done_dict + ) + return self.env.log_returns(episode_id, reward, info) + + self._send( + { + "command": Commands.LOG_RETURNS, + "reward": reward, + "info": info, + "episode_id": episode_id, + "done": multiagent_done_dict, + } + ) + + def end_episode( + self, episode_id: str, observation: Union[EnvObsType, MultiAgentDict] + ) -> None: + if self.local: + self._update_local_policy() + return self.env.end_episode(episode_id, observation) + + self._send( + { + "command": Commands.END_EPISODE, + "observation": observation, + "episode_id": episode_id, + } + ) + + def update_policy_weights(self) -> None: + """Query the server for new policy weights, if local inference is enabled.""" + self._update_local_policy(force=True) + + def _send(self, data): + payload = pickle.dumps(data) + + if self.session is None: + response = requests.post(self.address, data=payload) + else: + response = self.session.post(self.address, data=payload) + + if response.status_code != 200: + logger.error("Request failed {}: {}".format(response.text, data)) + response.raise_for_status() + parsed = pickle.loads(response.content) + return parsed + + def _setup_local_rollout_worker(self, update_interval): + self.update_interval = update_interval + self.last_updated = 0 + + logger.info("Querying server for rollout worker settings.") + kwargs = self._send( + { + "command": Commands.GET_WORKER_ARGS, + } + )["worker_args"] + (self.rollout_worker, self.inference_thread) = _create_embedded_rollout_worker( + kwargs, self._send + ) + self.env = self.rollout_worker.env + + def _update_local_policy(self, force=False): + assert self.inference_thread.is_alive() + if ( + self.update_interval + and time.time() - self.last_updated > self.update_interval + ) or force: + logger.info("Querying server for new policy weights.") + resp = self._send( + { + "command": Commands.GET_WEIGHTS, + } + ) + weights = resp["weights"] + global_vars = resp["global_vars"] + logger.info( + "Updating rollout worker weights and global vars {}.".format( + global_vars + ) + ) + self.rollout_worker.set_weights(weights, global_vars) + self.last_updated = time.time() + + +@OldAPIStack +class _LocalInferenceThread(threading.Thread): + def __init__(self, rollout_worker, send_fn): + super().__init__() + self.daemon = True + self.rollout_worker = rollout_worker + self.send_fn = send_fn + + def run(self): + try: + while True: + logger.info("Generating new batch of experiences.") + samples = self.rollout_worker.sample() + metrics = self.rollout_worker.get_metrics() + if isinstance(samples, MultiAgentBatch): + logger.info( + "Sending batch of {} env steps ({} agent steps) to " + "server.".format(samples.env_steps(), samples.agent_steps()) + ) + else: + logger.info( + "Sending batch of {} steps back to server.".format( + samples.count + ) + ) + self.send_fn( + { + "command": Commands.REPORT_SAMPLES, + "samples": samples, + "metrics": metrics, + } + ) + except Exception as e: + logger.error("Error: inference worker thread died!", e) + + +@OldAPIStack +def _auto_wrap_external(real_env_creator): + def wrapped_creator(env_config): + real_env = real_env_creator(env_config) + if not isinstance(real_env, (ExternalEnv, ExternalMultiAgentEnv)): + logger.info( + "The env you specified is not a supported (sub-)type of " + "ExternalEnv. Attempting to convert it automatically to " + "ExternalEnv." + ) + + if isinstance(real_env, MultiAgentEnv): + external_cls = ExternalMultiAgentEnv + else: + external_cls = ExternalEnv + + class _ExternalEnvWrapper(external_cls): + def __init__(self, real_env): + super().__init__( + observation_space=real_env.observation_space, + action_space=real_env.action_space, + ) + + def run(self): + # Since we are calling methods on this class in the + # client, run doesn't need to do anything. + time.sleep(999999) + + return _ExternalEnvWrapper(real_env) + return real_env + + return wrapped_creator + + +@OldAPIStack +def _create_embedded_rollout_worker(kwargs, send_fn): + # Since the server acts as an input datasource, we have to reset the + # input config to the default, which runs env rollouts. + kwargs = kwargs.copy() + kwargs["config"] = kwargs["config"].copy(copy_frozen=False) + config = kwargs["config"] + config.output = None + config.input_ = "sampler" + config.input_config = {} + + # If server has no env (which is the expected case): + # Generate a dummy ExternalEnv here using RandomEnv and the + # given observation/action spaces. + if config.env is None: + from ray.rllib.examples.envs.classes.random_env import ( + RandomEnv, + RandomMultiAgentEnv, + ) + + env_config = { + "action_space": config.action_space, + "observation_space": config.observation_space, + } + is_ma = config.is_multi_agent + kwargs["env_creator"] = _auto_wrap_external( + lambda _: (RandomMultiAgentEnv if is_ma else RandomEnv)(env_config) + ) + # kwargs["config"].env = True + # Otherwise, use the env specified by the server args. + else: + real_env_creator = kwargs["env_creator"] + kwargs["env_creator"] = _auto_wrap_external(real_env_creator) + + logger.info("Creating rollout worker with kwargs={}".format(kwargs)) + from ray.rllib.evaluation.rollout_worker import RolloutWorker + + rollout_worker = RolloutWorker(**kwargs) + + inference_thread = _LocalInferenceThread(rollout_worker, send_fn) + inference_thread.start() + + return rollout_worker, inference_thread diff --git a/lib/python3.12/site-packages/ray/rllib/env/remote_base_env.py b/lib/python3.12/site-packages/ray/rllib/env/remote_base_env.py new file mode 100644 index 0000000000000000000000000000000000000000..5d48de3098ec06249ec18ffe4590986aeaf2ab38 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/remote_base_env.py @@ -0,0 +1,463 @@ +import logging +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple + +import gymnasium as gym + +import ray +from ray.rllib.env.base_env import _DUMMY_AGENT_ID, ASYNC_RESET_RETURN, BaseEnv +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.typing import AgentID, EnvID, EnvType, MultiEnvDict +from ray.util import log_once + +if TYPE_CHECKING: + from ray.rllib.evaluation.rollout_worker import RolloutWorker + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class RemoteBaseEnv(BaseEnv): + """BaseEnv that executes its sub environments as @ray.remote actors. + + This provides dynamic batching of inference as observations are returned + from the remote simulator actors. Both single and multi-agent child envs + are supported, and envs can be stepped synchronously or asynchronously. + + NOTE: This class implicitly assumes that the remote envs are gym.Env's + + You shouldn't need to instantiate this class directly. It's automatically + inserted when you use the `remote_worker_envs=True` option in your + Algorithm's config. + """ + + def __init__( + self, + make_env: Callable[[int], EnvType], + num_envs: int, + multiagent: bool, + remote_env_batch_wait_ms: int, + existing_envs: Optional[List[ray.actor.ActorHandle]] = None, + worker: Optional["RolloutWorker"] = None, + restart_failed_sub_environments: bool = False, + ): + """Initializes a RemoteVectorEnv instance. + + Args: + make_env: Callable that produces a single (non-vectorized) env, + given the vector env index as only arg. + num_envs: The number of sub-environments to create for the + vectorization. + multiagent: Whether this is a multiagent env or not. + remote_env_batch_wait_ms: Time to wait for (ray.remote) + sub-environments to have new observations available when + polled. Only when none of the sub-environments is ready, + repeat the `ray.wait()` call until at least one sub-env + is ready. Then return only the observations of the ready + sub-environment(s). + existing_envs: Optional list of already created sub-environments. + These will be used as-is and only as many new sub-envs as + necessary (`num_envs - len(existing_envs)`) will be created. + worker: An optional RolloutWorker that owns the env. This is only + used if `remote_worker_envs` is True in your config and the + `on_sub_environment_created` custom callback needs to be + called on each created actor. + restart_failed_sub_environments: If True and any sub-environment (within + a vectorized env) throws any error during env stepping, the + Sampler will try to restart the faulty sub-environment. This is done + without disturbing the other (still intact) sub-environment and without + the RolloutWorker crashing. + """ + + # Could be creating local or remote envs. + self.make_env = make_env + self.num_envs = num_envs + self.multiagent = multiagent + self.poll_timeout = remote_env_batch_wait_ms / 1000 + self.worker = worker + self.restart_failed_sub_environments = restart_failed_sub_environments + + # Already existing env objects (generated by the RolloutWorker). + existing_envs = existing_envs or [] + + # Whether the given `make_env` callable already returns ActorHandles + # (@ray.remote class instances) or not. + self.make_env_creates_actors = False + + self._observation_space = None + self._action_space = None + + # List of ray actor handles (each handle points to one @ray.remote + # sub-environment). + self.actors: Optional[List[ray.actor.ActorHandle]] = None + + # `self.make_env` already produces Actors: Use it directly. + if len(existing_envs) > 0 and isinstance( + existing_envs[0], ray.actor.ActorHandle + ): + self.make_env_creates_actors = True + self.actors = existing_envs + while len(self.actors) < self.num_envs: + self.actors.append(self._make_sub_env(len(self.actors))) + + # `self.make_env` produces gym.Envs (or children thereof, such + # as MultiAgentEnv): Need to auto-wrap it here. The problem with + # this is that custom methods wil get lost. If you would like to + # keep your custom methods in your envs, you should provide the + # env class directly in your config (w/o tune.register_env()), + # such that your class can directly be made a @ray.remote + # (w/o the wrapping via `_Remote[Multi|Single]AgentEnv`). + # Also, if `len(existing_envs) > 0`, we have to throw those away + # as we need to create ray actors here. + else: + self.actors = [self._make_sub_env(i) for i in range(self.num_envs)] + # Utilize existing envs for inferring observation/action spaces. + if len(existing_envs) > 0: + self._observation_space = existing_envs[0].observation_space + self._action_space = existing_envs[0].action_space + # Have to call actors' remote methods to get observation/action spaces. + else: + self._observation_space, self._action_space = ray.get( + [ + self.actors[0].observation_space.remote(), + self.actors[0].action_space.remote(), + ] + ) + + # Dict mapping object refs (return values of @ray.remote calls), + # whose actual values we are waiting for (via ray.wait in + # `self.poll()`) to their corresponding actor handles (the actors + # that created these return values). + # Call `reset()` on all @ray.remote sub-environment actors. + self.pending: Dict[ray.actor.ActorHandle] = { + a.reset.remote(): a for a in self.actors + } + + @override(BaseEnv) + def poll( + self, + ) -> Tuple[ + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + ]: + + # each keyed by env_id in [0, num_remote_envs) + obs, rewards, terminateds, truncateds, infos = {}, {}, {}, {}, {} + ready = [] + + # Wait for at least 1 env to be ready here. + while not ready: + ready, _ = ray.wait( + list(self.pending), + num_returns=len(self.pending), + timeout=self.poll_timeout, + ) + + # Get and return observations for each of the ready envs + env_ids = set() + for obj_ref in ready: + # Get the corresponding actor handle from our dict and remove the + # object ref (we will call `ray.get()` on it and it will no longer + # be "pending"). + actor = self.pending.pop(obj_ref) + env_id = self.actors.index(actor) + env_ids.add(env_id) + # Get the ready object ref (this may be return value(s) of + # `reset()` or `step()`). + try: + ret = ray.get(obj_ref) + except Exception as e: + # Something happened on the actor during stepping/resetting. + # Restart sub-environment (create new actor; close old one). + if self.restart_failed_sub_environments: + logger.exception(e.args[0]) + self.try_restart(env_id) + # Always return multi-agent data. + # Set the observation to the exception, no rewards, + # terminated[__all__]=True (episode will be discarded anyways), + # no infos. + ret = ( + e, + {}, + {"__all__": True}, + {"__all__": False}, + {}, + ) + # Do not try to restart. Just raise the error. + else: + raise e + + # Our sub-envs are simple Actor-turned gym.Envs or MultiAgentEnvs. + if self.make_env_creates_actors: + rew, terminated, truncated, info = None, None, None, None + if self.multiagent: + if isinstance(ret, tuple): + # Gym >= 0.26: `step()` result: Obs, reward, terminated, + # truncated, info. + if len(ret) == 5: + ob, rew, terminated, truncated, info = ret + # Gym >= 0.26: `reset()` result: Obs and infos. + elif len(ret) == 2: + ob = ret[0] + info = ret[1] + # Gym < 0.26? Something went wrong. + else: + raise AssertionError( + "Your gymnasium.Env seems to NOT return the correct " + "number of return values for `step()` (needs to return" + " 5 values: obs, reward, terminated, truncated and " + "info) or `reset()` (needs to return 2 values: obs and " + "info)!" + ) + # Gym < 0.26: `reset()` result: Only obs. + else: + raise AssertionError( + "Your gymnasium.Env seems to only return a single value " + "upon `reset()`! Must return 2 (obs AND infos)." + ) + else: + if isinstance(ret, tuple): + # `step()` result: Obs, reward, terminated, truncated, info. + if len(ret) == 5: + ob = {_DUMMY_AGENT_ID: ret[0]} + rew = {_DUMMY_AGENT_ID: ret[1]} + terminated = {_DUMMY_AGENT_ID: ret[2], "__all__": ret[2]} + truncated = {_DUMMY_AGENT_ID: ret[3], "__all__": ret[3]} + info = {_DUMMY_AGENT_ID: ret[4]} + # `reset()` result: Obs and infos. + elif len(ret) == 2: + ob = {_DUMMY_AGENT_ID: ret[0]} + info = {_DUMMY_AGENT_ID: ret[1]} + # Gym < 0.26? Something went wrong. + else: + raise AssertionError( + "Your gymnasium.Env seems to NOT return the correct " + "number of return values for `step()` (needs to return" + " 5 values: obs, reward, terminated, truncated and " + "info) or `reset()` (needs to return 2 values: obs and " + "info)!" + ) + # Gym < 0.26? + else: + raise AssertionError( + "Your gymnasium.Env seems to only return a single value " + "upon `reset()`! Must return 2 (obs and infos)." + ) + + # If this is a `reset()` return value, we only have the initial + # observations and infos: Set rewards, terminateds, and truncateds to + # dummy values. + if rew is None: + rew = {agent_id: 0 for agent_id in ob.keys()} + terminated = {"__all__": False} + truncated = {"__all__": False} + + # Our sub-envs are auto-wrapped (by `_RemoteSingleAgentEnv` or + # `_RemoteMultiAgentEnv`) and already behave like multi-agent + # envs. + else: + ob, rew, terminated, truncated, info = ret + obs[env_id] = ob + rewards[env_id] = rew + terminateds[env_id] = terminated + truncateds[env_id] = truncated + infos[env_id] = info + + logger.debug(f"Got obs batch for actors {env_ids}") + return obs, rewards, terminateds, truncateds, infos, {} + + @override(BaseEnv) + def send_actions(self, action_dict: MultiEnvDict) -> None: + for env_id, actions in action_dict.items(): + actor = self.actors[env_id] + # `actor` is a simple single-agent (remote) env, e.g. a gym.Env + # that was made a @ray.remote. + if not self.multiagent and self.make_env_creates_actors: + obj_ref = actor.step.remote(actions[_DUMMY_AGENT_ID]) + # `actor` is already a _RemoteSingleAgentEnv or + # _RemoteMultiAgentEnv wrapper + # (handles the multi-agent action_dict automatically). + else: + obj_ref = actor.step.remote(actions) + self.pending[obj_ref] = actor + + @override(BaseEnv) + def try_reset( + self, + env_id: Optional[EnvID] = None, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Tuple[MultiEnvDict, MultiEnvDict]: + actor = self.actors[env_id] + obj_ref = actor.reset.remote(seed=seed, options=options) + + self.pending[obj_ref] = actor + # Because this env type does not support synchronous reset requests (with + # immediate return value), we return ASYNC_RESET_RETURN here to indicate + # that the reset results will be available via the next `poll()` call. + return ASYNC_RESET_RETURN, ASYNC_RESET_RETURN + + @override(BaseEnv) + def try_restart(self, env_id: Optional[EnvID] = None) -> None: + # Try closing down the old (possibly faulty) sub-env, but ignore errors. + try: + # Close the env on the remote side. + self.actors[env_id].close.remote() + except Exception as e: + if log_once("close_sub_env"): + logger.warning( + "Trying to close old and replaced sub-environment (at vector " + f"index={env_id}), but closing resulted in error:\n{e}" + ) + + # Terminate the actor itself to free up its resources. + self.actors[env_id].__ray_terminate__.remote() + + # Re-create a new sub-environment. + self.actors[env_id] = self._make_sub_env(env_id) + + @override(BaseEnv) + def stop(self) -> None: + if self.actors is not None: + for actor in self.actors: + actor.__ray_terminate__.remote() + + @override(BaseEnv) + def get_sub_environments(self, as_dict: bool = False) -> List[EnvType]: + if as_dict: + return dict(enumerate(self.actors)) + return self.actors + + @property + @override(BaseEnv) + def observation_space(self) -> gym.spaces.Dict: + return self._observation_space + + @property + @override(BaseEnv) + def action_space(self) -> gym.Space: + return self._action_space + + def _make_sub_env(self, idx: Optional[int] = None): + """Re-creates a sub-environment at the new index.""" + + # Our `make_env` creates ray actors directly. + if self.make_env_creates_actors: + sub_env = self.make_env(idx) + if self.worker is not None: + self.worker.callbacks.on_sub_environment_created( + worker=self.worker, + sub_environment=self.actors[idx], + env_context=self.worker.env_context.copy_with_overrides( + vector_index=idx + ), + ) + + # Our `make_env` returns actual envs -> Have to convert them into actors + # using our utility wrapper classes. + else: + + def make_remote_env(i): + logger.info("Launching env {} in remote actor".format(i)) + if self.multiagent: + sub_env = _RemoteMultiAgentEnv.remote(self.make_env, i) + else: + sub_env = _RemoteSingleAgentEnv.remote(self.make_env, i) + + if self.worker is not None: + self.worker.callbacks.on_sub_environment_created( + worker=self.worker, + sub_environment=sub_env, + env_context=self.worker.env_context.copy_with_overrides( + vector_index=i + ), + ) + + return sub_env + + sub_env = make_remote_env(idx) + + return sub_env + + @override(BaseEnv) + def get_agent_ids(self) -> Set[AgentID]: + if self.multiagent: + return ray.get(self.actors[0].get_agent_ids.remote()) + else: + return {_DUMMY_AGENT_ID} + + +@ray.remote(num_cpus=0) +class _RemoteMultiAgentEnv: + """Wrapper class for making a multi-agent env a remote actor.""" + + def __init__(self, make_env, i): + self.env = make_env(i) + self.agent_ids = set() + + def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): + obs, info = self.env.reset(seed=seed, options=options) + + # each keyed by agent_id in the env + rew = {} + for agent_id in obs.keys(): + self.agent_ids.add(agent_id) + rew[agent_id] = 0.0 + terminated = {"__all__": False} + truncated = {"__all__": False} + return obs, rew, terminated, truncated, info + + def step(self, action_dict): + return self.env.step(action_dict) + + # Defining these 2 functions that way this information can be queried + # with a call to ray.get(). + def observation_space(self): + return self.env.observation_space + + def action_space(self): + return self.env.action_space + + def get_agent_ids(self) -> Set[AgentID]: + return self.agent_ids + + +@ray.remote(num_cpus=0) +class _RemoteSingleAgentEnv: + """Wrapper class for making a gym env a remote actor.""" + + def __init__(self, make_env, i): + self.env = make_env(i) + + def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): + obs_and_info = self.env.reset(seed=seed, options=options) + + obs = {_DUMMY_AGENT_ID: obs_and_info[0]} + info = {_DUMMY_AGENT_ID: obs_and_info[1]} + + rew = {_DUMMY_AGENT_ID: 0.0} + terminated = {"__all__": False} + truncated = {"__all__": False} + return obs, rew, terminated, truncated, info + + def step(self, action): + results = self.env.step(action[_DUMMY_AGENT_ID]) + + obs, rew, terminated, truncated, info = [{_DUMMY_AGENT_ID: x} for x in results] + + terminated["__all__"] = terminated[_DUMMY_AGENT_ID] + truncated["__all__"] = truncated[_DUMMY_AGENT_ID] + + return obs, rew, terminated, truncated, info + + # Defining these 2 functions that way this information can be queried + # with a call to ray.get(). + def observation_space(self): + return self.env.observation_space + + def action_space(self): + return self.env.action_space diff --git a/lib/python3.12/site-packages/ray/rllib/env/single_agent_env_runner.py b/lib/python3.12/site-packages/ray/rllib/env/single_agent_env_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..4e74a11aab6bff479b16a251e787686ddf5c02d7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/single_agent_env_runner.py @@ -0,0 +1,927 @@ +import logging +import math +import time +from collections import defaultdict +from typing import Collection, DefaultDict, List, Optional, Union + +import gymnasium as gym +from gymnasium.wrappers.vector import DictInfoToList + +import ray +from ray._common.deprecation import Deprecated +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.callbacks.utils import make_callback +from ray.rllib.core import ( + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + COMPONENT_RL_MODULE, + DEFAULT_AGENT_ID, + DEFAULT_MODULE_ID, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec +from ray.rllib.env import INPUT_ENV_SINGLE_SPACES, INPUT_ENV_SPACES +from ray.rllib.env.env_context import EnvContext +from ray.rllib.env.env_runner import ENV_STEP_FAILURE, EnvRunner +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import override +from ray.rllib.utils.checkpoints import Checkpointable +from ray.rllib.utils.error import ERR_MSG_INVALID_ENV_DESCRIPTOR, EnvError +from ray.rllib.utils.framework import get_device +from ray.rllib.utils.metrics import ( + ENV_TO_MODULE_CONNECTOR, + EPISODE_DURATION_SEC_MEAN, + EPISODE_LEN_MAX, + EPISODE_LEN_MEAN, + EPISODE_LEN_MIN, + EPISODE_RETURN_MAX, + EPISODE_RETURN_MEAN, + EPISODE_RETURN_MIN, + MODULE_TO_ENV_CONNECTOR, + NUM_AGENT_STEPS_SAMPLED, + NUM_AGENT_STEPS_SAMPLED_LIFETIME, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_EPISODES, + NUM_EPISODES_LIFETIME, + NUM_MODULE_STEPS_SAMPLED, + NUM_MODULE_STEPS_SAMPLED_LIFETIME, + RLMODULE_INFERENCE_TIMER, + SAMPLE_TIMER, + TIME_BETWEEN_SAMPLING, + WEIGHTS_SEQ_NO, +) +from ray.rllib.utils.spaces.space_utils import unbatch +from ray.rllib.utils.typing import EpisodeID, ResultDict, StateDict +from ray.tune.registry import ENV_CREATOR, _global_registry +from ray.util.annotations import PublicAPI + +logger = logging.getLogger("ray.rllib") + + +# TODO (sven): As soon as RolloutWorker is no longer supported, make `EnvRunner` itself +# a Checkpointable. Currently, only some of its subclasses are Checkpointables. +@PublicAPI(stability="alpha") +class SingleAgentEnvRunner(EnvRunner, Checkpointable): + """The generic environment runner for the single agent case.""" + + @override(EnvRunner) + def __init__(self, *, config: AlgorithmConfig, **kwargs): + """Initializes a SingleAgentEnvRunner instance. + + Args: + config: An `AlgorithmConfig` object containing all settings needed to + build this `EnvRunner` class. + """ + super().__init__(config=config, **kwargs) + + self.tune_trial_id: str = kwargs.get("tune_trial_id") + self.spaces = kwargs.get("spaces", {}) + + # Create our callbacks object. + self._callbacks: List[RLlibCallback] = [ + cls() for cls in force_list(self.config.callbacks_class) + ] + + # Set device. + self._device = get_device( + self.config, + 0 if not self.worker_index else self.config.num_gpus_per_env_runner, + ) + + # Create the vectorized gymnasium env. + self.env: Optional[gym.vector.VectorEnv] = None + self.num_envs: int = 0 + if ( + self.worker_index is None + or self.worker_index > 0 + or self.config.create_env_on_local_worker + or self.config.num_env_runners == 0 + ): + self.make_env() + + # Create the env-to-module connector pipeline. + self._env_to_module = self.config.build_env_to_module_connector( + env=self.env, spaces=self.spaces, device=self._device + ) + # Cached env-to-module results taken at the end of a `_sample_timesteps()` + # call to make sure the final observation (before an episode cut) gets properly + # processed (and maybe postprocessed and re-stored into the episode). + # For example, if we had a connector that normalizes observations and directly + # re-inserts these new obs back into the episode, the last observation in each + # sample call would NOT be processed, which could be very harmful in cases, + # in which value function bootstrapping of those (truncation) observations is + # required in the learning step. + self._cached_to_module = None + + # Create the RLModule. + self.module: Optional[RLModule] = None + self.make_module() + + # Create the module-to-env connector pipeline. + self._module_to_env = self.config.build_module_to_env_connector( + env=self.env, spaces=self.spaces + ) + + # This should be the default. + self._needs_initial_reset: bool = True + self._episodes: List[Optional[SingleAgentEpisode]] = [ + None for _ in range(self.num_envs) + ] + self._shared_data = None + + self._done_episodes_for_metrics: List[SingleAgentEpisode] = [] + self._ongoing_episodes_for_metrics: DefaultDict[ + EpisodeID, List[SingleAgentEpisode] + ] = defaultdict(list) + self._weights_seq_no: int = 0 + + # Measures the time passed between returning from `sample()` + # and receiving the next `sample()` request from the user. + self._time_after_sampling = None + + # Save whether to convert episodes to numpy during sample + # In `OfflineSingleAgentEnvRunner`, this result is set to False + # during initialisation + self.episodes_to_numpy = self.config.episodes_to_numpy + + @override(EnvRunner) + def sample( + self, + *, + num_timesteps: int = None, + num_episodes: int = None, + explore: bool = None, + random_actions: bool = False, + force_reset: bool = False, + ) -> List[SingleAgentEpisode]: + """Runs and returns a sample (n timesteps or m episodes) on the env(s). + + If neither `num_timesteps` nor `num_episodes` are provided and the config + `batch_mode` is "truncate_episodes" then + `config.get_rollout_fragment_length(self.worker_index) * self.num_envs` + timesteps will be sampled. + + Args: + num_timesteps: The number of timesteps to sample during this call. + The episodes returned will contain the total timesteps greater than or + equal to num_timesteps and less than num_timesteps + num_envs_per_env_runner. + Note that only one of `num_timesteps` or `num_episodes` may be provided. + num_episodes: The number of episodes to sample during this call. + Note that only one of `num_timesteps` or `num_episodes` may be provided. + explore: If True, will use the RLModule's `forward_exploration()` + method to compute actions. If False, will use the RLModule's + `forward_inference()` method. If None (default), will use the `explore` + boolean setting from `self.config` passed into this EnvRunner's + constructor. You can change this setting in your config via + `config.env_runners(explore=True|False)`. + random_actions: If True, actions will be sampled randomly (from the action + space of the environment). If False (default), actions or action + distribution parameters are computed by the RLModule. + force_reset: Whether to force-reset all vectorized environments before + sampling. Useful if you would like to collect a clean slate of new + episodes via this call. Note that when sampling n episodes + (`num_episodes != None`), this is fixed to True. + + Returns: + A list of `SingleAgentEpisode` instances, carrying the sampled data. + """ + if self.env is None: + raise ValueError( + f"{self} doesn't have an env! Can't call `sample()` on it." + ) + + assert not (num_timesteps is not None and num_episodes is not None) + + # Log time between `sample()` requests. + if self._time_after_sampling is not None: + self.metrics.log_value( + key=TIME_BETWEEN_SAMPLING, + value=time.perf_counter() - self._time_after_sampling, + ) + + # Log current weight seq no. + self.metrics.log_value( + key=WEIGHTS_SEQ_NO, + value=self._weights_seq_no, + window=1, + ) + + with self.metrics.log_time(SAMPLE_TIMER): + # If no execution details are provided, use the config to try to infer the + # desired timesteps/episodes to sample and exploration behavior. + if explore is None: + explore = self.config.explore + + if ( + num_timesteps is None + and num_episodes is None + and self.config.batch_mode == "truncate_episodes" + ): + num_timesteps = ( + self.config.get_rollout_fragment_length(self.worker_index) + * self.num_envs + ) + + # Sample n timesteps. + if num_timesteps is not None: + assert num_timesteps >= 0 + samples = self._sample( + num_timesteps=num_timesteps, + explore=explore, + random_actions=random_actions, + force_reset=force_reset, + ) + # Sample m episodes. + elif num_episodes is not None: + assert num_episodes >= 0 + samples = self._sample( + num_episodes=num_episodes, + explore=explore, + random_actions=random_actions, + ) + # For complete episodes mode, sample as long as the number of timesteps + # done is smaller than the `train_batch_size`. + else: + samples = self._sample( + num_episodes=self.num_envs, + explore=explore, + random_actions=random_actions, + ) + + # Make the `on_sample_end` callback. + make_callback( + "on_sample_end", + callbacks_objects=self._callbacks, + callbacks_functions=self.config.callbacks_on_sample_end, + kwargs=dict( + env_runner=self, + metrics_logger=self.metrics, + samples=samples, + ), + ) + + self._time_after_sampling = time.perf_counter() + + return samples + + def _sample( + self, + *, + num_timesteps: Optional[int] = None, + num_episodes: Optional[int] = None, + explore: bool, + random_actions: bool = False, + force_reset: bool = False, + ) -> List[SingleAgentEpisode]: + """Helper method to sample n timesteps or m episodes.""" + + done_episodes_to_return: List[SingleAgentEpisode] = [] + + # Have to reset the env (on all vector sub_envs). + if force_reset or num_episodes is not None or self._needs_initial_reset: + episodes = self._episodes = [None for _ in range(self.num_envs)] + shared_data = self._shared_data = {} + self._reset_envs(episodes, shared_data, explore) + # We just reset the env. Don't have to force this again in the next + # call to `self._sample_timesteps()`. + self._needs_initial_reset = False + else: + episodes = self._episodes + shared_data = self._shared_data + + if num_episodes is not None: + self._needs_initial_reset = True + + # Loop through `num_timesteps` timesteps or `num_episodes` episodes. + ts = 0 + eps = 0 + while ( + (ts < num_timesteps) if num_timesteps is not None else (eps < num_episodes) + ): + # Act randomly. + if random_actions: + to_env = { + Columns.ACTIONS: self.env.action_space.sample(), + } + # Compute an action using the RLModule. + else: + # Env-to-module connector (already cached). + to_module = self._cached_to_module + assert to_module is not None + self._cached_to_module = None + + # RLModule forward pass: Explore or not. + if explore: + # Global env steps sampled are (roughly) this EnvRunner's lifetime + # count times the number of env runners in the algo. + global_env_steps_lifetime = ( + self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME, default=0) + + ts + ) * (self.config.num_env_runners or 1) + with self.metrics.log_time(RLMODULE_INFERENCE_TIMER): + to_env = self.module.forward_exploration( + to_module, t=global_env_steps_lifetime + ) + else: + with self.metrics.log_time(RLMODULE_INFERENCE_TIMER): + to_env = self.module.forward_inference(to_module) + + # Module-to-env connector. + to_env = self._module_to_env( + rl_module=self.module, + batch=to_env, + episodes=episodes, + explore=explore, + shared_data=shared_data, + metrics=self.metrics, + metrics_prefix_key=(MODULE_TO_ENV_CONNECTOR,), + ) + + # Extract the (vectorized) actions (to be sent to the env) from the + # module/connector output. Note that these actions are fully ready (e.g. + # already unsquashed/clipped) to be sent to the environment and might not + # be identical to the actions produced by the RLModule/distribution, which + # are the ones stored permanently in the episode objects. + actions = to_env.pop(Columns.ACTIONS) + actions_for_env = to_env.pop(Columns.ACTIONS_FOR_ENV, actions) + # Try stepping the environment. + results = self._try_env_step(actions_for_env) + if results == ENV_STEP_FAILURE: + return self._sample( + num_timesteps=num_timesteps, + num_episodes=num_episodes, + explore=explore, + random_actions=random_actions, + force_reset=True, + ) + observations, rewards, terminateds, truncateds, infos = results + observations, actions = unbatch(observations), unbatch(actions) + + call_on_episode_start = set() + for env_index in range(self.num_envs): + extra_model_output = {k: v[env_index] for k, v in to_env.items()} + extra_model_output[WEIGHTS_SEQ_NO] = self._weights_seq_no + + # Episode has no data in it yet -> Was just reset and needs to be called + # with its `add_env_reset()` method. + if not self._episodes[env_index].is_reset: + episodes[env_index].add_env_reset( + observation=observations[env_index], + infos=infos[env_index], + ) + call_on_episode_start.add(env_index) + + # Call `add_env_step()` method on episode. + else: + # Only increase ts when we actually stepped (not reset as a reset + # does not count as a timestep). + ts += 1 + episodes[env_index].add_env_step( + observation=observations[env_index], + action=actions[env_index], + reward=rewards[env_index], + infos=infos[env_index], + terminated=terminateds[env_index], + truncated=truncateds[env_index], + extra_model_outputs=extra_model_output, + ) + + # Env-to-module connector pass cache results as we will do the RLModule + # forward pass only in the next `while`-iteration. + if self.module is not None: + self._cached_to_module = self._env_to_module( + batch={}, + episodes=episodes, + explore=explore, + rl_module=self.module, + shared_data=shared_data, + metrics=self.metrics, + metrics_prefix_key=(ENV_TO_MODULE_CONNECTOR,), + ) + + for env_index in range(self.num_envs): + # Call `on_episode_start()` callback (always after reset). + if env_index in call_on_episode_start: + self._make_on_episode_callback( + "on_episode_start", env_index, episodes + ) + # Make the `on_episode_step` callbacks. + else: + self._make_on_episode_callback( + "on_episode_step", env_index, episodes + ) + + # Episode is done. + if episodes[env_index].is_done: + eps += 1 + + # Make the `on_episode_end` callbacks (before finalizing the episode + # object). + self._make_on_episode_callback( + "on_episode_end", env_index, episodes + ) + + # Numpy'ize the episode. + if self.episodes_to_numpy: + # Any possibly compress observations. + done_episodes_to_return.append(episodes[env_index].to_numpy()) + # Leave episode as lists of individual (obs, action, etc..) items. + else: + done_episodes_to_return.append(episodes[env_index]) + + # Also early-out if we reach the number of episodes within this + # for-loop. + if eps == num_episodes: + break + + # Create a new episode object with no data in it and execute + # `on_episode_created` callback (before the `env.reset()` call). + self._new_episode(env_index, episodes) + + # Return done episodes ... + self._done_episodes_for_metrics.extend(done_episodes_to_return) + # ... and all ongoing episode chunks. + + # Also, make sure we start new episode chunks (continuing the ongoing episodes + # from the to-be-returned chunks). + ongoing_episodes_to_return = [] + # Only if we are doing individual timesteps: We have to maybe cut an ongoing + # episode and continue building it on the next call to `sample()`. + if num_timesteps is not None: + ongoing_episodes_continuations = [ + eps.cut(len_lookback_buffer=self.config.episode_lookback_horizon) + for eps in self._episodes + ] + + for eps in self._episodes: + # Just started episodes do not have to be returned. There is no data + # in them anyway. + if eps.t == 0: + continue + eps.validate() + self._ongoing_episodes_for_metrics[eps.id_].append(eps) + + # Numpy'ize the episode. + if self.episodes_to_numpy: + # Any possibly compress observations. + ongoing_episodes_to_return.append(eps.to_numpy()) + # Leave episode as lists of individual (obs, action, etc..) items. + else: + ongoing_episodes_to_return.append(eps) + + # Continue collecting into the cut Episode chunks. + self._episodes = ongoing_episodes_continuations + + # Ray metrics + self._log_env_steps(metric=self._metrics_num_env_steps_sampled, num_steps=ts) + + self._increase_sampled_metrics(ts, len(done_episodes_to_return)) + + # Return collected episode data. + return done_episodes_to_return + ongoing_episodes_to_return + + @override(EnvRunner) + def get_spaces(self): + if self.env is None: + return self.spaces + return { + INPUT_ENV_SPACES: (self.env.observation_space, self.env.action_space), + INPUT_ENV_SINGLE_SPACES: ( + self.env.single_observation_space, + self.env.single_action_space, + ), + DEFAULT_MODULE_ID: ( + self._env_to_module.observation_space, + self.env.single_action_space, + ), + } + + @override(EnvRunner) + def get_metrics(self) -> ResultDict: + # Compute per-episode metrics (only on already completed episodes). + for eps in self._done_episodes_for_metrics: + assert eps.is_done + episode_length = len(eps) + episode_return = eps.get_return() + episode_duration_s = eps.get_duration_s() + # Don't forget about the already returned chunks of this episode. + if eps.id_ in self._ongoing_episodes_for_metrics: + for eps2 in self._ongoing_episodes_for_metrics[eps.id_]: + episode_length += len(eps2) + episode_return += eps2.get_return() + episode_duration_s += eps2.get_duration_s() + del self._ongoing_episodes_for_metrics[eps.id_] + + self._log_episode_metrics( + episode_length, episode_return, episode_duration_s + ) + + # Now that we have logged everything, clear cache of done episodes. + self._done_episodes_for_metrics.clear() + + # Return reduced metrics. + return self.metrics.reduce() + + @override(Checkpointable) + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + state = { + NUM_ENV_STEPS_SAMPLED_LIFETIME: ( + self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME, default=0) + ), + } + + if self._check_component(COMPONENT_RL_MODULE, components, not_components): + state[COMPONENT_RL_MODULE] = self.module.get_state( + components=self._get_subcomponents(COMPONENT_RL_MODULE, components), + not_components=self._get_subcomponents( + COMPONENT_RL_MODULE, not_components + ), + **kwargs, + ) + state[WEIGHTS_SEQ_NO] = self._weights_seq_no + if self._check_component( + COMPONENT_ENV_TO_MODULE_CONNECTOR, components, not_components + ): + state[COMPONENT_ENV_TO_MODULE_CONNECTOR] = self._env_to_module.get_state() + if self._check_component( + COMPONENT_MODULE_TO_ENV_CONNECTOR, components, not_components + ): + state[COMPONENT_MODULE_TO_ENV_CONNECTOR] = self._module_to_env.get_state() + + return state + + @override(Checkpointable) + def set_state(self, state: StateDict) -> None: + if COMPONENT_ENV_TO_MODULE_CONNECTOR in state: + self._env_to_module.set_state(state[COMPONENT_ENV_TO_MODULE_CONNECTOR]) + if COMPONENT_MODULE_TO_ENV_CONNECTOR in state: + self._module_to_env.set_state(state[COMPONENT_MODULE_TO_ENV_CONNECTOR]) + + # Update the RLModule state. + if COMPONENT_RL_MODULE in state: + # A missing value for WEIGHTS_SEQ_NO or a value of 0 means: Force the + # update. + weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0) + + # Only update the weights, if this is the first synchronization or + # if the weights of this `EnvRunner` lag behind the actual ones. + if weights_seq_no == 0 or self._weights_seq_no < weights_seq_no: + rl_module_state = state[COMPONENT_RL_MODULE] + if isinstance(rl_module_state, ray.ObjectRef): + rl_module_state = ray.get(rl_module_state) + if ( + isinstance(rl_module_state, dict) + and DEFAULT_MODULE_ID in rl_module_state + ): + rl_module_state = rl_module_state[DEFAULT_MODULE_ID] + self.module.set_state(rl_module_state) + + # Update our weights_seq_no, if the new one is > 0. + if weights_seq_no > 0: + self._weights_seq_no = weights_seq_no + + # Update our lifetime counters. + if NUM_ENV_STEPS_SAMPLED_LIFETIME in state: + self.metrics.set_value( + key=NUM_ENV_STEPS_SAMPLED_LIFETIME, + value=state[NUM_ENV_STEPS_SAMPLED_LIFETIME], + reduce="sum", + with_throughput=True, + ) + + @override(Checkpointable) + def get_ctor_args_and_kwargs(self): + return ( + (), # *args + {"config": self.config}, # **kwargs + ) + + @override(Checkpointable) + def get_metadata(self): + metadata = Checkpointable.get_metadata(self) + metadata.update( + { + # TODO (sven): Maybe add serialized (JSON-writable) config here? + } + ) + return metadata + + @override(Checkpointable) + def get_checkpointable_components(self): + return [ + (COMPONENT_RL_MODULE, self.module), + (COMPONENT_ENV_TO_MODULE_CONNECTOR, self._env_to_module), + (COMPONENT_MODULE_TO_ENV_CONNECTOR, self._module_to_env), + ] + + @override(EnvRunner) + def assert_healthy(self): + """Checks that self.__init__() has been completed properly. + + Ensures that the instance has a `MultiRLModule` and an + environment defined. + + Raises: + AssertionError: If the EnvRunner Actor has NOT been properly initialized. + """ + # Make sure we have built our gym.vector.Env and RLModule properly. + assert self.env and hasattr(self, "module") + + @override(EnvRunner) + def make_env(self) -> None: + """Creates a vectorized gymnasium env and stores it in `self.env`. + + Note that users can change the EnvRunner's config (e.g. change + `self.config.env_config`) and then call this method to create new environments + with the updated configuration. + """ + # If an env already exists, try closing it first + # to allow it to properly clean up. + if self.env is not None: + try: + self.env.close() + except Exception as e: + logger.warning( + "Tried closing the existing env, but failed with error: " + f"{e.args[0]}" + ) + + env_config = self.config.env_config + if not isinstance(env_config, EnvContext): + env_ctx = EnvContext( + env_config, + worker_index=self.worker_index, + num_workers=self.num_workers, + remote=self.config.remote_worker_envs, + ) + else: + env_ctx = env_config + + # No env provided -> Error. + if not self.config.env: + raise ValueError( + "`config.env` is not provided! " + "You should provide a valid environment to your config through " + "`config.environment([env descriptor e.g. 'CartPole-v1'])`." + ) + # Register env for the local context. + # Note, `gym.register` has to be called on each worker. + elif isinstance(self.config.env, str) and _global_registry.contains( + ENV_CREATOR, self.config.env + ): + env_name = "rllib-single-agent-env-v0" + entry_point = _global_registry.get(ENV_CREATOR, self.config.env) + gym.register( + env_name, + entry_point=lambda: entry_point(env_ctx), + vector_entry_point=lambda num_envs: entry_point( + env_ctx | {"num_envs": num_envs} + ), + ) + env_config = {} + elif callable(self.config.env): + env_name = "rllib-single-agent-env-v0" + gym.register( + env_name, + entry_point=lambda: self.config.env(env_ctx), + vector_entry_point=lambda num_envs: self.config.env( + env_ctx | {"num_envs": num_envs} + ), + ) + env_config = {} + else: + env_name = self.config.env + + vectorize_mode = gym.VectorizeMode(self.config.gym_env_vectorize_mode) + try: + self.env = DictInfoToList( + gym.make_vec( + env_name, + num_envs=self.config.num_envs_per_env_runner, + vectorization_mode=vectorize_mode, + **env_config, + ) + ) + except gym.error.Error as e: + raise EnvError( + ERR_MSG_INVALID_ENV_DESCRIPTOR.format(self.config.env) + ) from e + + self.num_envs: int = self.env.num_envs + assert self.num_envs == self.config.num_envs_per_env_runner + + # Set the flag to reset all envs upon the next `sample()` call. + self._needs_initial_reset = True + + # Call the `on_environment_created` callback. + make_callback( + "on_environment_created", + callbacks_objects=self._callbacks, + callbacks_functions=self.config.callbacks_on_environment_created, + kwargs=dict( + env_runner=self, + metrics_logger=self.metrics, + env=self.env.unwrapped, + env_context=env_ctx, + ), + ) + + @override(EnvRunner) + def make_module(self): + env = self.env.unwrapped if self.env is not None else None + try: + module_spec: RLModuleSpec = self.config.get_rl_module_spec( + env=env, spaces=self.get_spaces(), inference_only=True + ) + # Build the module from its spec. + self.module = module_spec.build() + + # Move the RLModule to our device. + # TODO (sven): In order to make this framework-agnostic, we should maybe + # make the RLModule.build() method accept a device OR create an additional + # `RLModule.to()` override. + self.module.to(self._device) + + # If `AlgorithmConfig.get_rl_module_spec()` is not implemented, this env runner + # will not have an RLModule, but might still be usable with random actions. + except NotImplementedError: + self.module = None + + @override(EnvRunner) + def stop(self): + # Close our env object via gymnasium's API. + if self.env is not None: + self.env.close() + + def _reset_envs(self, episodes, shared_data, explore): + # Create n new episodes and make the `on_episode_created` callbacks. + for env_index in range(self.num_envs): + self._new_episode(env_index, episodes) + + # Erase all cached ongoing episodes (these will never be completed and + # would thus never be returned/cleaned by `get_metrics` and cause a memory + # leak). + self._ongoing_episodes_for_metrics.clear() + + # Try resetting the environment. + observations, infos = self._try_env_reset( + # Only seed (if seed provided) upon initial reset. + seed=self._seed if self._needs_initial_reset else None, + # TODO (sven): Support options? + options=None, + ) + observations = unbatch(observations) + + # Set initial obs and infos in the episodes. + for env_index in range(self.num_envs): + episodes[env_index].add_env_reset( + observation=observations[env_index], + infos=infos[env_index], + ) + + # Run the env-to-module connector to make sure the reset-obs/infos have + # properly been processed (if applicable). + self._cached_to_module = None + if self.module: + self._cached_to_module = self._env_to_module( + rl_module=self.module, + episodes=episodes, + explore=explore, + shared_data=shared_data, + metrics=self.metrics, + metrics_prefix_key=(ENV_TO_MODULE_CONNECTOR,), + ) + + # Call `on_episode_start()` callbacks (always after reset). + for env_index in range(self.num_envs): + self._make_on_episode_callback("on_episode_start", env_index, episodes) + + def _new_episode(self, env_index, episodes=None): + episodes = episodes if episodes is not None else self._episodes + episodes[env_index] = 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) + + def _make_on_episode_callback( + self, which: str, idx: int, episodes: List[SingleAgentEpisode] + ): + kwargs = dict( + episode=episodes[idx], + env_runner=self, + metrics_logger=self.metrics, + env=self.env.unwrapped, + rl_module=self.module, + env_index=idx, + ) + if which == "on_episode_end": + kwargs["prev_episode_chunks"] = self._ongoing_episodes_for_metrics[ + episodes[idx].id_ + ] + + make_callback( + which, + callbacks_objects=self._callbacks, + callbacks_functions=getattr(self.config, f"callbacks_{which}"), + kwargs=kwargs, + ) + + def _increase_sampled_metrics(self, num_steps, num_episodes_completed): + # Per sample cycle stats. + self.metrics.log_value( + NUM_ENV_STEPS_SAMPLED, num_steps, reduce="sum", clear_on_reduce=True + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_SAMPLED, DEFAULT_AGENT_ID), + num_steps, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_SAMPLED, DEFAULT_MODULE_ID), + num_steps, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_EPISODES, + num_episodes_completed, + reduce="sum", + clear_on_reduce=True, + ) + # Lifetime stats. + self.metrics.log_value( + NUM_ENV_STEPS_SAMPLED_LIFETIME, + num_steps, + reduce="sum", + with_throughput=True, + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_SAMPLED_LIFETIME, DEFAULT_AGENT_ID), + num_steps, + reduce="sum", + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_SAMPLED_LIFETIME, DEFAULT_MODULE_ID), + num_steps, + reduce="sum", + ) + self.metrics.log_value( + NUM_EPISODES_LIFETIME, + num_episodes_completed, + reduce="sum", + ) + return num_steps + + def _log_episode_metrics(self, length, ret, sec): + # Log general episode metrics. + # Use the configured window, but factor in the parallelism of the EnvRunners. + # As a result, we only log the last `window / num_env_runners` steps here, + # because everything gets parallel-merged in the Algorithm process. + win = max( + 1, + int( + math.ceil( + self.config.metrics_num_episodes_for_smoothing + / (self.config.num_env_runners or 1) + ) + ), + ) + self.metrics.log_value(EPISODE_LEN_MEAN, length, window=win) + self.metrics.log_value(EPISODE_RETURN_MEAN, ret, window=win) + self.metrics.log_value(EPISODE_DURATION_SEC_MEAN, sec, window=win) + # Per-agent returns. + self.metrics.log_value( + ("agent_episode_return_mean", DEFAULT_AGENT_ID), ret, window=win + ) + # Per-RLModule returns. + self.metrics.log_value( + ("module_episode_return_mean", DEFAULT_MODULE_ID), ret, window=win + ) + + # For some metrics, log min/max as well. + self.metrics.log_value(EPISODE_LEN_MIN, length, reduce="min", window=win) + self.metrics.log_value(EPISODE_RETURN_MIN, ret, reduce="min", window=win) + self.metrics.log_value(EPISODE_LEN_MAX, length, reduce="max", window=win) + self.metrics.log_value(EPISODE_RETURN_MAX, ret, reduce="max", window=win) + + @Deprecated( + new="SingleAgentEnvRunner.get_state(components='rl_module')", + error=True, + ) + def get_weights(self, *args, **kwargs): + pass + + @Deprecated(new="SingleAgentEnvRunner.set_state()", error=True) + def set_weights(self, *args, **kwargs): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/env/single_agent_episode.py b/lib/python3.12/site-packages/ray/rllib/env/single_agent_episode.py new file mode 100644 index 0000000000000000000000000000000000000000..dffe8affa04dcca650bbccf3d1f04852a3a3373a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/single_agent_episode.py @@ -0,0 +1,1836 @@ +import copy +import functools +import time +import uuid +from collections import defaultdict +from typing import Any, Dict, List, Optional, SupportsFloat, Union + +import gymnasium as gym +import numpy as np +import tree +from gymnasium.core import ActType, ObsType + +from ray._common.deprecation import Deprecated +from ray.rllib.core.columns import Columns +from ray.rllib.env.utils.infinite_lookback_buffer import InfiniteLookbackBuffer +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.serialization import gym_space_from_dict, gym_space_to_dict +from ray.rllib.utils.typing import AgentID, ModuleID +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class SingleAgentEpisode: + """A class representing RL environment episodes for individual agents. + + SingleAgentEpisode stores observations, info dicts, actions, rewards, and all + module outputs (e.g. state outs, action logp, etc..) for an individual agent within + some single-agent or multi-agent environment. + The two main APIs to add data to an ongoing episode are the `add_env_reset()` + and `add_env_step()` methods, which should be called passing the outputs of the + respective gym.Env API calls: `env.reset()` and `env.step()`. + + A SingleAgentEpisode might also only represent a chunk of an episode, which is + useful for cases, in which partial (non-complete episode) sampling is performed + and collected episode data has to be returned before the actual gym.Env episode has + finished (see `SingleAgentEpisode.cut()`). In order to still maintain visibility + onto past experiences within such a "cut" episode, SingleAgentEpisode instances + can have a "lookback buffer" of n timesteps at their beginning (left side), which + solely exists for the purpose of compiling extra data (e.g. "prev. reward"), but + is not considered part of the finished/packaged episode (b/c the data in the + lookback buffer is already part of a previous episode chunk). + + Powerful getter methods, such as `get_observations()` help collect different types + of data from the episode at individual time indices or time ranges, including the + "lookback buffer" range described above. For example, to extract the last 4 rewards + of an ongoing episode, one can call `self.get_rewards(slice(-4, None))` or + `self.rewards[-4:]`. This would work, even if the ongoing SingleAgentEpisode is + a continuation chunk from a much earlier started episode, as long as it has a + lookback buffer size of sufficient size. + + Examples: + + .. testcode:: + + import gymnasium as gym + import numpy as np + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + + # Construct a new episode (without any data in it yet). + episode = SingleAgentEpisode() + assert len(episode) == 0 + + # Fill the episode with some data (10 timesteps). + env = gym.make("CartPole-v1") + obs, infos = env.reset() + episode.add_env_reset(obs, infos) + + # Even with the initial obs/infos, the episode is still considered len=0. + assert len(episode) == 0 + for _ in range(5): + action = env.action_space.sample() + obs, reward, term, trunc, infos = env.step(action) + episode.add_env_step( + observation=obs, + action=action, + reward=reward, + terminated=term, + truncated=trunc, + infos=infos, + ) + assert len(episode) == 5 + + # We can now access information from the episode via the getter APIs. + + # Get the last 3 rewards (in a batch of size 3). + episode.get_rewards(slice(-3, None)) # same as `episode.rewards[-3:]` + + # Get the most recent action (single item, not batched). + # This works regardless of the action space or whether the episode has + # been numpy'ized or not (see below). + episode.get_actions(-1) # same as episode.actions[-1] + + # Looking back from ts=1, get the previous 4 rewards AND fill with 0.0 + # in case we go over the beginning (ts=0). So we would expect + # [0.0, 0.0, 0.0, r0] to be returned here, where r0 is the very first received + # reward in the episode: + episode.get_rewards(slice(-4, 0), neg_index_as_lookback=True, fill=0.0) + + # Note the use of fill=0.0 here (fill everything that's out of range with this + # value) AND the argument `neg_index_as_lookback=True`, which interprets + # negative indices as being left of ts=0 (e.g. -1 being the timestep before + # ts=0). + + # Assuming we had a complex action space (nested gym.spaces.Dict) with one or + # more elements being Discrete or MultiDiscrete spaces: + # 1) The `fill=...` argument would still work, filling all spaces (Boxes, + # Discrete) with that provided value. + # 2) Setting the flag `one_hot_discrete=True` would convert those discrete + # sub-components automatically into one-hot (or multi-one-hot) tensors. + # This simplifies the task of having to provide the previous 4 (nested and + # partially discrete/multi-discrete) actions for each timestep within a training + # batch, thereby filling timesteps before the episode started with 0.0s and + # one-hot'ing the discrete/multi-discrete components in these actions: + episode = SingleAgentEpisode(action_space=gym.spaces.Dict({ + "a": gym.spaces.Discrete(3), + "b": gym.spaces.MultiDiscrete([2, 3]), + "c": gym.spaces.Box(-1.0, 1.0, (2,)), + })) + + # ... fill episode with data ... + episode.add_env_reset(observation=0) + # ... from a few steps. + episode.add_env_step( + observation=1, + action={"a":0, "b":np.array([1, 2]), "c":np.array([.5, -.5], np.float32)}, + reward=1.0, + ) + + # In your connector + prev_4_a = [] + # Note here that len(episode) does NOT include the lookback buffer. + for ts in range(len(episode)): + prev_4_a.append( + episode.get_actions( + indices=slice(ts - 4, ts), + # Make sure negative indices are interpreted as + # "into lookback buffer" + neg_index_as_lookback=True, + # Zero-out everything even further before the lookback buffer. + fill=0.0, + # Take care of discrete components (get ready as NN input). + one_hot_discrete=True, + ) + ) + + # Finally, convert from list of batch items to a struct (same as action space) + # of batched (numpy) arrays, in which all leafs have B==len(prev_4_a). + from ray.rllib.utils.spaces.space_utils import batch + + prev_4_actions_col = batch(prev_4_a) + """ + + __slots__ = ( + "actions", + "agent_id", + "extra_model_outputs", + "id_", + "infos", + "is_terminated", + "is_truncated", + "module_id", + "multi_agent_episode_id", + "observations", + "rewards", + "t", + "t_started", + "_action_space", + "_last_added_observation", + "_last_added_infos", + "_last_step_time", + "_observation_space", + "_start_time", + "_custom_data", + ) + + def __init__( + self, + id_: Optional[str] = None, + *, + observations: Optional[Union[List[ObsType], InfiniteLookbackBuffer]] = None, + observation_space: Optional[gym.Space] = None, + infos: Optional[Union[List[Dict], InfiniteLookbackBuffer]] = None, + actions: Optional[Union[List[ActType], InfiniteLookbackBuffer]] = None, + action_space: Optional[gym.Space] = None, + rewards: Optional[Union[List[SupportsFloat], InfiniteLookbackBuffer]] = None, + terminated: bool = False, + truncated: bool = False, + extra_model_outputs: Optional[Dict[str, Any]] = None, + t_started: Optional[int] = None, + len_lookback_buffer: Union[int, str] = "auto", + agent_id: Optional[AgentID] = None, + module_id: Optional[ModuleID] = None, + multi_agent_episode_id: Optional[int] = None, + ): + """Initializes a SingleAgentEpisode instance. + + This constructor can be called with or without already sampled data, part of + which might then go into the lookback buffer. + + Args: + id_: Unique identifier for this episode. If no ID is provided the + constructor generates a unique hexadecimal code for the id. + observations: Either a list of individual observations from a sampling or + an already instantiated `InfiniteLookbackBuffer` object (possibly + with observation data in it). If a list, will construct the buffer + automatically (given the data and the `len_lookback_buffer` argument). + observation_space: An optional gym.Space, which all individual observations + should abide to. If not None and this SingleAgentEpisode is numpy'ized + (via the `self.to_numpy()` method), and data is appended or set, the new + data will be checked for correctness. + infos: Either a list of individual info dicts from a sampling or + an already instantiated `InfiniteLookbackBuffer` object (possibly + with info dicts in it). If a list, will construct the buffer + automatically (given the data and the `len_lookback_buffer` argument). + actions: Either a list of individual info dicts from a sampling or + an already instantiated `InfiniteLookbackBuffer` object (possibly + with info dict] data in it). If a list, will construct the buffer + automatically (given the data and the `len_lookback_buffer` argument). + action_space: An optional gym.Space, which all individual actions + should abide to. If not None and this SingleAgentEpisode is numpy'ized + (via the `self.to_numpy()` method), and data is appended or set, the new + data will be checked for correctness. + rewards: Either a list of individual rewards from a sampling or + an already instantiated `InfiniteLookbackBuffer` object (possibly + with reward data in it). If a list, will construct the buffer + automatically (given the data and the `len_lookback_buffer` argument). + extra_model_outputs: A dict mapping string keys to either lists of + individual extra model output tensors (e.g. `action_logp` or + `state_outs`) from a sampling or to already instantiated + `InfiniteLookbackBuffer` object (possibly with extra model output data + in it). If mapping is to lists, will construct the buffers automatically + (given the data and the `len_lookback_buffer` argument). + terminated: A boolean indicating, if the episode is already terminated. + truncated: A boolean indicating, if the episode has been truncated. + t_started: Optional. The starting timestep of the episode. The default + is zero. If data is provided, the starting point is from the last + observation onwards (i.e. `t_started = len(observations) - 1`). If + this parameter is provided the episode starts at the provided value. + len_lookback_buffer: The size of the (optional) lookback buffers to keep in + front of this Episode for each type of data (observations, actions, + etc..). If larger than 0, the first `len_lookback_buffer` + items of each type of data are interpreted as NOT part of this actual + episode chunk, but instead serve as "historical" record that may be + viewed and used to derive new data from. For example, it might be + necessary to have a lookback buffer of four if you would like to do + observation frame stacking and your episode has been cut and you're now + operating on a new chunk (continuing from the cut one). Then, for the + first 3 items, you would have to be able to look back into the old + chunk's data. + If `len_lookback_buffer` is "auto" (default), will interpret all + provided data in the constructor as part of the lookback buffers. + agent_id: An optional AgentID indicating which agent this episode belongs + to. This information is stored under `self.agent_id` and only serves + reference purposes. + module_id: An optional ModuleID indicating which RLModule this episode + belongs to. Normally, this information is obtained by querying an + `agent_to_module_mapping_fn` with a given agent ID. This information + is stored under `self.module_id` and only serves reference purposes. + multi_agent_episode_id: An optional EpisodeID of the encapsulating + `MultiAgentEpisode` that this `SingleAgentEpisode` belongs to. + """ + self.id_ = id_ or uuid.uuid4().hex + + self.agent_id = agent_id + self.module_id = module_id + self.multi_agent_episode_id = multi_agent_episode_id + + # Lookback buffer length is not provided. Interpret already given data as + # lookback buffer lengths for all data types. + len_rewards = len(rewards) if rewards is not None else 0 + if len_lookback_buffer == "auto" or len_lookback_buffer > len_rewards: + len_lookback_buffer = len_rewards + + infos = infos or [{} for _ in range(len(observations or []))] + + # Observations: t0 (initial obs) to T. + self._observation_space = None + if isinstance(observations, InfiniteLookbackBuffer): + self.observations = observations + else: + self.observations = InfiniteLookbackBuffer( + data=observations, + lookback=len_lookback_buffer, + ) + self.observation_space = observation_space + + # Infos: t0 (initial info) to T. + if isinstance(infos, InfiniteLookbackBuffer): + self.infos = infos + else: + self.infos = InfiniteLookbackBuffer( + data=infos, + lookback=len_lookback_buffer, + ) + + # Actions: t1 to T. + self._action_space = None + if isinstance(actions, InfiniteLookbackBuffer): + self.actions = actions + else: + self.actions = InfiniteLookbackBuffer( + data=actions, + lookback=len_lookback_buffer, + ) + self.action_space = action_space + + # Rewards: t1 to T. + if isinstance(rewards, InfiniteLookbackBuffer): + self.rewards = rewards + else: + self.rewards = InfiniteLookbackBuffer( + data=rewards, + lookback=len_lookback_buffer, + space=gym.spaces.Box(float("-inf"), float("inf"), (), np.float32), + ) + + # obs[-1] is the final observation in the episode. + self.is_terminated = terminated + # obs[-1] is the last obs in a truncated-by-the-env episode (there will no more + # observations in following chunks for this episode). + self.is_truncated = truncated + + # Extra model outputs, e.g. `action_dist_input` needed in the batch. + self.extra_model_outputs = {} + for k, v in (extra_model_outputs or {}).items(): + if isinstance(v, InfiniteLookbackBuffer): + self.extra_model_outputs[k] = v + else: + # We cannot use the defaultdict's own constructor here as this would + # auto-set the lookback buffer to 0 (there is no data passed to that + # constructor). Then, when we manually have to set the data property, + # the lookback buffer would still be (incorrectly) 0. + self.extra_model_outputs[k] = InfiniteLookbackBuffer( + data=v, lookback=len_lookback_buffer + ) + + # The (global) timestep when this episode (possibly an episode chunk) started, + # excluding a possible lookback buffer. + self.t_started = t_started or 0 + # The current (global) timestep in the episode (possibly an episode chunk). + self.t = len(self.rewards) + self.t_started + + # Cache for custom data. May be used to store custom metrics from within a + # callback for the ongoing episode (e.g. render images). + self._custom_data = {} + + # Keep timer stats on deltas between steps. + self._start_time = None + self._last_step_time = None + + self._last_added_observation = None + self._last_added_infos = None + + # Validate the episode data thus far. + self.validate() + + def add_env_reset( + self, + observation: ObsType, + infos: Optional[Dict] = None, + ) -> None: + """Adds the initial data (after an `env.reset()`) to the episode. + + This data consists of initial observations and initial infos. + + Args: + observation: The initial observation returned by `env.reset()`. + infos: An (optional) info dict returned by `env.reset()`. + """ + assert not self.is_reset + assert not self.is_done + assert len(self.observations) == 0 + # Assume that this episode is completely empty and has not stepped yet. + # Leave self.t (and self.t_started) at 0. + assert self.t == self.t_started == 0 + + infos = infos or {} + + self.observations.append(observation) + self.infos.append(infos) + + self._last_added_observation = observation + self._last_added_infos = infos + + # Validate our data. + self.validate() + + # Start the timer for this episode. + self._start_time = time.perf_counter() + + def add_env_step( + self, + observation: ObsType, + action: ActType, + reward: SupportsFloat, + infos: Optional[Dict[str, Any]] = None, + *, + terminated: bool = False, + truncated: bool = False, + extra_model_outputs: Optional[Dict[str, Any]] = None, + ) -> None: + """Adds results of an `env.step()` call (including the action) to this episode. + + This data consists of an observation and info dict, an action, a reward, + terminated/truncated flags, and extra model outputs (e.g. action probabilities + or RNN internal state outputs). + + Args: + observation: The next observation received from the environment after(!) + taking `action`. + action: The last action used by the agent during the call to `env.step()`. + reward: The last reward received by the agent after taking `action`. + infos: The last info received from the environment after taking `action`. + terminated: A boolean indicating, if the environment has been + terminated (after taking `action`). + truncated: A boolean indicating, if the environment has been + truncated (after taking `action`). + extra_model_outputs: The last timestep's specific model outputs. + These are normally outputs of an RLModule that were computed along with + `action`, e.g. `action_logp` or `action_dist_inputs`. + """ + # Cannot add data to an already done episode. + assert ( + not self.is_done + ), "The agent is already done: no data can be added to its episode." + + self.observations.append(observation) + self.actions.append(action) + self.rewards.append(reward) + infos = infos or {} + self.infos.append(infos) + self.t += 1 + if extra_model_outputs is not None: + for k, v in extra_model_outputs.items(): + if k not in self.extra_model_outputs: + self.extra_model_outputs[k] = InfiniteLookbackBuffer([v]) + else: + self.extra_model_outputs[k].append(v) + self.is_terminated = terminated + self.is_truncated = truncated + + self._last_added_observation = observation + self._last_added_infos = infos + + # Only check spaces if numpy'ized AND every n timesteps. + if self.is_numpy and self.t % 100: + if self.observation_space is not None: + assert self.observation_space.contains(observation), ( + f"`observation` {observation} does NOT fit SingleAgentEpisode's " + f"observation_space: {self.observation_space}!" + ) + if self.action_space is not None: + assert self.action_space.contains(action), ( + f"`action` {action} does NOT fit SingleAgentEpisode's " + f"action_space: {self.action_space}!" + ) + + # Validate our data. + self.validate() + + # Step time stats. + self._last_step_time = time.perf_counter() + if self._start_time is None: + self._start_time = self._last_step_time + + def validate(self) -> None: + """Validates the episode's data. + + This function ensures that the data stored to a `SingleAgentEpisode` is + in order (e.g. that the correct number of observations, actions, rewards + are there). + """ + assert len(self.observations) == len(self.infos) + if len(self.observations) == 0: + assert len(self.infos) == len(self.rewards) == len(self.actions) == 0 + for k, v in self.extra_model_outputs.items(): + assert len(v) == 0, (k, v, v.data, len(v)) + # Make sure we always have one more obs stored than rewards (and actions) + # due to the reset/last-obs logic of an MDP. + else: + assert ( + len(self.observations) + == len(self.infos) + == len(self.rewards) + 1 + == len(self.actions) + 1 + ), ( + len(self.observations), + len(self.infos), + len(self.rewards), + len(self.actions), + ) + for k, v in self.extra_model_outputs.items(): + assert len(v) == len(self.observations) - 1 + + @property + def custom_data(self): + return self._custom_data + + @property + def is_reset(self) -> bool: + """Returns True if `self.add_env_reset()` has already been called.""" + return len(self.observations) > 0 + + @property + def is_numpy(self) -> bool: + """True, if the data in this episode is already stored as numpy arrays.""" + # If rewards are still a list, return False. + # Otherwise, rewards should already be a (1D) numpy array. + return self.rewards.finalized + + @property + def is_done(self) -> bool: + """Whether the episode is actually done (terminated or truncated). + + A done episode cannot be continued via `self.add_timestep()` or being + concatenated on its right-side with another episode chunk or being + succeeded via `self.create_successor()`. + """ + return self.is_terminated or self.is_truncated + + def to_numpy(self) -> "SingleAgentEpisode": + """Converts this Episode's list attributes to numpy arrays. + + This means in particular that this episodes' lists of (possibly complex) + data (e.g. if we have a dict obs space) will be converted to (possibly complex) + structs, whose leafs are now numpy arrays. Each of these leaf numpy arrays will + have the same length (batch dimension) as the length of the original lists. + + Note that the data under the Columns.INFOS are NEVER numpy'ized and will remain + a list (normally, a list of the original, env-returned dicts). This is due to + the herterogenous nature of INFOS returned by envs, which would make it unwieldy + to convert this information to numpy arrays. + + After calling this method, no further data may be added to this episode via + the `self.add_env_step()` method. + + Examples: + + .. testcode:: + + import numpy as np + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + + episode = SingleAgentEpisode( + observations=[0, 1, 2, 3], + actions=[1, 2, 3], + rewards=[1, 2, 3], + # Note: terminated/truncated have nothing to do with an episode + # being numpy'ized or not (via the `self.to_numpy()` method)! + terminated=False, + len_lookback_buffer=0, # no lookback; all data is actually "in" episode + ) + # Episode has not been numpy'ized yet. + assert not episode.is_numpy + # We are still operating on lists. + assert episode.get_observations([1]) == [1] + assert episode.get_observations(slice(None, 2)) == [0, 1] + # We can still add data (and even add the terminated=True flag). + episode.add_env_step( + observation=4, + action=4, + reward=4, + terminated=True, + ) + # Still NOT numpy'ized. + assert not episode.is_numpy + + # Numpy'ized the episode. + episode.to_numpy() + assert episode.is_numpy + + # We cannot add data anymore. The following would crash. + # episode.add_env_step(observation=5, action=5, reward=5) + + # Everything is now numpy arrays (with 0-axis of size + # B=[len of requested slice]). + assert isinstance(episode.get_observations([1]), np.ndarray) # B=1 + assert isinstance(episode.actions[0:2], np.ndarray) # B=2 + assert isinstance(episode.rewards[1:4], np.ndarray) # B=3 + + Returns: + This `SingleAgentEpisode` object with the converted numpy data. + """ + + self.observations.finalize() + if len(self) > 0: + self.actions.finalize() + self.rewards.finalize() + for k, v in self.extra_model_outputs.items(): + self.extra_model_outputs[k].finalize() + + return self + + def concat_episode(self, other: "SingleAgentEpisode") -> None: + """Adds the given `other` SingleAgentEpisode to the right side of `self`. + + In order for this to work, both chunks (`self` and `other`) must fit + together. This is checked by the IDs (must be identical), the time step counters + (`self.env_t` must be the same as `episode_chunk.env_t_started`), as well as the + observations/infos at the concatenation boundaries. Also, `self.is_done` must + not be True, meaning `self.is_terminated` and `self.is_truncated` are both + False. + + Args: + other: The other `SingleAgentEpisode` to be concatenated to this one. + + Returns: + A `SingleAgentEpisode` instance containing the concatenated data + from both episodes (`self` and `other`). + """ + assert other.id_ == self.id_ + # NOTE (sven): This is what we agreed on. As the replay buffers must be + # able to concatenate. + assert not self.is_done + # Make sure the timesteps match. + assert self.t == other.t_started + # Validate `other`. + other.validate() + + # Make sure, end matches other episode chunk's beginning. + tree.assert_same_structure(other.observations[0], self.observations[-1]) + # Use tree.map_structure with np.array_equal to check every leaf node are equivalent + # then np.all on flatten to validate all are tree + assert np.all( + tree.flatten( + tree.map_structure( + np.array_equal, other.observations[0], self.observations[-1] + ) + ) + ) + + # Pop out our last observations and infos (as these are identical + # to the first obs and infos in the next episode). + self.observations.pop() + self.infos.pop() + + # Extend ourselves. In case, episode_chunk is already terminated and numpy'ized + # we need to convert to lists (as we are ourselves still filling up lists). + self.observations.extend(other.get_observations()) + self.actions.extend(other.get_actions()) + self.rewards.extend(other.get_rewards()) + self.infos.extend(other.get_infos()) + self.t = other.t + + if other.is_terminated: + self.is_terminated = True + elif other.is_truncated: + self.is_truncated = True + + for key in other.extra_model_outputs.keys(): + assert key in self.extra_model_outputs + self.extra_model_outputs[key].extend(other.get_extra_model_outputs(key)) + + # Merge with `other`'s custom_data, but give `other` priority b/c we assume + # that as a follow-up chunk of `self` other has a more complete version of + # `custom_data`. + self.custom_data.update(other.custom_data) + + # Validate. + self.validate() + + def cut(self, len_lookback_buffer: int = 0) -> "SingleAgentEpisode": + """Returns a successor episode chunk (of len=0) continuing from this Episode. + + The successor will have the same ID as `self`. + If no lookback buffer is requested (len_lookback_buffer=0), the successor's + observations will be the last observation(s) of `self` and its length will + therefore be 0 (no further steps taken yet). If `len_lookback_buffer` > 0, + the returned successor will have `len_lookback_buffer` observations (and + actions, rewards, etc..) taken from the right side (end) of `self`. For example + if `len_lookback_buffer=2`, the returned successor's lookback buffer actions + will be identical to `self.actions[-2:]`. + + This method is useful if you would like to discontinue building an episode + chunk (b/c you have to return it from somewhere), but would like to have a new + episode instance to continue building the actual gym.Env episode at a later + time. Vie the `len_lookback_buffer` argument, the continuing chunk (successor) + will still be able to "look back" into this predecessor episode's data (at + least to some extend, depending on the value of `len_lookback_buffer`). + + Args: + len_lookback_buffer: The number of timesteps to take along into the new + chunk as "lookback buffer". A lookback buffer is additional data on + the left side of the actual episode data for visibility purposes + (but without actually being part of the new chunk). For example, if + `self` ends in actions 5, 6, 7, and 8, and we call + `self.cut(len_lookback_buffer=2)`, the returned chunk will have + actions 7 and 8 already in it, but still `t_started`==t==8 (not 7!) and + a length of 0. If there is not enough data in `self` yet to fulfil + the `len_lookback_buffer` request, the value of `len_lookback_buffer` + is automatically adjusted (lowered). + + Returns: + The successor Episode chunk of this one with the same ID and state and the + only observation being the last observation in self. + """ + assert not self.is_done and len_lookback_buffer >= 0 + + # Initialize this chunk with the most recent obs and infos (even if lookback is + # 0). Similar to an initial `env.reset()`. + indices_obs_and_infos = slice(-len_lookback_buffer - 1, None) + indices_rest = ( + slice(-len_lookback_buffer, None) + if len_lookback_buffer > 0 + else slice(None, 0) + ) + + sa_episode = SingleAgentEpisode( + # Same ID. + id_=self.id_, + observations=self.get_observations(indices=indices_obs_and_infos), + observation_space=self.observation_space, + infos=self.get_infos(indices=indices_obs_and_infos), + actions=self.get_actions(indices=indices_rest), + action_space=self.action_space, + rewards=self.get_rewards(indices=indices_rest), + extra_model_outputs={ + k: self.get_extra_model_outputs(k, indices_rest) + for k in self.extra_model_outputs.keys() + }, + # Continue with self's current timestep. + t_started=self.t, + # Use the length of the provided data as lookback buffer. + len_lookback_buffer="auto", + ) + # Deepcopy all custom data in `self` to be continued in the cut episode. + sa_episode._custom_data = copy.deepcopy(self.custom_data) + + return sa_episode + + # TODO (sven): Distinguish between: + # - global index: This is the absolute, global timestep whose values always + # start from 0 (at the env reset). So doing get_observations(0, global_ts=True) + # should always return the exact 1st observation (reset obs), no matter what. In + # case we are in an episode chunk and `fill` or a sufficient lookback buffer is + # provided, this should yield a result. Otherwise, error. + # - global index=False -> indices are relative to the chunk start. If a chunk has + # t_started=6 and we ask for index=0, then return observation at timestep 6 + # (t_started). + def get_observations( + self, + indices: Optional[Union[int, List[int], slice]] = None, + *, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + one_hot_discrete: bool = False, + ) -> Any: + """Returns individual observations or batched ranges thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual observation stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual observations in a batch of size len(indices). + A slice object is interpreted as a range of observations to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all observations (from ts=0 to the end). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with observations [4, 5, 6, 7, 8, 9], + where [4, 5, 6] is the lookback buffer range (ts=0 item is 7), will + respond to `get_observations(-1, neg_index_as_lookback=True)` + with `6` and to + `get_observations(slice(-2, 1), neg_index_as_lookback=True)` with + `[5, 6, 7]`. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with observations [10, 11, 12, 13, 14] and + lookback buffer size of 2 (meaning observations `10` and `11` are part + of the lookback buffer) will respond to + `get_observations(slice(-7, -2), fill=0.0)` with + `[0.0, 0.0, 10, 11, 12]`. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) observation + space that are Discrete or MultiDiscrete. Note that if `fill=0` and the + requested `indices` are out of the range of our data, the returned + one-hot vectors will actually be zero-hot (all slots zero). + + Examples: + + .. testcode:: + + import gymnasium as gym + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + episode = SingleAgentEpisode( + # Discrete(4) observations (ints between 0 and 4 (excl.)) + observation_space=gym.spaces.Discrete(4), + observations=[0, 1, 2, 3], + actions=[1, 2, 3], rewards=[1, 2, 3], # <- not relevant for this demo + len_lookback_buffer=0, # no lookback; all data is actually "in" episode + ) + # Plain usage (`indices` arg only). + check(episode.get_observations(-1), 3) + check(episode.get_observations(0), 0) + check(episode.get_observations([0, 2]), [0, 2]) + check(episode.get_observations([-1, 0]), [3, 0]) + check(episode.get_observations(slice(None, 2)), [0, 1]) + check(episode.get_observations(slice(-2, None)), [2, 3]) + # Using `fill=...` (requesting slices beyond the boundaries). + check(episode.get_observations(slice(-6, -2), fill=-9), [-9, -9, 0, 1]) + check(episode.get_observations(slice(2, 5), fill=-7), [2, 3, -7]) + # Using `one_hot_discrete=True`. + check(episode.get_observations(2, one_hot_discrete=True), [0, 0, 1, 0]) + check(episode.get_observations(3, one_hot_discrete=True), [0, 0, 0, 1]) + check(episode.get_observations( + slice(0, 3), + one_hot_discrete=True, + ), [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0]]) + # Special case: Using `fill=0.0` AND `one_hot_discrete=True`. + check(episode.get_observations( + -1, + neg_index_as_lookback=True, # -1 means one left of ts=0 + fill=0.0, + one_hot_discrete=True, + ), [0, 0, 0, 0]) # <- all 0s one-hot tensor (note difference to [1 0 0 0]!) + + Returns: + The collected observations. + As a 0-axis batch, if there are several `indices` or a list of exactly one + index provided OR `indices` is a slice object. + As single item (B=0 -> no additional 0-axis) if `indices` is a single int. + """ + return self.observations.get( + indices=indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + one_hot_discrete=one_hot_discrete, + ) + + def get_infos( + self, + indices: Optional[Union[int, List[int], slice]] = None, + *, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + ) -> Any: + """Returns individual info dicts or list (ranges) thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual info dict stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual info dicts in a list of size len(indices). + A slice object is interpreted as a range of info dicts to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all infos (from ts=0 to the end). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with infos + [{"l":4}, {"l":5}, {"l":6}, {"a":7}, {"b":8}, {"c":9}], where the + first 3 items are the lookback buffer (ts=0 item is {"a": 7}), will + respond to `get_infos(-1, neg_index_as_lookback=True)` with + `{"l":6}` and to + `get_infos(slice(-2, 1), neg_index_as_lookback=True)` with + `[{"l":5}, {"l":6}, {"a":7}]`. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to + auto-fill. For example, an episode with infos + [{"l":10}, {"l":11}, {"a":12}, {"b":13}, {"c":14}] and lookback buffer + size of 2 (meaning infos {"l":10}, {"l":11} are part of the lookback + buffer) will respond to `get_infos(slice(-7, -2), fill={"o": 0.0})` + with `[{"o":0.0}, {"o":0.0}, {"l":10}, {"l":11}, {"a":12}]`. + + Examples: + + .. testcode:: + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + + episode = SingleAgentEpisode( + infos=[{"a":0}, {"b":1}, {"c":2}, {"d":3}], + # The following is needed, but not relevant for this demo. + observations=[0, 1, 2, 3], actions=[1, 2, 3], rewards=[1, 2, 3], + len_lookback_buffer=0, # no lookback; all data is actually "in" episode + ) + # Plain usage (`indices` arg only). + episode.get_infos(-1) # {"d":3} + episode.get_infos(0) # {"a":0} + episode.get_infos([0, 2]) # [{"a":0},{"c":2}] + episode.get_infos([-1, 0]) # [{"d":3},{"a":0}] + episode.get_infos(slice(None, 2)) # [{"a":0},{"b":1}] + episode.get_infos(slice(-2, None)) # [{"c":2},{"d":3}] + # Using `fill=...` (requesting slices beyond the boundaries). + # TODO (sven): This would require a space being provided. Maybe we can + # skip this check for infos, which don't have a space anyways. + # episode.get_infos(slice(-5, -3), fill={"o":-1}) # [{"o":-1},{"a":0}] + # episode.get_infos(slice(3, 5), fill={"o":-2}) # [{"d":3},{"o":-2}] + + Returns: + The collected info dicts. + As a 0-axis batch, if there are several `indices` or a list of exactly one + index provided OR `indices` is a slice object. + As single item (B=0 -> no additional 0-axis) if `indices` is a single int. + """ + return self.infos.get( + indices=indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + ) + + def get_actions( + self, + indices: Optional[Union[int, List[int], slice]] = None, + *, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + one_hot_discrete: bool = False, + ) -> Any: + """Returns individual actions or batched ranges thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual action stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual actions in a batch of size len(indices). + A slice object is interpreted as a range of actions to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all actions (from ts=0 to the end). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with actions [4, 5, 6, 7, 8, 9], where + [4, 5, 6] is the lookback buffer range (ts=0 item is 7), will respond + to `get_actions(-1, neg_index_as_lookback=True)` with `6` and + to `get_actions(slice(-2, 1), neg_index_as_lookback=True)` with + `[5, 6, 7]`. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with actions [10, 11, 12, 13, 14] and + lookback buffer size of 2 (meaning actions `10` and `11` are part + of the lookback buffer) will respond to + `get_actions(slice(-7, -2), fill=0.0)` with `[0.0, 0.0, 10, 11, 12]`. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) action + space that are Discrete or MultiDiscrete. Note that if `fill=0` and the + requested `indices` are out of the range of our data, the returned + one-hot vectors will actually be zero-hot (all slots zero). + + Examples: + + .. testcode:: + + import gymnasium as gym + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + + episode = SingleAgentEpisode( + # Discrete(4) actions (ints between 0 and 4 (excl.)) + action_space=gym.spaces.Discrete(4), + actions=[1, 2, 3], + observations=[0, 1, 2, 3], rewards=[1, 2, 3], # <- not relevant here + len_lookback_buffer=0, # no lookback; all data is actually "in" episode + ) + # Plain usage (`indices` arg only). + episode.get_actions(-1) # 3 + episode.get_actions(0) # 1 + episode.get_actions([0, 2]) # [1, 3] + episode.get_actions([-1, 0]) # [3, 1] + episode.get_actions(slice(None, 2)) # [1, 2] + episode.get_actions(slice(-2, None)) # [2, 3] + # Using `fill=...` (requesting slices beyond the boundaries). + episode.get_actions(slice(-5, -2), fill=-9) # [-9, -9, 1, 2] + episode.get_actions(slice(1, 5), fill=-7) # [2, 3, -7, -7] + # Using `one_hot_discrete=True`. + episode.get_actions(1, one_hot_discrete=True) # [0 0 1 0] (action=2) + episode.get_actions(2, one_hot_discrete=True) # [0 0 0 1] (action=3) + episode.get_actions( + slice(0, 2), + one_hot_discrete=True, + ) # [[0 1 0 0], [0 0 0 1]] (actions=1 and 3) + # Special case: Using `fill=0.0` AND `one_hot_discrete=True`. + episode.get_actions( + -1, + neg_index_as_lookback=True, # -1 means one left of ts=0 + fill=0.0, + one_hot_discrete=True, + ) # [0 0 0 0] <- all 0s one-hot tensor (note difference to [1 0 0 0]!) + + Returns: + The collected actions. + As a 0-axis batch, if there are several `indices` or a list of exactly one + index provided OR `indices` is a slice object. + As single item (B=0 -> no additional 0-axis) if `indices` is a single int. + """ + return self.actions.get( + indices=indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + one_hot_discrete=one_hot_discrete, + ) + + def get_rewards( + self, + indices: Optional[Union[int, List[int], slice]] = None, + *, + neg_index_as_lookback: bool = False, + fill: Optional[float] = None, + ) -> Any: + """Returns individual rewards or batched ranges thereof from this episode. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual reward stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual rewards in a batch of size len(indices). + A slice object is interpreted as a range of rewards to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + If None, will return all rewards (from ts=0 to the end). + neg_index_as_lookback: Negative values in `indices` are interpreted as + as "before ts=0", meaning going back into the lookback buffer. + For example, an episode with rewards [4, 5, 6, 7, 8, 9], where + [4, 5, 6] is the lookback buffer range (ts=0 item is 7), will respond + to `get_rewards(-1, neg_index_as_lookback=True)` with `6` and + to `get_rewards(slice(-2, 1), neg_index_as_lookback=True)` with + `[5, 6, 7]`. + fill: An optional float value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with rewards [10, 11, 12, 13, 14] and + lookback buffer size of 2 (meaning rewards `10` and `11` are part + of the lookback buffer) will respond to + `get_rewards(slice(-7, -2), fill=0.0)` with `[0.0, 0.0, 10, 11, 12]`. + + Examples: + + .. testcode:: + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + + episode = SingleAgentEpisode( + rewards=[1.0, 2.0, 3.0], + observations=[0, 1, 2, 3], actions=[1, 2, 3], # <- not relevant here + len_lookback_buffer=0, # no lookback; all data is actually "in" episode + ) + # Plain usage (`indices` arg only). + episode.get_rewards(-1) # 3.0 + episode.get_rewards(0) # 1.0 + episode.get_rewards([0, 2]) # [1.0, 3.0] + episode.get_rewards([-1, 0]) # [3.0, 1.0] + episode.get_rewards(slice(None, 2)) # [1.0, 2.0] + episode.get_rewards(slice(-2, None)) # [2.0, 3.0] + # Using `fill=...` (requesting slices beyond the boundaries). + episode.get_rewards(slice(-5, -2), fill=0.0) # [0.0, 0.0, 1.0, 2.0] + episode.get_rewards(slice(1, 5), fill=0.0) # [2.0, 3.0, 0.0, 0.0] + + Returns: + The collected rewards. + As a 0-axis batch, if there are several `indices` or a list of exactly one + index provided OR `indices` is a slice object. + As single item (B=0 -> no additional 0-axis) if `indices` is a single int. + """ + return self.rewards.get( + indices=indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + ) + + def get_extra_model_outputs( + self, + key: str, + indices: Optional[Union[int, List[int], slice]] = None, + *, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + ) -> Any: + """Returns extra model outputs (under given key) from this episode. + + Args: + key: The `key` within `self.extra_model_outputs` to extract data for. + indices: A single int is interpreted as an index, from which to return an + individual extra model output stored under `key` at index. + A list of ints is interpreted as a list of indices from which to gather + individual actions in a batch of size len(indices). + A slice object is interpreted as a range of extra model outputs to be + returned. Thereby, negative indices by default are interpreted as + "before the end" unless the `neg_index_as_lookback=True` option is + used, in which case negative indices are interpreted as "before ts=0", + meaning going back into the lookback buffer. + If None, will return all extra model outputs (from ts=0 to the end). + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + extra_model_outputs['a'] = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will respond to + `get_extra_model_outputs("a", -1, neg_index_as_lookback=True)` with + `6` and to `get_extra_model_outputs("a", slice(-2, 1), + neg_index_as_lookback=True)` with `[5, 6, 7]`. + fill: An optional value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the episode's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, an episode with + extra_model_outputs["b"] = [10, 11, 12, 13, 14] and lookback buffer + size of 2 (meaning `10` and `11` are part of the lookback buffer) will + respond to + `get_extra_model_outputs("b", slice(-7, -2), fill=0.0)` with + `[0.0, 0.0, 10, 11, 12]`. + TODO (sven): This would require a space being provided. Maybe we can + automatically infer the space from existing data? + + Examples: + + .. testcode:: + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + + episode = SingleAgentEpisode( + extra_model_outputs={"mo": [1, 2, 3]}, + len_lookback_buffer=0, # no lookback; all data is actually "in" episode + # The following is needed, but not relevant for this demo. + observations=[0, 1, 2, 3], actions=[1, 2, 3], rewards=[1, 2, 3], + ) + + # Plain usage (`indices` arg only). + episode.get_extra_model_outputs("mo", -1) # 3 + episode.get_extra_model_outputs("mo", 1) # 0 + episode.get_extra_model_outputs("mo", [0, 2]) # [1, 3] + episode.get_extra_model_outputs("mo", [-1, 0]) # [3, 1] + episode.get_extra_model_outputs("mo", slice(None, 2)) # [1, 2] + episode.get_extra_model_outputs("mo", slice(-2, None)) # [2, 3] + # Using `fill=...` (requesting slices beyond the boundaries). + # TODO (sven): This would require a space being provided. Maybe we can + # automatically infer the space from existing data? + # episode.get_extra_model_outputs("mo", slice(-5, -2), fill=0) # [0, 0, 1] + # episode.get_extra_model_outputs("mo", slice(2, 5), fill=-1) # [3, -1, -1] + + Returns: + The collected extra_model_outputs[`key`]. + As a 0-axis batch, if there are several `indices` or a list of exactly one + index provided OR `indices` is a slice object. + As single item (B=0 -> no additional 0-axis) if `indices` is a single int. + """ + value = self.extra_model_outputs[key] + # The expected case is: `value` is a `InfiniteLookbackBuffer`. + if isinstance(value, InfiniteLookbackBuffer): + return value.get( + indices=indices, + neg_index_as_lookback=neg_index_as_lookback, + fill=fill, + ) + # TODO (sven): This does not seem to be solid yet. Users should NOT be able + # to just write directly into our buffers. Instead, use: + # `self.set_extra_model_outputs(key, new_data, at_indices=...)` and if key + # is not known, add a new buffer to the `extra_model_outputs` dict. + assert False + # It might be that the user has added new key/value pairs in their custom + # postprocessing/connector logic. The values are then most likely numpy + # arrays. We convert them automatically to buffers and get the requested + # indices (with the given options) from there. + return InfiniteLookbackBuffer(value).get( + indices, fill=fill, neg_index_as_lookback=neg_index_as_lookback + ) + + def set_observations( + self, + *, + new_data, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of this Episode's observations with the provided data. + + Note that an episode's observation data cannot be written to directly as it is + managed by a `InfiniteLookbackBuffer` object. Normally, individual, current + observations are added to the episode either by calling `self.add_env_step` or + more directly (and manually) via `self.observations.append|extend()`. + However, for certain postprocessing steps, the entirety (or a slice) of an + episode's observations might have to be rewritten, which is when + `self.set_observations()` should be used. + + Args: + new_data: The new observation data to overwrite existing data with. + This may be a list of individual observation(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be (possibly complex) struct matching the + observation space and with a batch size of its leafs exactly the size + of the to-be-overwritten slice or segment (provided by `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single observation). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + observations = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_observations(individual_observation, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + observations buffer with the provided "individual_observation". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + self.observations.set( + new_data=new_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def set_actions( + self, + *, + new_data, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of this Episode's actions with the provided data. + + Note that an episode's action data cannot be written to directly as it is + managed by a `InfiniteLookbackBuffer` object. Normally, individual, current + actions are added to the episode either by calling `self.add_env_step` or + more directly (and manually) via `self.actions.append|extend()`. + However, for certain postprocessing steps, the entirety (or a slice) of an + episode's actions might have to be rewritten, which is when + `self.set_actions()` should be used. + + Args: + new_data: The new action data to overwrite existing data with. + This may be a list of individual action(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be (possibly complex) struct matching the + action space and with a batch size of its leafs exactly the size + of the to-be-overwritten slice or segment (provided by `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single action). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + actions = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_actions(individual_action, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + actions buffer with the provided "individual_action". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + self.actions.set( + new_data=new_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def set_rewards( + self, + *, + new_data, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of this Episode's rewards with the provided data. + + Note that an episode's reward data cannot be written to directly as it is + managed by a `InfiniteLookbackBuffer` object. Normally, individual, current + rewards are added to the episode either by calling `self.add_env_step` or + more directly (and manually) via `self.rewards.append|extend()`. + However, for certain postprocessing steps, the entirety (or a slice) of an + episode's rewards might have to be rewritten, which is when + `self.set_rewards()` should be used. + + Args: + new_data: The new reward data to overwrite existing data with. + This may be a list of individual reward(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be a np.ndarray with a length exactly + the size of the to-be-overwritten slice or segment (provided by + `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single reward). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + rewards = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_rewards(individual_reward, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + rewards buffer with the provided "individual_reward". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + self.rewards.set( + new_data=new_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def set_extra_model_outputs( + self, + *, + key, + new_data, + at_indices: Optional[Union[int, List[int], slice]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of this Episode's extra model outputs with `new_data`. + + Note that an episode's `extra_model_outputs` data cannot be written to directly + as it is managed by a `InfiniteLookbackBuffer` object. Normally, individual, + current `extra_model_output` values are added to the episode either by calling + `self.add_env_step` or more directly (and manually) via + `self.extra_model_outputs[key].append|extend()`. However, for certain + postprocessing steps, the entirety (or a slice) of an episode's + `extra_model_outputs` might have to be rewritten or a new key (a new type of + `extra_model_outputs`) must be inserted, which is when + `self.set_extra_model_outputs()` should be used. + + Args: + key: The `key` within `self.extra_model_outputs` to override data on or + to insert as a new key into `self.extra_model_outputs`. + new_data: The new data to overwrite existing data with. + This may be a list of individual reward(s) in case this episode + is still not numpy'ized yet. In case this episode has already been + numpy'ized, this should be a np.ndarray with a length exactly + the size of the to-be-overwritten slice or segment (provided by + `at_indices`). + at_indices: A single int is interpreted as one index, which to overwrite + with `new_data` (which is expected to be a single reward). + A list of ints is interpreted as a list of indices, all of which to + overwrite with `new_data` (which is expected to be of the same size + as `len(at_indices)`). + A slice object is interpreted as a range of indices to be overwritten + with `new_data` (which is expected to be of the same size as the + provided slice). + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, an episode with + rewards = [4, 5, 6, 7, 8, 9], where [4, 5, 6] is the + lookback buffer range (ts=0 item is 7), will handle a call to + `set_rewards(individual_reward, -1, + neg_index_as_lookback=True)` by overwriting the value of 6 in our + rewards buffer with the provided "individual_reward". + + Raises: + IndexError: If the provided `at_indices` do not match the size of + `new_data`. + """ + # Record already exists -> Set existing record's data to new values. + assert key in self.extra_model_outputs + self.extra_model_outputs[key].set( + new_data=new_data, + at_indices=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def slice( + self, + slice_: slice, + *, + len_lookback_buffer: Optional[int] = None, + ) -> "SingleAgentEpisode": + """Returns a slice of this episode with the given slice object. + + For example, if `self` contains o0 (the reset observation), o1, o2, o3, and o4 + and the actions a1, a2, a3, and a4 (len of `self` is 4), then a call to + `self.slice(slice(1, 3))` would return a new SingleAgentEpisode with + observations o1, o2, and o3, and actions a2 and a3. Note here that there is + always one observation more in an episode than there are actions (and rewards + and extra model outputs) due to the initial observation received after an env + reset. + + .. testcode:: + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.test_utils import check + + # Generate a simple multi-agent episode. + observations = [0, 1, 2, 3, 4, 5] + actions = [1, 2, 3, 4, 5] + rewards = [0.1, 0.2, 0.3, 0.4, 0.5] + episode = SingleAgentEpisode( + observations=observations, + actions=actions, + rewards=rewards, + len_lookback_buffer=0, # all given data is part of the episode + ) + slice_1 = episode[:1] + check(slice_1.observations, [0, 1]) + check(slice_1.actions, [1]) + check(slice_1.rewards, [0.1]) + + slice_2 = episode[-2:] + check(slice_2.observations, [3, 4, 5]) + check(slice_2.actions, [4, 5]) + check(slice_2.rewards, [0.4, 0.5]) + + Args: + slice_: The slice object to use for slicing. This should exclude the + lookback buffer, which will be prepended automatically to the returned + slice. + len_lookback_buffer: If not None, forces the returned slice to try to have + this number of timesteps in its lookback buffer (if available). If None + (default), tries to make the returned slice's lookback as large as the + current lookback buffer of this episode (`self`). + + Returns: + The new SingleAgentEpisode representing the requested slice. + """ + # Translate `slice_` into one that only contains 0-or-positive ints and will + # NOT contain any None. + start = slice_.start + stop = slice_.stop + + # Start is None -> 0. + if start is None: + start = 0 + # Start is negative -> Interpret index as counting "from end". + elif start < 0: + start = len(self) + start + + # Stop is None -> Set stop to our len (one ts past last valid index). + if stop is None: + stop = len(self) + # Stop is negative -> Interpret index as counting "from end". + elif stop < 0: + stop = len(self) + stop + + step = slice_.step if slice_.step is not None else 1 + + # Figure out, whether slicing stops at the very end of this episode to know + # whether `self.is_terminated/is_truncated` should be kept as-is. + keep_done = stop == len(self) + # Provide correct timestep- and pre-buffer information. + t_started = self.t_started + start + + _lb = ( + len_lookback_buffer + if len_lookback_buffer is not None + else self.observations.lookback + ) + if ( + start >= 0 + and start - _lb < 0 + and self.observations.lookback < (_lb - start) + ): + _lb = self.observations.lookback + start + observations = InfiniteLookbackBuffer( + data=self.get_observations( + slice(start - _lb, stop + 1, step), + neg_index_as_lookback=True, + ), + lookback=_lb, + space=self.observation_space, + ) + + _lb = ( + len_lookback_buffer + if len_lookback_buffer is not None + else self.infos.lookback + ) + if start >= 0 and start - _lb < 0 and self.infos.lookback < (_lb - start): + _lb = self.infos.lookback + start + infos = InfiniteLookbackBuffer( + data=self.get_infos( + slice(start - _lb, stop + 1, step), + neg_index_as_lookback=True, + ), + lookback=_lb, + ) + + _lb = ( + len_lookback_buffer + if len_lookback_buffer is not None + else self.actions.lookback + ) + if start >= 0 and start - _lb < 0 and self.actions.lookback < (_lb - start): + _lb = self.actions.lookback + start + actions = InfiniteLookbackBuffer( + data=self.get_actions( + slice(start - _lb, stop, step), + neg_index_as_lookback=True, + ), + lookback=_lb, + space=self.action_space, + ) + + _lb = ( + len_lookback_buffer + if len_lookback_buffer is not None + else self.rewards.lookback + ) + if start >= 0 and start - _lb < 0 and self.rewards.lookback < (_lb - start): + _lb = self.rewards.lookback + start + rewards = InfiniteLookbackBuffer( + data=self.get_rewards( + slice(start - _lb, stop, step), + neg_index_as_lookback=True, + ), + lookback=_lb, + ) + + extra_model_outputs = {} + for k, v in self.extra_model_outputs.items(): + _lb = len_lookback_buffer if len_lookback_buffer is not None else v.lookback + if start >= 0 and start - _lb < 0 and v.lookback < (_lb - start): + _lb = v.lookback + start + extra_model_outputs[k] = InfiniteLookbackBuffer( + data=self.get_extra_model_outputs( + key=k, + indices=slice(start - _lb, stop, step), + neg_index_as_lookback=True, + ), + lookback=_lb, + ) + + return SingleAgentEpisode( + id_=self.id_, + # In the following, offset `start`s automatically by lookbacks. + observations=observations, + observation_space=self.observation_space, + infos=infos, + actions=actions, + action_space=self.action_space, + rewards=rewards, + extra_model_outputs=extra_model_outputs, + terminated=(self.is_terminated if keep_done else False), + truncated=(self.is_truncated if keep_done else False), + t_started=t_started, + ) + + def get_data_dict(self): + """Converts a SingleAgentEpisode into a data dict mapping str keys to data. + + The keys used are: + Columns.EPS_ID, T, OBS, INFOS, ACTIONS, REWARDS, TERMINATEDS, TRUNCATEDS, + and those in `self.extra_model_outputs`. + + Returns: + A data dict mapping str keys to data records. + """ + t = list(range(self.t_started, self.t)) + terminateds = [False] * (len(self) - 1) + [self.is_terminated] + truncateds = [False] * (len(self) - 1) + [self.is_truncated] + eps_id = [self.id_] * len(self) + + if self.is_numpy: + t = np.array(t) + terminateds = np.array(terminateds) + truncateds = np.array(truncateds) + eps_id = np.array(eps_id) + + return dict( + { + # Trivial 1D data (compiled above). + Columns.TERMINATEDS: terminateds, + Columns.TRUNCATEDS: truncateds, + Columns.T: t, + Columns.EPS_ID: eps_id, + # Retrieve obs, infos, actions, rewards using our get_... APIs, + # which return all relevant timesteps (excluding the lookback + # buffer!). Slice off last obs and infos to have the same number + # of them as we have actions and rewards. + Columns.OBS: self.get_observations(slice(None, -1)), + Columns.INFOS: self.get_infos(slice(None, -1)), + Columns.ACTIONS: self.get_actions(), + Columns.REWARDS: self.get_rewards(), + }, + # All `extra_model_outs`: Same as obs: Use get_... API. + **{ + k: self.get_extra_model_outputs(k) + for k in self.extra_model_outputs.keys() + }, + ) + + def get_sample_batch(self) -> SampleBatch: + """Converts this `SingleAgentEpisode` into a `SampleBatch`. + + Returns: + A SampleBatch containing all of this episode's data. + """ + return SampleBatch(self.get_data_dict()) + + def get_return(self) -> float: + """Calculates an episode's return, excluding the lookback buffer's rewards. + + The return is computed by a simple sum, neglecting the discount factor. + Note that if `self` is a continuation chunk (resulting from a call to + `self.cut()`), the previous chunk's rewards are NOT counted and thus NOT + part of the returned reward sum. + + Returns: + The sum of rewards collected during this episode, excluding possible data + inside the lookback buffer and excluding possible data in a predecessor + chunk. + """ + return sum(self.get_rewards()) + + def get_duration_s(self) -> float: + """Returns the duration of this Episode (chunk) in seconds.""" + if self._last_step_time is None: + return 0.0 + return self._last_step_time - self._start_time + + def env_steps(self) -> int: + """Returns the number of environment steps. + + Note, this episode instance could be a chunk of an actual episode. + + Returns: + An integer that counts the number of environment steps this episode instance + has seen. + """ + return len(self) + + def agent_steps(self) -> int: + """Returns the number of agent steps. + + Note, these are identical to the environment steps for a single-agent episode. + + Returns: + An integer counting the number of agent steps executed during the time this + episode instance records. + """ + return self.env_steps() + + def get_state(self) -> Dict[str, Any]: + """Returns the pickable state of an episode. + + The data in the episode is stored into a dictionary. Note that episodes + can also be generated from states (see `SingleAgentEpisode.from_state()`). + + Returns: + A dict containing all the data from the episode. + """ + infos = self.infos.get_state() + infos["data"] = np.array([info if info else None for info in infos["data"]]) + return { + "id_": self.id_, + "agent_id": self.agent_id, + "module_id": self.module_id, + "multi_agent_episode_id": self.multi_agent_episode_id, + # Note, all data is stored in `InfiniteLookbackBuffer`s. + "observations": self.observations.get_state(), + "actions": self.actions.get_state(), + "rewards": self.rewards.get_state(), + "infos": self.infos.get_state(), + "extra_model_outputs": { + k: v.get_state() if v else v + for k, v in self.extra_model_outputs.items() + } + if len(self.extra_model_outputs) > 0 + else None, + "is_terminated": self.is_terminated, + "is_truncated": self.is_truncated, + "t_started": self.t_started, + "t": self.t, + "_observation_space": gym_space_to_dict(self._observation_space) + if self._observation_space + else None, + "_action_space": gym_space_to_dict(self._action_space) + if self._action_space + else None, + "_start_time": self._start_time, + "_last_step_time": self._last_step_time, + "custom_data": self.custom_data, + } + + @staticmethod + def from_state(state: Dict[str, Any]) -> "SingleAgentEpisode": + """Creates a new `SingleAgentEpisode` instance from a state dict. + + Args: + state: The state dict, as returned by `self.get_state()`. + + Returns: + A new `SingleAgentEpisode` instance with the data from the state dict. + """ + # Create an empy episode instance. + episode = SingleAgentEpisode(id_=state["id_"]) + # Load all the data from the state dict into the episode. + episode.agent_id = state["agent_id"] + episode.module_id = state["module_id"] + episode.multi_agent_episode_id = state["multi_agent_episode_id"] + # Convert data back to `InfiniteLookbackBuffer`s. + episode.observations = InfiniteLookbackBuffer.from_state(state["observations"]) + episode.actions = InfiniteLookbackBuffer.from_state(state["actions"]) + episode.rewards = InfiniteLookbackBuffer.from_state(state["rewards"]) + episode.infos = InfiniteLookbackBuffer.from_state(state["infos"]) + episode.extra_model_outputs = ( + defaultdict( + functools.partial( + InfiniteLookbackBuffer, lookback=episode.observations.lookback + ), + { + k: InfiniteLookbackBuffer.from_state(v) + for k, v in state["extra_model_outputs"].items() + }, + ) + if state["extra_model_outputs"] + else defaultdict( + functools.partial( + InfiniteLookbackBuffer, lookback=episode.observations.lookback + ), + ) + ) + episode.is_terminated = state["is_terminated"] + episode.is_truncated = state["is_truncated"] + episode.t_started = state["t_started"] + episode.t = state["t"] + # We need to convert the spaces to dictionaries for serialization. + episode._observation_space = ( + gym_space_from_dict(state["_observation_space"]) + if state["_observation_space"] + else None + ) + episode._action_space = ( + gym_space_from_dict(state["_action_space"]) + if state["_action_space"] + else None + ) + episode._start_time = state["_start_time"] + episode._last_step_time = state["_last_step_time"] + episode._custom_data = state.get("custom_data", {}) + # Validate the episode. + episode.validate() + + return episode + + @property + def observation_space(self): + return self._observation_space + + @observation_space.setter + def observation_space(self, value): + self._observation_space = self.observations.space = value + + @property + def action_space(self): + return self._action_space + + @action_space.setter + def action_space(self, value): + self._action_space = self.actions.space = value + + def __len__(self) -> int: + """Returning the length of an episode. + + The length of an episode is defined by the length of its data, excluding + the lookback buffer data. The length is the number of timesteps an agent has + stepped through an environment thus far. + + The length is 0 in case of an episode whose env has NOT been reset yet, but + also 0 right after the `env.reset()` data has been added via + `self.add_env_reset()`. Only after the first call to `env.step()` (and + `self.add_env_step()`, the length will be 1. + + Returns: + An integer, defining the length of an episode. + """ + return self.t - self.t_started + + def __repr__(self): + return ( + f"SAEps(len={len(self)} done={self.is_done} " + f"R={self.get_return()} id_={self.id_})" + ) + + def __getitem__(self, item: slice) -> "SingleAgentEpisode": + """Enable squared bracket indexing- and slicing syntax, e.g. episode[-4:].""" + if isinstance(item, slice): + return self.slice(slice_=item) + else: + raise NotImplementedError( + f"SingleAgentEpisode does not support getting item '{item}'! " + "Only slice objects allowed with the syntax: `episode[a:b]`." + ) + + @Deprecated(new="SingleAgentEpisode.custom_data[some-key] = ...", error=True) + def add_temporary_timestep_data(self): + pass + + @Deprecated(new="SingleAgentEpisode.custom_data[some-key]", error=True) + def get_temporary_timestep_data(self): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/env/tcp_client_inference_env_runner.py b/lib/python3.12/site-packages/ray/rllib/env/tcp_client_inference_env_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..09f8f4a2e715ce46116d94b792106bf31797aa4a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/tcp_client_inference_env_runner.py @@ -0,0 +1,6 @@ +from ray.rllib.env.external.env_runner_server_for_external_inference import ( + EnvRunnerServerForExternalInference, +) + +# @Deprecated +TcpClientInferenceEnvRunner = EnvRunnerServerForExternalInference diff --git a/lib/python3.12/site-packages/ray/rllib/env/utils/__init__.py b/lib/python3.12/site-packages/ray/rllib/env/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f324c3c3bf4f50396bd117b756606216d3e7c5a8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/utils/__init__.py @@ -0,0 +1,114 @@ +import logging +from typing import Type, Union + +import gymnasium as gym + +from ray.rllib.env.env_context import EnvContext +from ray.rllib.utils.error import ( + ERR_MSG_INVALID_ENV_DESCRIPTOR, + EnvError, +) +from ray.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + + +@PublicAPI +def try_import_pyspiel(error: bool = False): + """Tries importing pyspiel and returns the module (or None). + + Args: + error: Whether to raise an error if pyspiel cannot be imported. + + Returns: + The pyspiel module. + + Raises: + ImportError: If error=True and pyspiel is not installed. + """ + try: + import pyspiel + + return pyspiel + except ImportError: + if error: + raise ImportError( + "Could not import pyspiel! Pyspiel is not a dependency of RLlib " + "and RLlib requires you to install pyspiel separately: " + "`pip install open_spiel`." + ) + return None + + +@PublicAPI +def try_import_open_spiel(error: bool = False): + """Tries importing open_spiel and returns the module (or None). + + Args: + error: Whether to raise an error if open_spiel cannot be imported. + + Returns: + The open_spiel module. + + Raises: + ImportError: If error=True and open_spiel is not installed. + """ + try: + import open_spiel + + return open_spiel + except ImportError: + if error: + raise ImportError( + "Could not import open_spiel! open_spiel is not a dependency of RLlib " + "and RLlib requires you to install open_spiel separately: " + "`pip install open_spiel`." + ) + return None + + +def _gym_env_creator( + env_context: EnvContext, + env_descriptor: Union[str, Type[gym.Env]], +) -> gym.Env: + """Tries to create a gym env given an EnvContext object and descriptor. + + Note: This function tries to construct the env from a string descriptor + only using possibly installed RL env packages (such as gymnasium). + These packages are no installation requirements for RLlib. In case + you would like to support more such env packages, add the necessary imports + and construction logic below. + + Args: + env_context: The env context object to configure the env. + Note that this is a config dict, plus the properties: + `worker_index`, `vector_index`, and `remote`. + env_descriptor: The env descriptor as a gym-registered string, e.g. + "CartPole-v1", "ale_py:ALE/Breakout-v5". + Alternatively, the gym.Env subclass to use. + + Returns: + The actual gym environment object. + + Raises: + gym.error.Error: If the env cannot be constructed. + """ + # If env descriptor is a str, starting with "ale_py:ALE/", for now, register all ALE + # envs from ale_py. + if isinstance(env_descriptor, str) and env_descriptor.startswith("ale_py:ALE/"): + import ale_py + + gym.register_envs(ale_py) + + # Try creating a gym env. If this fails we can output a + # decent error message. + try: + # If class provided, call constructor directly. + if isinstance(env_descriptor, type): + env = env_descriptor(env_context) + else: + env = gym.make(env_descriptor, **env_context) + except gym.error.Error: + raise EnvError(ERR_MSG_INVALID_ENV_DESCRIPTOR.format(env_descriptor)) + + return env diff --git a/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24b3770f953ccb3dd848d4a934ae1974af72d42d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/external_env_protocol.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/external_env_protocol.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..74117227af962ce7efcc9417963308d9896c9068 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/external_env_protocol.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/infinite_lookback_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/infinite_lookback_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58f5c9385b4408a00c223e3562e3870d818fddca Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/utils/__pycache__/infinite_lookback_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/utils/external_env_protocol.py b/lib/python3.12/site-packages/ray/rllib/env/utils/external_env_protocol.py new file mode 100644 index 0000000000000000000000000000000000000000..3356a87da30ac27923962aa636db0ae801a9c009 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/utils/external_env_protocol.py @@ -0,0 +1,8 @@ +from ray.rllib.env.external.rllink import RLlink # noqa +from ray._common.deprecation import deprecation_warning + +deprecation_warning( + old="ray.rllib.env.utils.external_env_protocol", + new="ray.rllib.env.external.rllink", + error=False, +) diff --git a/lib/python3.12/site-packages/ray/rllib/env/utils/infinite_lookback_buffer.py b/lib/python3.12/site-packages/ray/rllib/env/utils/infinite_lookback_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..bbe468f9a6d3475f50ee75b5455e3cbb677c68e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/utils/infinite_lookback_buffer.py @@ -0,0 +1,715 @@ +from typing import Any, Dict, List, Optional, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.utils.env_checker import data_equivalence + +from ray.rllib.utils.numpy import LARGE_INTEGER, one_hot, one_hot_multidiscrete +from ray.rllib.utils.serialization import gym_space_from_dict, gym_space_to_dict +from ray.rllib.utils.spaces.space_utils import ( + batch, + get_base_struct_from_space, + get_dummy_batch_for_space, +) +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class InfiniteLookbackBuffer: + def __init__( + self, + data: Optional[Union[List, np.ndarray]] = None, + lookback: int = 0, + space: Optional[gym.Space] = None, + ): + self.data = data if data is not None else [] + self.lookback = min(lookback, len(self.data)) + self.finalized = not isinstance(self.data, list) + + self.space = space + + def __eq__( + self, + other: "InfiniteLookbackBuffer", + ) -> bool: + """Compares two `InfiniteLookbackBuffers. + + Args: + other: Another object. If another `LookbackBuffer` instance all + their attributes are compared. + + Returns: + `True`, if `other` is an `InfiniteLookbackBuffer` instance and all + attributes are identical. Otherwise, returns `False`. + """ + return ( + isinstance(other, InfiniteLookbackBuffer) + # Todo (mark): Replace `data_equivalence` with ray / rllib implementation similar to `check` without asserts + and data_equivalence(self.data, other.data) + and self.lookback == other.lookback + and self.finalized == other.finalized + and self.space_struct == other.space_struct + and self.space == other.space + ) + + @property + def space(self): + return self._space + + @space.setter + def space(self, value): + self._space = value + self._space_struct = get_base_struct_from_space(value) + + @property + def space_struct(self): + return self._space_struct + + def get_state(self) -> Dict[str, Any]: + """Returns the pickable state of a buffer. + + The data in the buffer is stored into a dictionary. Note that + buffers can also be generated from pickable states (see + `InfiniteLookbackBuffer.from_state`) + + Returns: + A dict containing all the data and metadata from the buffer. + """ + return { + "data": self.data, + "lookback": self.lookback, + "finalized": self.finalized, + "space": gym_space_to_dict(self.space) if self.space else None, + } + + @staticmethod + def from_state(state: Dict[str, Any]) -> "InfiniteLookbackBuffer": + """Creates a new `InfiniteLookbackBuffer` from a state dict. + + Args: + state: The state dict, as returned by `self.get_state`. + + Returns: + A new `InfiniteLookbackBuffer` instance with the data and metadata + from the state dict. + """ + buffer = InfiniteLookbackBuffer() + buffer.lookback = state["lookback"] + buffer.finalized = state["finalized"] + buffer.space = gym_space_from_dict(state["space"]) if state["space"] else None + # space_struct is set when space is assigned + buffer.data = state["data"] + + return buffer + + def append(self, item) -> None: + """Appends the given item to the end of this buffer.""" + if self.finalized: + self.data = tree.map_structure( + lambda d, i: np.concatenate([d, [i]], axis=0), self.data, item + ) + else: + self.data.append(item) + + def extend(self, items) -> None: + """Appends all items in `items` to the end of this buffer.""" + if self.finalized: + # TODO (sven): When extending with a list of structs, we should + # probably rather do: `tree.map_structure(..., self.data, + # tree.map_structure(lambda *s: np.array(*s), *items)`)?? + self.data = tree.map_structure( + lambda d, i: np.concatenate([d, i], axis=0), + self.data, + # Note, we could have dictionaries here. + np.array(items) if isinstance(items, list) else items, + ) + else: + for item in items: + self.append(item) + + def concat(self, other: "InfiniteLookbackBuffer") -> None: + """Concatenates the data of `other` (w/o its lookback) to `self`. + + Args: + other: The other InfiniteLookbackBuffer to be concatenated to self. + """ + self.data.extend(other.get()) + + def pop(self, index: int = -1) -> None: + """Removes the item at `index` from this buffer, but does NOT return it. + + Args: + index: The index to pop out of this buffer (w/o returning it from this + method). + """ + if self.finalized: + self.data = tree.map_structure( + lambda s: np.delete(s, index, axis=0), self.data + ) + else: + self.data.pop(index) + + def finalize(self) -> None: + """Finalizes this buffer by converting internal data lists into numpy arrays. + + Thereby, if the individual items in the list are nested structures, the + resulting buffer content will be a nested struct of np.ndarrays (leafs). + """ + if not self.finalized: + self.data = batch(self.data) + self.finalized = True + + def get( + self, + indices: Optional[Union[int, slice, List[int]]] = None, + *, + neg_index_as_lookback: bool = False, + fill: Optional[Any] = None, + one_hot_discrete: bool = False, + _ignore_last_ts: bool = False, + _add_last_ts_value: Optional[Any] = None, + ) -> Any: + """Returns data, based on the given args, from this buffer. + + Args: + indices: A single int is interpreted as an index, from which to return the + individual data stored at this index. + A list of ints is interpreted as a list of indices from which to gather + individual data in a batch of size len(indices). + A slice object is interpreted as a range of data to be returned. + Thereby, negative indices by default are interpreted as "before the end" + unless the `neg_index_as_lookback=True` option is used, in which case + negative indices are interpreted as "before ts=0", meaning going back + into the lookback buffer. + neg_index_as_lookback: If True, negative values in `indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, a buffer with data [4, 5, 6, 7, 8, 9], + where [4, 5, 6] is the lookback buffer range (ts=0 item is 7), will + respond to `get(-1, neg_index_as_lookback=True)` with `6` and to + `get(slice(-2, 1), neg_index_as_lookback=True)` with `[5, 6, 7]`. + fill: An optional float value to use for filling up the returned results at + the boundaries. This filling only happens if the requested index range's + start/stop boundaries exceed the buffer's boundaries (including the + lookback buffer on the left side). This comes in very handy, if users + don't want to worry about reaching such boundaries and want to zero-pad. + For example, a buffer with data [10, 11, 12, 13, 14] and lookback + buffer size of 2 (meaning `10` and `11` are part of the lookback buffer) + will respond to `get(slice(-7, -2), fill=0.0)` + with `[0.0, 0.0, 10, 11, 12]`. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) space + that are Discrete or MultiDiscrete. Note that if `fill=0` and the + requested `indices` are out of the range of our data, the returned + one-hot vectors will actually be zero-hot (all slots zero). + _ignore_last_ts: Whether to ignore the last record in our internal + `self.data` when getting the provided indices. + _add_last_ts_value: Whether to add the value of this arg to the end of + the internal `self.data` buffer (just for the duration of this get + operation, not permanently). + """ + if indices is None: + data = self._get_all_data( + one_hot_discrete=one_hot_discrete, + _ignore_last_ts=_ignore_last_ts, + ) + elif isinstance(indices, slice): + data = self._get_slice( + indices, + fill=fill, + neg_index_as_lookback=neg_index_as_lookback, + one_hot_discrete=one_hot_discrete, + _ignore_last_ts=_ignore_last_ts, + _add_last_ts_value=_add_last_ts_value, + ) + elif isinstance(indices, list): + data = [ + self._get_int_index( + idx, + fill=fill, + neg_index_as_lookback=neg_index_as_lookback, + one_hot_discrete=one_hot_discrete, + _ignore_last_ts=_ignore_last_ts, + _add_last_ts_value=_add_last_ts_value, + ) + for idx in indices + ] + if self.finalized: + data = batch(data) + else: + assert isinstance(indices, int) + data = self._get_int_index( + indices, + fill=fill, + neg_index_as_lookback=neg_index_as_lookback, + one_hot_discrete=one_hot_discrete, + _ignore_last_ts=_ignore_last_ts, + _add_last_ts_value=_add_last_ts_value, + ) + + return data + + def __add__( + self, other: Union[List, "InfiniteLookbackBuffer", int, float, complex] + ) -> "InfiniteLookbackBuffer": + """Adds another InfiniteLookbackBuffer object or list to the end of this one. + + Args: + other: Another `InfiniteLookbackBuffer` or a `list` or a number. + If a `InfiniteLookbackBuffer` its data (w/o its lookback buffer) gets + concatenated to self's data. If a `list`, we concat it to self's data. + If a number, we add this number to each element of self (if possible). + + Returns: + A new `InfiniteLookbackBuffer` instance `self.data` containing + concatenated data from `self` and `other` (or adding `other` to each element + in self's data). + """ + + if self.finalized: + raise RuntimeError(f"Cannot `add` to a finalized {type(self).__name__}.") + else: + # If `other` is an int, simply add it to all our values (if possible) and + # use the result as the underlying data for the returned buffer. + if isinstance(other, (int, float, complex)): + data = [ + (d + other) if isinstance(d, (int, float, complex)) else d + for d in self.data + ] + # If `other` is a InfiniteLookbackBuffer itself, do NOT include its + # lookback buffer anymore. We assume that `other`'s lookback buffer i + # already at the end of `self`. + elif isinstance(other, InfiniteLookbackBuffer): + data = self.data + other.data[other.lookback :] + # `other` is a list, simply concat the two lists and use the result as + # the underlying data for the returned buffer. + else: + data = self.data + other + + return InfiniteLookbackBuffer( + data=data, + lookback=self.lookback, + space=self.space, + ) + + def __getitem__(self, item): + """Support squared bracket syntax, e.g. buffer[:5].""" + return self.get(item) + + def __setitem__(self, key, value): + self.set(new_data=value, at_indices=key) + + def set( + self, + new_data, + *, + at_indices: Optional[Union[int, slice, List[int]]] = None, + neg_index_as_lookback: bool = False, + ) -> None: + """Overwrites all or some of the data in this buffer with the provided data. + + Args: + new_data: The new data to overwrite existing records with. + at_indices: A single int is interpreted as an index, at which to overwrite + the individual record stored at this index with `new_data`. + A list of ints is interpreted as a list of indices, which to overwrite + with `new_data`, which must be a batch of size `len(at_indices)`. + A slice object is interpreted as a range, which to overwrite with + `new_data`. Thereby, negative indices by default are interpreted as + "before the end" unless the `neg_index_as_lookback=True` option is + used, in which case negative indices are interpreted as + "before ts=0", meaning going back into the lookback buffer. + neg_index_as_lookback: If True, negative values in `at_indices` are + interpreted as "before ts=0", meaning going back into the lookback + buffer. For example, a buffer with data [4, 5, 6, 7, 8, 9], + where [4, 5, 6] is the lookback buffer range (ts=0 item is 7), will + handle a call `set(99, at_indices=-1, neg_index_as_lookback=True)` + with `6` being replaced by 99 and to `set([98, 99, 100], + at_indices=slice(-2, 1), neg_index_as_lookback=True)` with + `[5, 6, 7]` being replaced by `[98, 99, 100]`. + """ + # `at_indices` is None -> Override all our data (excluding the lookback buffer). + if at_indices is None: + self._set_all_data(new_data) + + elif isinstance(at_indices, slice): + self._set_slice( + new_data, + slice_=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + elif isinstance(at_indices, list): + for i, idx in enumerate(at_indices): + self._set_int_index( + new_data[i], + idx=idx, + neg_index_as_lookback=neg_index_as_lookback, + ) + else: + assert isinstance(at_indices, int) + self._set_int_index( + new_data, + idx=at_indices, + neg_index_as_lookback=neg_index_as_lookback, + ) + + def __len__(self): + """Return the length of our data, excluding the lookback buffer.""" + len_ = self.len_incl_lookback() + # Only count the data after the lookback. + return max(len_ - self.lookback, 0) + + def len_incl_lookback(self): + if self.finalized: + return len(tree.flatten(self.data)[0]) + else: + return len(self.data) + + def __repr__(self): + return ( + f"{type(self).__name__}({self.data[:self.lookback]} <- " + f"lookback({self.lookback}) | {self.data[self.lookback:]})" + ) + + def _get_all_data(self, one_hot_discrete=False, _ignore_last_ts=False): + data = self[: (None if not _ignore_last_ts else -1)] + if one_hot_discrete: + data = self._one_hot(data, space_struct=self.space_struct) + return data + + def _set_all_data(self, new_data): + self._set_slice(new_data, slice(0, None)) + + def _get_slice( + self, + slice_, + fill=None, + neg_index_as_lookback=False, + one_hot_discrete=False, + _ignore_last_ts=False, + _add_last_ts_value=None, + ): + data_to_use = self.data + if _ignore_last_ts: + if self.finalized: + data_to_use = tree.map_structure(lambda s: s[:-1], self.data) + else: + data_to_use = self.data[:-1] + if _add_last_ts_value is not None: + if self.finalized: + data_to_use = tree.map_structure( + lambda s, t: np.append(s, t), + data_to_use.copy(), + _add_last_ts_value, + ) + else: + data_to_use = np.append(data_to_use.copy(), _add_last_ts_value) + + slice_, slice_len, fill_left_count, fill_right_count = self._interpret_slice( + slice_, + neg_index_as_lookback, + len_self_plus_lookback=( + self.len_incl_lookback() + + int(_add_last_ts_value is not None) + - int(_ignore_last_ts) + ), + ) + + # Perform the actual slice. + data_slice = None + if slice_len > 0: + if self.finalized: + data_slice = tree.map_structure(lambda s: s[slice_], data_to_use) + else: + data_slice = data_to_use[slice_] + + if one_hot_discrete: + data_slice = self._one_hot(data_slice, space_struct=self.space_struct) + + # Data is shorter than the range requested -> Fill the rest with `fill` data. + if fill is not None and (fill_right_count > 0 or fill_left_count > 0): + if self.finalized: + if fill_left_count: + if self.space is None: + fill_batch = np.array([fill] * fill_left_count) + else: + fill_batch = get_dummy_batch_for_space( + self.space, + fill_value=fill, + batch_size=fill_left_count, + one_hot_discrete=one_hot_discrete, + ) + if data_slice is not None: + data_slice = tree.map_structure( + lambda s0, s: np.concatenate([s0, s]), + fill_batch, + data_slice, + ) + else: + data_slice = fill_batch + if fill_right_count: + if self.space is None: + fill_batch = np.array([fill] * fill_right_count) + else: + fill_batch = get_dummy_batch_for_space( + self.space, + fill_value=fill, + batch_size=fill_right_count, + one_hot_discrete=one_hot_discrete, + ) + if data_slice is not None: + data_slice = tree.map_structure( + lambda s0, s: np.concatenate([s, s0]), + fill_batch, + data_slice, + ) + else: + data_slice = fill_batch + + else: + if self.space is None: + fill_batch = [fill] + else: + fill_batch = [ + get_dummy_batch_for_space( + self.space, + fill_value=fill, + batch_size=0, + one_hot_discrete=one_hot_discrete, + ) + ] + data_slice = ( + fill_batch * fill_left_count + + (data_slice if data_slice is not None else []) + + fill_batch * fill_right_count + ) + + if data_slice is None: + if self.finalized: + return tree.map_structure(lambda s: s[slice_], data_to_use) + else: + return data_to_use[slice_] + return data_slice + + def _set_slice( + self, + new_data, + slice_, + neg_index_as_lookback=False, + ): + slice_, _, _, _ = self._interpret_slice(slice_, neg_index_as_lookback) + + # Check, whether the setting to new_data changes the length of self + # (it shouldn't). If it does, raise an error. + try: + if self.finalized: + + def __set(s, n): + if self.space: + assert self.space.contains(n[0]) + assert len(s[slice_]) == len(n) + s[slice_] = n + + tree.map_structure(__set, self.data, new_data) + else: + assert len(self.data[slice_]) == len(new_data) + self.data[slice_] = new_data + except AssertionError: + raise IndexError( + f"Cannot `set()` value via at_indices={slice_} (option " + f"neg_index_as_lookback={neg_index_as_lookback})! Slice of data " + "does NOT have the same size as `new_data`." + ) + + def _get_int_index( + self, + idx: int, + fill=None, + neg_index_as_lookback=False, + one_hot_discrete=False, + _ignore_last_ts=False, + _add_last_ts_value=None, + ): + data_to_use = self.data + if _ignore_last_ts: + if self.finalized: + data_to_use = tree.map_structure(lambda s: s[:-1], self.data) + else: + data_to_use = self.data[:-1] + if _add_last_ts_value is not None: + if self.finalized: + data_to_use = tree.map_structure( + lambda s, last: np.append(s, last), data_to_use, _add_last_ts_value + ) + else: + data_to_use = data_to_use.copy() + data_to_use.append(_add_last_ts_value) + + # If index >= 0 -> Ignore lookback buffer. + # Otherwise, include lookback buffer. + if idx >= 0 or neg_index_as_lookback: + idx = self.lookback + idx + # Negative indices mean: Go to left into lookback buffer starting from idx=0. + # But if we pass the lookback buffer, the index should be invalid and we will + # have to fill, if required. Invalidate the index by setting it to one larger + # than max. + if neg_index_as_lookback and idx < 0: + idx = len(self) + self.lookback - (_ignore_last_ts is True) + + try: + if self.finalized: + data = tree.map_structure(lambda s: s[idx], data_to_use) + else: + data = data_to_use[idx] + # Out of range index -> If `fill`, use a fill dummy (B=0), if not, error out. + except IndexError as e: + if fill is not None: + if self.space is None: + return fill + return get_dummy_batch_for_space( + self.space, + fill_value=fill, + batch_size=0, + one_hot_discrete=one_hot_discrete, + ) + else: + raise e + + # Convert discrete/multi-discrete components to one-hot vectors, if required. + if one_hot_discrete: + data = self._one_hot(data, self.space_struct) + return data + + def _set_int_index(self, new_data, idx, neg_index_as_lookback): + actual_idx = idx + # If index >= 0 -> Ignore lookback buffer. + # Otherwise, include lookback buffer. + if actual_idx >= 0 or neg_index_as_lookback: + actual_idx = self.lookback + actual_idx + # Negative indices mean: Go to left into lookback buffer starting from idx=0. + # But if we pass the lookback buffer, the index should be invalid and we will + # have to fill, if required. Invalidate the index by setting it to one larger + # than max. + if neg_index_as_lookback and actual_idx < 0: + actual_idx = len(self) + self.lookback + + try: + if self.finalized: + + def __set(s, n): + if self.space: + assert self.space.contains(n), n + s[actual_idx] = n + + tree.map_structure(__set, self.data, new_data) + else: + self.data[actual_idx] = new_data + except IndexError: + raise IndexError( + f"Cannot `set()` value at index {idx} (option " + f"neg_index_as_lookback={neg_index_as_lookback})! Out of range " + f"of buffer data." + ) + + def _interpret_slice( + self, + slice_, + neg_index_as_lookback, + len_self_plus_lookback=None, + ): + if len_self_plus_lookback is None: + len_self_plus_lookback = len(self) + self.lookback + + # Re-interpret slice bounds as absolute positions (>=0) within our + # internal data. + start = slice_.start + stop = slice_.stop + + # Start is None -> Exclude lookback buffer. + if start is None: + start = self.lookback + # Start is negative. + elif start < 0: + # `neg_index_as_lookback=True` -> User wants to index into the lookback + # range. + if neg_index_as_lookback: + start = self.lookback + start + # Interpret index as counting "from end". + else: + start = len_self_plus_lookback + start + # Start is 0 or positive -> timestep right after lookback is interpreted as 0. + else: + start = self.lookback + start + + # Stop is None -> Set stop to very last index + 1 of our internal data. + if stop is None: + stop = len_self_plus_lookback + # Stop is negative. + elif stop < 0: + # `neg_index_as_lookback=True` -> User wants to index into the lookback + # range. Set to 0 (beginning of lookback buffer) if result is a negative + # index. + if neg_index_as_lookback: + stop = self.lookback + stop + # Interpret index as counting "from end". Set to 0 (beginning of actual + # episode) if result is a negative index. + else: + stop = len_self_plus_lookback + stop + # Stop is positive -> Add lookback range to it. + else: + stop = self.lookback + stop + + fill_left_count = fill_right_count = 0 + # Both start and stop are on left side. + if start < 0 and stop < 0: + fill_left_count = abs(start - stop) + fill_right_count = 0 + start = stop = 0 + # Both start and stop are on right side. + elif start >= len_self_plus_lookback and stop >= len_self_plus_lookback: + fill_right_count = abs(start - stop) + fill_left_count = 0 + start = stop = len_self_plus_lookback + # Set to 0 (beginning of actual episode) if result is a negative index. + elif start < 0: + fill_left_count = -start + start = 0 + elif stop >= len_self_plus_lookback: + fill_right_count = stop - len_self_plus_lookback + stop = len_self_plus_lookback + # Only `stop` might be < 0, when slice has negative step and start is > 0. + elif stop < 0: + if start >= len_self_plus_lookback: + fill_left_count = start - len_self_plus_lookback + 1 + start = len_self_plus_lookback - 1 + fill_right_count = -stop - 1 + stop = -LARGE_INTEGER + + assert start >= 0 and (stop >= 0 or stop == -LARGE_INTEGER), (start, stop) + + step = slice_.step if slice_.step is not None else 1 + slice_ = slice(start, stop, step) + slice_len = max(0, (stop - start + (step - (1 if step > 0 else -1))) // step) + return slice_, slice_len, fill_left_count, fill_right_count + + def _one_hot(self, data, space_struct): + if space_struct is None: + raise ValueError( + f"Cannot `one_hot` data in `{type(self).__name__}` if a " + "gym.Space was NOT provided during construction!" + ) + + def _convert(dat_, space): + if isinstance(space, gym.spaces.Discrete): + return one_hot(dat_, depth=space.n) + elif isinstance(space, gym.spaces.MultiDiscrete): + return one_hot_multidiscrete(dat_, depths=space.nvec) + return dat_ + + if isinstance(data, list): + data = [ + tree.map_structure(_convert, dslice, space_struct) for dslice in data + ] + else: + data = tree.map_structure(_convert, data, space_struct) + return data diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/__init__.py b/lib/python3.12/site-packages/ray/rllib/env/vector/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b56a482787f3b79e6a20069745a6839cbde14d6d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/registration.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/registration.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27a23454ad9a183ab9d8d312ca814d58ef6cf0b7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/registration.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/sync_vector_multi_agent_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/sync_vector_multi_agent_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7b0cf8ee3983831cc3e31445df6265e46f15539 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/sync_vector_multi_agent_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/vector_multi_agent_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/vector_multi_agent_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6aa7971c72a02f30db14369b16bacd61f89408c1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/vector/__pycache__/vector_multi_agent_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/registration.py b/lib/python3.12/site-packages/ray/rllib/env/vector/registration.py new file mode 100644 index 0000000000000000000000000000000000000000..ec0a8b43e6331abcbc868176ef473b112bdf8149 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/vector/registration.py @@ -0,0 +1,91 @@ +import copy +import logging +from typing import Any, Dict, Optional + +import gymnasium as gym +from gymnasium.envs.registration import VectorizeMode + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.vector.sync_vector_multi_agent_env import SyncVectorMultiAgentEnv +from ray.rllib.env.vector.vector_multi_agent_env import VectorMultiAgentEnv + +logger = logging.getLogger(__file__) + + +def make_vec( + id: str, + num_envs: int = 1, + vectorization_mode: Optional[VectorizeMode] = None, + vector_kwargs: Optional[Dict[str, Any]] = None, + # TODO (simon): Add wrappers? + **kwargs, +) -> VectorMultiAgentEnv: + + if vector_kwargs is None: + vector_kwargs = {} + + if vectorization_mode is None: + vectorization_mode = "sync" + + # Create an `gymnasium.envs.registration.EnvSpec` to properly + # initialize the sub-environments. + if isinstance(id, gym.envs.registration.EnvSpec): + env_spec = id + elif isinstance(id, str): + env_spec = gym.envs.registration._find_spec(id) + else: + raise ValueError(f"Invalid id type: {type(id)}. Expected `str` or `EnvSpec`.") + + env_spec = copy.deepcopy(env_spec) + env_spec_kwargs = env_spec.kwargs + env_spec.kwargs = dict() + + num_envs = env_spec.kwargs.get("num_envs", num_envs) + vectorization_mode = env_spec_kwargs.pop("vectorization_mode", vectorization_mode) + vector_kwargs = env_spec_kwargs.pop("vector_kwargs", vector_kwargs) + env_spec_kwargs.update(kwargs) + + # Specify the vectorization mode. + if vectorization_mode is None: + vectorization_mode = VectorizeMode.SYNC + else: + try: + vectorization_mode = VectorizeMode(vectorization_mode) + except ValueError: + raise ValueError( + f"Invalid vectorization mode: {vectorization_mode!r}, " + f"valid modes: {[mode.value for mode in VectorizeMode]}." + ) + assert isinstance(vectorization_mode, VectorizeMode) + + def create_single_env() -> MultiAgentEnv: + single_env = gym.make(env_spec, **env_spec_kwargs.copy()) + + return single_env + + # Check, the vectorization mode. + if vectorization_mode == VectorizeMode.SYNC: + # Create the synchronized vector environemnt. + env = SyncVectorMultiAgentEnv( + env_fns=(create_single_env for _ in range(num_envs)), + **vector_kwargs, + ) + # Other modes are not implemented, yet. + else: + raise ValueError( + "For `MultiAgentEnv` only synchronous environment vectorization " + "is implemented. Use `gym_env_vectorize_mode='sync'`." + ) + + # Add all creation specifications to the environment. + copied_id_spec = copy.deepcopy(env_spec) + copied_id_spec.kwargs = env_spec_kwargs.copy() + if num_envs != 1: + copied_id_spec.kwargs["num_envs"] = num_envs + copied_id_spec.kwargs["vectorization_mode"] = vectorization_mode.value + if len(vector_kwargs) > 0: + copied_id_spec.kwargs["vector_kwargs"] = vector_kwargs + env.unwrapped.spec = copied_id_spec + + # Return the `VectorMultiAgentEnv`. + return env diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/sync_vector_multi_agent_env.py b/lib/python3.12/site-packages/ray/rllib/env/vector/sync_vector_multi_agent_env.py new file mode 100644 index 0000000000000000000000000000000000000000..e5c375526f76b8a831f0621bddc531395106c6ad --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/vector/sync_vector_multi_agent_env.py @@ -0,0 +1,215 @@ +from copy import deepcopy +from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union + +import gymnasium as gym +import numpy as np +from gymnasium.core import ActType, RenderFrame + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.vector.vector_multi_agent_env import ArrayType, VectorMultiAgentEnv +from ray.rllib.utils.typing import AgentID + + +class SyncVectorMultiAgentEnv(VectorMultiAgentEnv): + def __init__( + self, + env_fns: Union[ + Iterator[Callable[[], MultiAgentEnv]], Sequence[Callable[[], MultiAgentEnv]] + ], + copy: bool = True, + autoreset_mode: str = "next_step", + ): + super().__init__() + + self.env_fns = env_fns + self.copy = copy + self.autoreset_mode = autoreset_mode + + # Initialize all sub-environments. + self.envs = [env_fn() for env_fn in self.env_fns] + + self.num_envs = len(self.envs) + self.metadata = self.envs[0].metadata + self.metadata["autoreset_mode"] = self.autoreset_mode + self.render_mode = self.envs[0].render_mode + + if self.envs[0].unwrapped.observation_spaces is not None: + self.single_observation_spaces = self.envs[0].unwrapped.observation_spaces + elif self.envs[0].unwrapped.observation_space is not None: + self.single_observation_spaces = dict( + self.envs[0].unwrapped.observation_space + ) + else: + self.single_observation_spaces = { + aid: self.envs[0].unwrapped.get_observation_space(aid) + for aid in self.envs[0].unwrapped.possible_agents + } + self.single_observation_space = gym.spaces.Dict(self.single_observation_spaces) + + if self.envs[0].unwrapped.action_spaces is not None: + self.single_action_spaces = self.envs[0].unwrapped.action_spaces + elif self.envs[0].unwrapped.action_space is not None: + self.single_action_spaces = dict(self.envs[0].unwrapped.action_space) + else: + self.single_action_spaces = { + aid: self.envs[0].unwrapped.get_action_space(aid) + for aid in self.envs[0].unwrapped.possible_agents + } + self.single_action_space = gym.spaces.Dict(self.single_action_spaces) + + # TODO (simon): Decide if we want to include a spaces check here. + + # Note, if `single_observation_spaces` are not defined, this will + # raise an error. + self._observations = self.create_empty_array( + self.single_observation_spaces, + n=self.num_envs, + ) + self._rewards = [{} for _ in range(self.num_envs)] + self._terminations = [{} for _ in range(self.num_envs)] + self._truncations = [{} for _ in range(self.num_envs)] + self._infos = [{} for _ in range(self.num_envs)] + + def reset( + self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None + ) -> Tuple[ArrayType, ArrayType]: + + # Check the `seed`. + if seed is None: + seed = [None for _ in range(self.num_envs)] + elif isinstance(seed, int): + seed = [seed + i for i in range(self.num_envs)] + # If `seed` is a sequence, check if length is correct. + else: + assert len(seed) == self.num_envs, ( + f"If seeds are passed as a list the length must match num_envs={self.num_envs} " + f"but got length={len(seed)}.", + ) + + # Set auto-reset to `False` for all sub-environments. This will be + # set to `True` during `VectorMultiAgentEnv.step`. + self._autoreset_envs = np.zeros((self.num_envs,), dtype=np.bool_) + + # Force reset on all sub-environments. + for i, (env, single_seed) in enumerate(zip(self.envs, seed)): + self._observations[i], self._infos[i] = env.reset( + seed=single_seed, options=options + ) + + return ( + deepcopy(self._observations) if self.copy else self._observations, + self._infos, + ) + + def step( + self, actions: List[Dict[AgentID, ActType]] + ) -> Tuple[ArrayType, ArrayType, ArrayType, ArrayType, ArrayType]: + + for i, action in enumerate(actions): + # Note, this will be a feature coming in `gymnasium`s next release. + if self.autoreset_mode == "next_step": + # Auto-reset environments that terminated or truncated. + if self._autoreset_envs[i]: + self._observations[i], self._infos[i] = self.envs[i].reset() + + self._rewards[i] = { + aid: 0.0 for aid in self.envs[i].unwrapped.possible_agents + } + self._terminations[i] = { + aid: False for aid in self.envs[i].unwrapped.possible_agents + } + self._truncations[i] = { + aid: False for aid in self.envs[i].unwrapped.possible_agents + } + # Note, this needs to be added for `MultiAgentEpisode`s to work. + self._terminations[i].update({"__all__": False}) + self._truncations[i].update({"__all__": False}) + # Otherwise make a regular step. + else: + ( + self._observations[i], + self._rewards[i], + self._terminations[i], + self._truncations[i], + self._infos[i], + ) = self.envs[i].step(action) + # Other autoreset modes are not implemented, yet. + else: + raise ValueError(f"Unexpected autoreset mode: {self.autoreset_mode}") + + # Is any sub-environment terminated or truncated? + self._autoreset_envs = np.logical_or( + np.array([t["__all__"] for t in self._terminations]), + # Note, some environments just return `terminated`. + np.array([t.get("__all__", False) for t in self._truncations]), + ) + + return ( + deepcopy(self._observations) if self.copy else self._observations, + np.copy(self._rewards), + np.copy(self._terminations), + np.copy(self._truncations), + np.copy(self._infos), + ) + + def render(self) -> Optional[Tuple[RenderFrame, ...]]: + """Returns the rendered frames from all environments.""" + return tuple(env.render() for env in self.envs) + + def call(self, name: str, *args: Any, **kwargs: Any) -> tuple[Any, ...]: + """Calls a sub-environment method with name and applies args and kwargs. + + Args: + name: The method name + *args: The method args + **kwargs: The method kwargs + + Returns: + Tuple of results + """ + results = [] + for env in self.envs: + function = env.get_wrapper_attr(name) + + if callable(function): + results.append(function(*args, **kwargs)) + else: + results.append(function) + + return tuple(results) + + def get_attr(self, name: str) -> tuple[Any, ...]: + """Get a property from each parallel environment. + + Args: + name: Name of the property to get from each individual environment. + + Returns: + The property with name + """ + return self.call(name) + + def close_extras(self, **kwargs): + """Close all environments.""" + if hasattr(self, "envs"): + [env.close() for env in self.envs] + + @property + def np_random_seed(self) -> Tuple[int, ...]: + """Returns a tuple of random seeds for the wrapped envs.""" + return self.get_attr("np_random_seed") + + @property + def np_random(self) -> Tuple[np.random.Generator, ...]: + """Returns a tuple of numpy random generators for the wrapped envs.""" + return self.get_attr("np_random") + + # TODO (simon): Move to space_utils. + def create_empty_array( + self, space: Dict[str, gym.Space], n: int = 1 + ) -> List[Dict[str, gym.Space]]: + + return [ + {aid: agent_space.sample() for aid, agent_space in space.items()} + for _ in range(n) + ] diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector/vector_multi_agent_env.py b/lib/python3.12/site-packages/ray/rllib/env/vector/vector_multi_agent_env.py new file mode 100644 index 0000000000000000000000000000000000000000..a4e712643c18d43037d39a04d0231a65bb043f9d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/vector/vector_multi_agent_env.py @@ -0,0 +1,87 @@ +from typing import Any, Dict, List, Optional, Tuple, TypeVar + +import gymnasium as gym +import numpy as np +from gymnasium.core import RenderFrame +from gymnasium.envs.registration import EnvSpec +from gymnasium.utils import seeding + +ArrayType = TypeVar("ArrayType") + + +class VectorMultiAgentEnv: + + metadata: Dict[str, Any] = {} + spec: Optional[EnvSpec] = None + render_mode: Optional[str] = None + closed: bool = False + + envs: Optional[List] = None + + # TODO (simon, sven): We could think about enabling here different + # spaces for different envs (e.g. different high/lows). In this + # case we would need here actually "batched" spaces and not a + # single on that holds for all sub-envs. + single_observation_spaces: Optional[Dict[str, gym.Space]] = None + single_action_spaces: Optional[Dict[str, gym.Space]] = None + # Note, the proper `gym` spaces are needed for the connector pipeline. + single_observation_space: Optional[gym.spaces.Dict] = None + single_action_space: Optional[gym.spaces.Dict] = None + + num_envs: int + + _np_random: Optional[np.random.Generator] = None + _np_random_seed: Optional[int] = None + + # @OldAPIStack, use `observation_spaces` and `action_spaces`, instead. + observation_space: Optional[gym.Space] = None + action_space: Optional[gym.Space] = None + + # TODO (simon): Add docstrings, when final design is clear. + def reset( + self, *, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None + ) -> Tuple[ArrayType, ArrayType]: + # Set random generators with the provided seeds. + if seed is not None: + self._np_random, self._np_random_seed = seeding.np_random(seed) + + def step( + self, actions: ArrayType + ) -> Tuple[ArrayType, ArrayType, ArrayType, ArrayType, ArrayType]: + raise NotImplementedError(f"{self.__str__()} step function is not implemented.") + + def render(self) -> Optional[Tuple[RenderFrame, ...]]: + raise NotImplementedError( + f"{self.__str__()} render function is not implemented." + ) + + def close(self, **kwargs: Any): + + # If already closed, there is nothing more to do. + if self.closed: + return + + # Otherwise close environments gracefully. + self.close_extras(**kwargs) + self.closed = True + + def close_extras(self, **kwargs: Any): + # Users must not implement this. + pass + + @property + def unwrapped(self): + return self + + def __del__(self): + # Close environemnts, if necessary when deleting instances. + if not getattr(self, "closed", True): + self.close() + + def __repr__(self): + if self.spec is None: + return f"{self.__class__.__name__}(num_envs={self.num_envs})" + else: + return ( + f"{self.__class__.__name__}({self.spec.id}, num_envs={self.num_envs})" + ) diff --git a/lib/python3.12/site-packages/ray/rllib/env/vector_env.py b/lib/python3.12/site-packages/ray/rllib/env/vector_env.py new file mode 100644 index 0000000000000000000000000000000000000000..ed1e2dfdb70c1fa051bc761641a3989fd84b2451 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/vector_env.py @@ -0,0 +1,542 @@ +import logging +from typing import Callable, List, Optional, Set, Tuple, Union + +import gymnasium as gym +import numpy as np + +from ray.rllib.env.base_env import _DUMMY_AGENT_ID, BaseEnv +from ray.rllib.utils.annotations import Deprecated, OldAPIStack, override +from ray.rllib.utils.typing import ( + AgentID, + EnvActionType, + EnvID, + EnvInfoDict, + EnvObsType, + EnvType, + MultiEnvDict, +) +from ray.util import log_once + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class VectorEnv: + """An environment that supports batch evaluation using clones of sub-envs.""" + + def __init__( + self, observation_space: gym.Space, action_space: gym.Space, num_envs: int + ): + """Initializes a VectorEnv instance. + + Args: + observation_space: The observation Space of a single + sub-env. + action_space: The action Space of a single sub-env. + num_envs: The number of clones to make of the given sub-env. + """ + self.observation_space = observation_space + self.action_space = action_space + self.num_envs = num_envs + + @staticmethod + def vectorize_gym_envs( + make_env: Optional[Callable[[int], EnvType]] = None, + existing_envs: Optional[List[gym.Env]] = None, + num_envs: int = 1, + action_space: Optional[gym.Space] = None, + observation_space: Optional[gym.Space] = None, + restart_failed_sub_environments: bool = False, + # Deprecated. These seem to have never been used. + env_config=None, + policy_config=None, + ) -> "_VectorizedGymEnv": + """Translates any given gym.Env(s) into a VectorizedEnv object. + + Args: + make_env: Factory that produces a new gym.Env taking the sub-env's + vector index as only arg. Must be defined if the + number of `existing_envs` is less than `num_envs`. + existing_envs: Optional list of already instantiated sub + environments. + num_envs: Total number of sub environments in this VectorEnv. + action_space: The action space. If None, use existing_envs[0]'s + action space. + observation_space: The observation space. If None, use + existing_envs[0]'s observation space. + restart_failed_sub_environments: If True and any sub-environment (within + a vectorized env) throws any error during env stepping, the + Sampler will try to restart the faulty sub-environment. This is done + without disturbing the other (still intact) sub-environment and without + the RolloutWorker crashing. + + Returns: + The resulting _VectorizedGymEnv object (subclass of VectorEnv). + """ + return _VectorizedGymEnv( + make_env=make_env, + existing_envs=existing_envs or [], + num_envs=num_envs, + observation_space=observation_space, + action_space=action_space, + restart_failed_sub_environments=restart_failed_sub_environments, + ) + + def vector_reset( + self, *, seeds: Optional[List[int]] = None, options: Optional[List[dict]] = None + ) -> Tuple[List[EnvObsType], List[EnvInfoDict]]: + """Resets all sub-environments. + + Args: + seed: The list of seeds to be passed to the sub-environments' when resetting + them. If None, will not reset any existing PRNGs. If you pass + integers, the PRNGs will be reset even if they already exists. + options: The list of options dicts to be passed to the sub-environments' + when resetting them. + + Returns: + Tuple consitsing of a list of observations from each environment and + a list of info dicts from each environment. + """ + raise NotImplementedError + + def reset_at( + self, + index: Optional[int] = None, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Union[Tuple[EnvObsType, EnvInfoDict], Exception]: + """Resets a single sub-environment. + + Args: + index: An optional sub-env index to reset. + seed: The seed to be passed to the sub-environment at index `index` when + resetting it. If None, will not reset any existing PRNG. If you pass an + integer, the PRNG will be reset even if it already exists. + options: An options dict to be passed to the sub-environment at index + `index` when resetting it. + + Returns: + Tuple consisting of observations from the reset sub environment and + an info dict of the reset sub environment. Alternatively an Exception + can be returned, indicating that the reset operation on the sub environment + has failed (and why it failed). + """ + raise NotImplementedError + + def restart_at(self, index: Optional[int] = None) -> None: + """Restarts a single sub-environment. + + Args: + index: An optional sub-env index to restart. + """ + raise NotImplementedError + + def vector_step( + self, actions: List[EnvActionType] + ) -> Tuple[ + List[EnvObsType], List[float], List[bool], List[bool], List[EnvInfoDict] + ]: + """Performs a vectorized step on all sub environments using `actions`. + + Args: + actions: List of actions (one for each sub-env). + + Returns: + A tuple consisting of + 1) New observations for each sub-env. + 2) Reward values for each sub-env. + 3) Terminated values for each sub-env. + 4) Truncated values for each sub-env. + 5) Info values for each sub-env. + """ + raise NotImplementedError + + def get_sub_environments(self) -> List[EnvType]: + """Returns the underlying sub environments. + + Returns: + List of all underlying sub environments. + """ + return [] + + # TODO: (sven) Experimental method. Make @PublicAPI at some point. + def try_render_at(self, index: Optional[int] = None) -> Optional[np.ndarray]: + """Renders a single environment. + + Args: + index: An optional sub-env index to render. + + Returns: + Either a numpy RGB image (shape=(w x h x 3) dtype=uint8) or + None in case rendering is handled directly by this method. + """ + pass + + def to_base_env( + self, + make_env: Optional[Callable[[int], EnvType]] = None, + num_envs: int = 1, + remote_envs: bool = False, + remote_env_batch_wait_ms: int = 0, + restart_failed_sub_environments: bool = False, + ) -> "BaseEnv": + """Converts an RLlib MultiAgentEnv into a BaseEnv object. + + The resulting BaseEnv is always vectorized (contains n + sub-environments) to support batched forward passes, where n may + also be 1. BaseEnv also supports async execution via the `poll` and + `send_actions` methods and thus supports external simulators. + + Args: + make_env: A callable taking an int as input (which indicates + the number of individual sub-environments within the final + vectorized BaseEnv) and returning one individual + sub-environment. + num_envs: The number of sub-environments to create in the + resulting (vectorized) BaseEnv. The already existing `env` + will be one of the `num_envs`. + remote_envs: Whether each sub-env should be a @ray.remote + actor. You can set this behavior in your config via the + `remote_worker_envs=True` option. + remote_env_batch_wait_ms: The wait time (in ms) to poll remote + sub-environments for, if applicable. Only used if + `remote_envs` is True. + + Returns: + The resulting BaseEnv object. + """ + env = VectorEnvWrapper(self) + return env + + @Deprecated(new="vectorize_gym_envs", error=True) + def wrap(self, *args, **kwargs) -> "_VectorizedGymEnv": + pass + + @Deprecated(new="get_sub_environments", error=True) + def get_unwrapped(self) -> List[EnvType]: + pass + + +@OldAPIStack +class _VectorizedGymEnv(VectorEnv): + """Internal wrapper to translate any gym.Envs into a VectorEnv object.""" + + def __init__( + self, + make_env: Optional[Callable[[int], EnvType]] = None, + existing_envs: Optional[List[gym.Env]] = None, + num_envs: int = 1, + *, + observation_space: Optional[gym.Space] = None, + action_space: Optional[gym.Space] = None, + restart_failed_sub_environments: bool = False, + # Deprecated. These seem to have never been used. + env_config=None, + policy_config=None, + ): + """Initializes a _VectorizedGymEnv object. + + Args: + make_env: Factory that produces a new gym.Env taking the sub-env's + vector index as only arg. Must be defined if the + number of `existing_envs` is less than `num_envs`. + existing_envs: Optional list of already instantiated sub + environments. + num_envs: Total number of sub environments in this VectorEnv. + action_space: The action space. If None, use existing_envs[0]'s + action space. + observation_space: The observation space. If None, use + existing_envs[0]'s observation space. + restart_failed_sub_environments: If True and any sub-environment (within + a vectorized env) throws any error during env stepping, we will try to + restart the faulty sub-environment. This is done + without disturbing the other (still intact) sub-environments. + """ + self.envs = existing_envs + self.make_env = make_env + self.restart_failed_sub_environments = restart_failed_sub_environments + + # Fill up missing envs (so we have exactly num_envs sub-envs in this + # VectorEnv. + while len(self.envs) < num_envs: + self.envs.append(make_env(len(self.envs))) + + super().__init__( + observation_space=observation_space or self.envs[0].observation_space, + action_space=action_space or self.envs[0].action_space, + num_envs=num_envs, + ) + + @override(VectorEnv) + def vector_reset( + self, *, seeds: Optional[List[int]] = None, options: Optional[List[dict]] = None + ) -> Tuple[List[EnvObsType], List[EnvInfoDict]]: + seeds = seeds or [None] * self.num_envs + options = options or [None] * self.num_envs + # Use reset_at(index) to restart and retry until + # we successfully create a new env. + resetted_obs = [] + resetted_infos = [] + for i in range(len(self.envs)): + while True: + obs, infos = self.reset_at(i, seed=seeds[i], options=options[i]) + if not isinstance(obs, Exception): + break + resetted_obs.append(obs) + resetted_infos.append(infos) + return resetted_obs, resetted_infos + + @override(VectorEnv) + def reset_at( + self, + index: Optional[int] = None, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Tuple[Union[EnvObsType, Exception], Union[EnvInfoDict, Exception]]: + if index is None: + index = 0 + try: + obs_and_infos = self.envs[index].reset(seed=seed, options=options) + + except Exception as e: + if self.restart_failed_sub_environments: + logger.exception(e.args[0]) + self.restart_at(index) + obs_and_infos = e, {} + else: + raise e + + return obs_and_infos + + @override(VectorEnv) + def restart_at(self, index: Optional[int] = None) -> None: + if index is None: + index = 0 + + # Try closing down the old (possibly faulty) sub-env, but ignore errors. + try: + self.envs[index].close() + except Exception as e: + if log_once("close_sub_env"): + logger.warning( + "Trying to close old and replaced sub-environment (at vector " + f"index={index}), but closing resulted in error:\n{e}" + ) + env_to_del = self.envs[index] + self.envs[index] = None + del env_to_del + + # Re-create the sub-env at the new index. + logger.warning(f"Trying to restart sub-environment at index {index}.") + self.envs[index] = self.make_env(index) + logger.warning(f"Sub-environment at index {index} restarted successfully.") + + @override(VectorEnv) + def vector_step( + self, actions: List[EnvActionType] + ) -> Tuple[ + List[EnvObsType], List[float], List[bool], List[bool], List[EnvInfoDict] + ]: + obs_batch, reward_batch, terminated_batch, truncated_batch, info_batch = ( + [], + [], + [], + [], + [], + ) + for i in range(self.num_envs): + try: + results = self.envs[i].step(actions[i]) + except Exception as e: + if self.restart_failed_sub_environments: + logger.exception(e.args[0]) + self.restart_at(i) + results = e, 0.0, True, True, {} + else: + raise e + + obs, reward, terminated, truncated, info = results + + if not isinstance(info, dict): + raise ValueError( + "Info should be a dict, got {} ({})".format(info, type(info)) + ) + obs_batch.append(obs) + reward_batch.append(reward) + terminated_batch.append(terminated) + truncated_batch.append(truncated) + info_batch.append(info) + return obs_batch, reward_batch, terminated_batch, truncated_batch, info_batch + + @override(VectorEnv) + def get_sub_environments(self) -> List[EnvType]: + return self.envs + + @override(VectorEnv) + def try_render_at(self, index: Optional[int] = None): + if index is None: + index = 0 + return self.envs[index].render() + + +@OldAPIStack +class VectorEnvWrapper(BaseEnv): + """Internal adapter of VectorEnv to BaseEnv. + + We assume the caller will always send the full vector of actions in each + call to send_actions(), and that they call reset_at() on all completed + environments before calling send_actions(). + """ + + def __init__(self, vector_env: VectorEnv): + self.vector_env = vector_env + self.num_envs = vector_env.num_envs + self._observation_space = vector_env.observation_space + self._action_space = vector_env.action_space + + # Sub-environments' states. + self.new_obs = None + self.cur_rewards = None + self.cur_terminateds = None + self.cur_truncateds = None + self.cur_infos = None + # At first `poll()`, reset everything (all sub-environments). + self.first_reset_done = False + # Initialize sub-environments' state. + self._init_env_state(idx=None) + + @override(BaseEnv) + def poll( + self, + ) -> Tuple[ + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + MultiEnvDict, + ]: + from ray.rllib.env.base_env import with_dummy_agent_id + + if not self.first_reset_done: + self.first_reset_done = True + # TODO(sven): We probably would like to seed this call here as well. + self.new_obs, self.cur_infos = self.vector_env.vector_reset() + new_obs = dict(enumerate(self.new_obs)) + rewards = dict(enumerate(self.cur_rewards)) + terminateds = dict(enumerate(self.cur_terminateds)) + truncateds = dict(enumerate(self.cur_truncateds)) + infos = dict(enumerate(self.cur_infos)) + + # Empty all states (in case `poll()` gets called again). + self.new_obs = [] + self.cur_rewards = [] + self.cur_terminateds = [] + self.cur_truncateds = [] + self.cur_infos = [] + + return ( + with_dummy_agent_id(new_obs), + with_dummy_agent_id(rewards), + with_dummy_agent_id(terminateds, "__all__"), + with_dummy_agent_id(truncateds, "__all__"), + with_dummy_agent_id(infos), + {}, + ) + + @override(BaseEnv) + def send_actions(self, action_dict: MultiEnvDict) -> None: + from ray.rllib.env.base_env import _DUMMY_AGENT_ID + + action_vector = [None] * self.num_envs + for i in range(self.num_envs): + action_vector[i] = action_dict[i][_DUMMY_AGENT_ID] + ( + self.new_obs, + self.cur_rewards, + self.cur_terminateds, + self.cur_truncateds, + self.cur_infos, + ) = self.vector_env.vector_step(action_vector) + + @override(BaseEnv) + def try_reset( + self, + env_id: Optional[EnvID] = None, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> Tuple[MultiEnvDict, MultiEnvDict]: + from ray.rllib.env.base_env import _DUMMY_AGENT_ID + + if env_id is None: + env_id = 0 + assert isinstance(env_id, int) + obs, infos = self.vector_env.reset_at(env_id, seed=seed, options=options) + + # If exceptions were returned, return MultiEnvDict mapping env indices to + # these exceptions (for obs and infos). + if isinstance(obs, Exception): + return {env_id: obs}, {env_id: infos} + # Otherwise, return a MultiEnvDict (with single agent ID) and the actual + # obs and info dicts. + else: + return {env_id: {_DUMMY_AGENT_ID: obs}}, {env_id: {_DUMMY_AGENT_ID: infos}} + + @override(BaseEnv) + def try_restart(self, env_id: Optional[EnvID] = None) -> None: + assert env_id is None or isinstance(env_id, int) + # Restart the sub-env at the index. + self.vector_env.restart_at(env_id) + # Auto-reset (get ready for next `poll()`). + self._init_env_state(env_id) + + @override(BaseEnv) + def get_sub_environments(self, as_dict: bool = False) -> Union[List[EnvType], dict]: + if not as_dict: + return self.vector_env.get_sub_environments() + else: + return dict(enumerate(self.vector_env.get_sub_environments())) + + @override(BaseEnv) + def try_render(self, env_id: Optional[EnvID] = None) -> None: + assert env_id is None or isinstance(env_id, int) + return self.vector_env.try_render_at(env_id) + + @property + @override(BaseEnv) + def observation_space(self) -> gym.Space: + return self._observation_space + + @property + @override(BaseEnv) + def action_space(self) -> gym.Space: + return self._action_space + + @override(BaseEnv) + def get_agent_ids(self) -> Set[AgentID]: + return {_DUMMY_AGENT_ID} + + def _init_env_state(self, idx: Optional[int] = None) -> None: + """Resets all or one particular sub-environment's state (by index). + + Args: + idx: The index to reset at. If None, reset all the sub-environments' states. + """ + # If index is None, reset all sub-envs' states: + if idx is None: + self.new_obs = [None for _ in range(self.num_envs)] + self.cur_rewards = [0.0 for _ in range(self.num_envs)] + self.cur_terminateds = [False for _ in range(self.num_envs)] + self.cur_truncateds = [False for _ in range(self.num_envs)] + self.cur_infos = [{} for _ in range(self.num_envs)] + # Index provided, reset only the sub-env's state at the given index. + else: + self.new_obs[idx], self.cur_infos[idx] = self.vector_env.reset_at(idx) + # Reset all other states to null values. + self.cur_rewards[idx] = 0.0 + self.cur_terminateds[idx] = False + self.cur_truncateds[idx] = False diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__init__.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5729a2d0bc271e47035e46f28071abf36d9c1395 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/atari_wrappers.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/atari_wrappers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d56f569baac267a48237ce58422cd0f85da51ce4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/atari_wrappers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/dm_control_wrapper.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/dm_control_wrapper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c3da3bf77688962885040b46c722ab01a305448 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/dm_control_wrapper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/dm_env_wrapper.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/dm_env_wrapper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b5709a3cf25ec91c5a719e5bb7cc2ca1d687072 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/dm_env_wrapper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/group_agents_wrapper.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/group_agents_wrapper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a392d655b29117d6347c7e7549f09dd014c11845 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/group_agents_wrapper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/multi_agent_env_compatibility.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/multi_agent_env_compatibility.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d73b5cd43be733de46b84f50a9ecb19e6dd85ab7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/multi_agent_env_compatibility.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/open_spiel.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/open_spiel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7eb768f64c2573eb3cb9de110bb950526f8e3d81 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/open_spiel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/pettingzoo_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/pettingzoo_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47886498c28aafb74716561d6c9f9357fb358df4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/pettingzoo_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/unity3d_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/unity3d_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..449c3ea07d5f95b558e2892a21b92147446b9a05 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/env/wrappers/__pycache__/unity3d_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/atari_wrappers.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/atari_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..3b20ea23221e1e5b2c9326789ed8a03986144bae --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/atari_wrappers.py @@ -0,0 +1,412 @@ +from collections import deque +from typing import Optional, Union + +import gymnasium as gym +import numpy as np +from gymnasium import spaces + +from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.images import resize, rgb2gray + + +@PublicAPI +def is_atari(env: Union[gym.Env, str]) -> bool: + """Returns, whether a given env object or env descriptor (str) is an Atari env. + + Args: + env: The gym.Env object or a string descriptor of the env (for example, + "ale_py:ALE/Pong-v5"). + + Returns: + Whether `env` is an Atari environment. + """ + # If a gym.Env, check proper spaces as well as occurrence of the "Atari 0: + # for Qbert sometimes we stay in lives == 0 condtion for a few fr + # so its important to keep lives > 0, so that we only reset once + # the environment advertises `terminated`. + terminated = True + self.lives = lives + return obs, reward, terminated, truncated, info + + def reset(self, **kwargs): + """Reset only when lives are exhausted. + This way all states are still reachable even though lives are episodic, + and the learner need not know about any of this behind-the-scenes. + """ + if self.was_real_terminated: + obs, info = self.env.reset(**kwargs) + else: + # no-op step to advance from terminal/lost life state + obs, _, _, _, info = self.env.step(0) + self.lives = self.env.unwrapped.ale.lives() + return obs, info + + +@PublicAPI +class FireResetEnv(gym.Wrapper): + def __init__(self, env): + """Take action on reset. + + For environments that are fixed until firing.""" + gym.Wrapper.__init__(self, env) + assert env.unwrapped.get_action_meanings()[1] == "FIRE" + assert len(env.unwrapped.get_action_meanings()) >= 3 + + def reset(self, **kwargs): + self.env.reset(**kwargs) + obs, _, terminated, truncated, _ = self.env.step(1) + if terminated or truncated: + self.env.reset(**kwargs) + obs, _, terminated, truncated, info = self.env.step(2) + if terminated or truncated: + self.env.reset(**kwargs) + return obs, info + + def step(self, ac): + return self.env.step(ac) + + +@PublicAPI +class FrameStack(gym.Wrapper): + def __init__(self, env, k): + """Stack k last frames.""" + gym.Wrapper.__init__(self, env) + self.k = k + self.frames = deque([], maxlen=k) + shp = env.observation_space.shape + self.observation_space = spaces.Box( + low=np.repeat(env.observation_space.low, repeats=k, axis=-1), + high=np.repeat(env.observation_space.high, repeats=k, axis=-1), + shape=(shp[0], shp[1], shp[2] * k), + dtype=env.observation_space.dtype, + ) + + def reset(self, *, seed=None, options=None): + ob, infos = self.env.reset(seed=seed, options=options) + for _ in range(self.k): + self.frames.append(ob) + return self._get_ob(), infos + + def step(self, action): + ob, reward, terminated, truncated, info = self.env.step(action) + self.frames.append(ob) + return self._get_ob(), reward, terminated, truncated, info + + def _get_ob(self): + assert len(self.frames) == self.k + return np.concatenate(self.frames, axis=2) + + +@PublicAPI +class FrameStackTrajectoryView(gym.ObservationWrapper): + def __init__(self, env): + """No stacking. Trajectory View API takes care of this.""" + gym.Wrapper.__init__(self, env) + shp = env.observation_space.shape + assert shp[2] == 1 + self.observation_space = spaces.Box( + low=0, high=255, shape=(shp[0], shp[1]), dtype=env.observation_space.dtype + ) + + def observation(self, observation): + return np.squeeze(observation, axis=-1) + + +@PublicAPI +class MaxAndSkipEnv(gym.Wrapper): + def __init__(self, env, skip=4): + """Return only every `skip`-th frame""" + gym.Wrapper.__init__(self, env) + # most recent raw observations (for max pooling across time steps) + self._obs_buffer = np.zeros( + (2,) + env.observation_space.shape, dtype=env.observation_space.dtype + ) + self._skip = skip + + def step(self, action): + """Repeat action, sum reward, and max over last observations.""" + total_reward = 0.0 + terminated = truncated = info = None + for i in range(self._skip): + obs, reward, terminated, truncated, info = self.env.step(action) + if i == self._skip - 2: + self._obs_buffer[0] = obs + if i == self._skip - 1: + self._obs_buffer[1] = obs + total_reward += reward + if terminated or truncated: + break + # Note that the observation on the terminated|truncated=True frame + # doesn't matter + max_frame = self._obs_buffer.max(axis=0) + + return max_frame, total_reward, terminated, truncated, info + + def reset(self, **kwargs): + return self.env.reset(**kwargs) + + +@PublicAPI +class MonitorEnv(gym.Wrapper): + def __init__(self, env=None): + """Record episodes stats prior to EpisodicLifeEnv, etc.""" + gym.Wrapper.__init__(self, env) + self._current_reward = None + self._num_steps = None + self._total_steps = None + self._episode_rewards = [] + self._episode_lengths = [] + self._num_episodes = 0 + self._num_returned = 0 + + def reset(self, **kwargs): + obs, info = self.env.reset(**kwargs) + + if self._total_steps is None: + self._total_steps = sum(self._episode_lengths) + + if self._current_reward is not None: + self._episode_rewards.append(self._current_reward) + self._episode_lengths.append(self._num_steps) + self._num_episodes += 1 + + self._current_reward = 0 + self._num_steps = 0 + + return obs, info + + def step(self, action): + obs, rew, terminated, truncated, info = self.env.step(action) + self._current_reward += rew + self._num_steps += 1 + self._total_steps += 1 + return obs, rew, terminated, truncated, info + + def get_episode_rewards(self): + return self._episode_rewards + + def get_episode_lengths(self): + return self._episode_lengths + + def get_total_steps(self): + return self._total_steps + + def next_episode_results(self): + for i in range(self._num_returned, len(self._episode_rewards)): + yield (self._episode_rewards[i], self._episode_lengths[i]) + self._num_returned = len(self._episode_rewards) + + +@PublicAPI +class NoopResetEnv(gym.Wrapper): + def __init__(self, env, noop_max=30): + """Sample initial states by taking random number of no-ops on reset. + No-op is assumed to be action 0. + """ + gym.Wrapper.__init__(self, env) + self.noop_max = noop_max + self.override_num_noops = None + self.noop_action = 0 + assert env.unwrapped.get_action_meanings()[0] == "NOOP" + + def reset(self, **kwargs): + """Do no-op action for a number of steps in [1, noop_max].""" + self.env.reset(**kwargs) + if self.override_num_noops is not None: + noops = self.override_num_noops + else: + # This environment now uses the pcg64 random number generator which + # does not have `randint` as an attribute only has `integers`. + try: + noops = self.unwrapped.np_random.integers(1, self.noop_max + 1) + # Also still support older versions. + except AttributeError: + noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) + assert noops > 0 + obs = None + for _ in range(noops): + obs, _, terminated, truncated, info = self.env.step(self.noop_action) + if terminated or truncated: + obs, info = self.env.reset(**kwargs) + return obs, info + + def step(self, ac): + return self.env.step(ac) + + +@PublicAPI +class NormalizedImageEnv(gym.ObservationWrapper): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.observation_space = gym.spaces.Box( + -1.0, + 1.0, + shape=self.observation_space.shape, + dtype=np.float32, + ) + + # Divide by scale and center around 0.0, such that observations are in the range + # of -1.0 and 1.0. + def observation(self, observation): + return (observation.astype(np.float32) / 128.0) - 1.0 + + +@PublicAPI +class GrayScaleAndResize(gym.ObservationWrapper): + def __init__(self, env, dim, grayscale: bool = True): + """Warp frames to the specified size (dim x dim).""" + gym.ObservationWrapper.__init__(self, env) + self.width = dim + self.height = dim + self.grayscale = grayscale + self.observation_space = spaces.Box( + low=0, + high=255, + shape=(self.height, self.width, 1 if grayscale else 3), + dtype=np.uint8, + ) + + def observation(self, frame): + if self.grayscale: + frame = rgb2gray(frame) + frame = resize(frame, height=self.height, width=self.width) + return frame[:, :, None] + else: + return resize(frame, height=self.height, width=self.width) + + +WarpFrame = GrayScaleAndResize + + +@PublicAPI +def wrap_atari_for_new_api_stack( + env: gym.Env, + dim: int = 64, + frameskip: int = 4, + framestack: Optional[int] = None, + grayscale: bool = True, + # TODO (sven): Add option to NOT grayscale, in which case framestack must be None + # (b/c we are using the 3 color channels already as stacking frames). +) -> gym.Env: + """Wraps `env` for new-API-stack-friendly RLlib Atari experiments. + + Note that we assume reward clipping is done outside the wrapper. + + Args: + env: The env object to wrap. + dim: Dimension to resize observations to (dim x dim). + frameskip: Whether to skip n frames and max over them (keep brightest pixels). + framestack: Whether to stack the last n (grayscaled) frames. Note that this + step happens after(!) a possible frameskip step, meaning that if + frameskip=4 and framestack=2, we would perform the following over this + trajectory: + actual env timesteps: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 -> ... + frameskip: ( max ) ( max ) ( max ) ( max ) + framestack: ( stack ) (stack ) + + Returns: + The wrapped gym.Env. + """ + # Time limit. + env = gym.wrappers.TimeLimit(env, max_episode_steps=108000) + # Grayscale + resize. + env = WarpFrame(env, dim=dim, grayscale=grayscale) + # Normalize the image. + env = NormalizedImageEnv(env) + # Frameskip: Take max over these n frames. + if frameskip > 1: + assert env.spec is not None + env = MaxAndSkipEnv(env, skip=frameskip) + # Send n noop actions into env after reset to increase variance in the + # "start states" of the trajectories. These dummy steps are NOT included in the + # sampled data used for learning. + env = NoopResetEnv(env, noop_max=30) + # Each life is one episode. + env = EpisodicLifeEnv(env) + # Some envs only start playing after pressing fire. Unblock those. + if "FIRE" in env.unwrapped.get_action_meanings(): + env = FireResetEnv(env) + # Framestack. + if framestack: + env = FrameStack(env, k=framestack) + return env + + +@PublicAPI +def wrap_deepmind(env, dim=84, framestack=True, noframeskip=False): + """Configure environment for DeepMind-style Atari. + + Note that we assume reward clipping is done outside the wrapper. + + Args: + env: The env object to wrap. + dim: Dimension to resize observations to (dim x dim). + framestack: Whether to framestack observations. + """ + env = MonitorEnv(env) + env = NoopResetEnv(env, noop_max=30) + if env.spec is not None and noframeskip is True: + env = MaxAndSkipEnv(env, skip=4) + env = EpisodicLifeEnv(env) + if "FIRE" in env.unwrapped.get_action_meanings(): + env = FireResetEnv(env) + env = WarpFrame(env, dim) + # env = ClipRewardEnv(env) # reward clipping is handled by policy eval + # 4x image framestacking. + if framestack is True: + env = FrameStack(env, 4) + return env diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/dm_control_wrapper.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/dm_control_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..431e34a19037a2606edfa65a692c7f049f62f6e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/dm_control_wrapper.py @@ -0,0 +1,237 @@ +""" +DeepMind Control Suite Wrapper directly sourced from: +https://github.com/denisyarats/dmc2gym + +MIT License + +Copyright (c) 2020 Denis Yarats + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +""" +import gymnasium as gym + +try: + from dm_env import specs +except ImportError: + specs = None +try: + # Suppress MuJoCo warning (dm_control uses absl logging). + import absl.logging + + absl.logging.set_verbosity("error") + from dm_control import suite +except (ImportError, OSError): + suite = None +import numpy as np + +from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI + + +def _spec_to_box(spec): + def extract_min_max(s): + assert s.dtype == np.float64 or s.dtype == np.float32 + dim = np.int_(np.prod(s.shape)) + if type(s) is specs.Array: + bound = np.inf * np.ones(dim, dtype=np.float32) + return -bound, bound + elif type(s) is specs.BoundedArray: + zeros = np.zeros(dim, dtype=np.float32) + return s.minimum + zeros, s.maximum + zeros + + mins, maxs = [], [] + for s in spec: + mn, mx = extract_min_max(s) + mins.append(mn) + maxs.append(mx) + low = np.concatenate(mins, axis=0) + high = np.concatenate(maxs, axis=0) + assert low.shape == high.shape + return gym.spaces.Box(low, high, dtype=np.float32) + + +def _flatten_obs(obs): + obs_pieces = [] + for v in obs.values(): + flat = np.array([v]) if np.isscalar(v) else v.ravel() + obs_pieces.append(flat) + return np.concatenate(obs_pieces, axis=0) + + +@PublicAPI +class DMCEnv(gym.Env): + def __init__( + self, + domain_name, + task_name, + task_kwargs=None, + visualize_reward=False, + from_pixels=False, + height=64, + width=64, + camera_id=0, + frame_skip=2, + environment_kwargs=None, + channels_first=True, + preprocess=True, + ): + self._from_pixels = from_pixels + self._height = height + self._width = width + self._camera_id = camera_id + self._frame_skip = frame_skip + self._channels_first = channels_first + self.preprocess = preprocess + + if specs is None: + raise RuntimeError( + ( + "The `specs` module from `dm_env` was not imported. Make sure " + "`dm_env` is installed and visible in the current python " + "environment." + ) + ) + if suite is None: + raise RuntimeError( + ( + "The `suite` module from `dm_control` was not imported. Make " + "sure `dm_control` is installed and visible in the current " + "python enviornment." + ) + ) + + # create task + self._env = suite.load( + domain_name=domain_name, + task_name=task_name, + task_kwargs=task_kwargs, + visualize_reward=visualize_reward, + environment_kwargs=environment_kwargs, + ) + + # true and normalized action spaces + self._true_action_space = _spec_to_box([self._env.action_spec()]) + self._norm_action_space = gym.spaces.Box( + low=-1.0, high=1.0, shape=self._true_action_space.shape, dtype=np.float32 + ) + + # create observation space + if from_pixels: + shape = [3, height, width] if channels_first else [height, width, 3] + self._observation_space = gym.spaces.Box( + low=0, high=255, shape=shape, dtype=np.uint8 + ) + if preprocess: + self._observation_space = gym.spaces.Box( + low=-0.5, high=0.5, shape=shape, dtype=np.float32 + ) + else: + self._observation_space = _spec_to_box( + self._env.observation_spec().values() + ) + + self._state_space = _spec_to_box(self._env.observation_spec().values()) + + self.current_state = None + + def __getattr__(self, name): + return getattr(self._env, name) + + def _get_obs(self, time_step): + if self._from_pixels: + obs = self.render( + height=self._height, width=self._width, camera_id=self._camera_id + ) + if self._channels_first: + obs = obs.transpose(2, 0, 1).copy() + if self.preprocess: + obs = obs / 255.0 - 0.5 + else: + obs = _flatten_obs(time_step.observation) + return obs.astype(np.float32) + + def _convert_action(self, action): + action = action.astype(np.float64) + true_delta = self._true_action_space.high - self._true_action_space.low + norm_delta = self._norm_action_space.high - self._norm_action_space.low + action = (action - self._norm_action_space.low) / norm_delta + action = action * true_delta + self._true_action_space.low + action = action.astype(np.float32) + return action + + @property + def observation_space(self): + return self._observation_space + + @property + def state_space(self): + return self._state_space + + @property + def action_space(self): + return self._norm_action_space + + def step(self, action): + assert self._norm_action_space.contains(action) + action = self._convert_action(action) + assert self._true_action_space.contains(action) + reward = 0.0 + extra = {"internal_state": self._env.physics.get_state().copy()} + + terminated = truncated = False + for _ in range(self._frame_skip): + time_step = self._env.step(action) + reward += time_step.reward or 0.0 + terminated = False + truncated = time_step.last() + if terminated or truncated: + break + obs = self._get_obs(time_step) + self.current_state = _flatten_obs(time_step.observation) + extra["discount"] = time_step.discount + return obs, reward, terminated, truncated, extra + + def reset(self, *, seed=None, options=None): + time_step = self._env.reset() + self.current_state = _flatten_obs(time_step.observation) + obs = self._get_obs(time_step) + return obs, {} + + def render(self, mode="rgb_array", height=None, width=None, camera_id=0): + assert mode == "rgb_array", "only support for rgb_array mode" + height = height or self._height + width = width or self._width + camera_id = camera_id or self._camera_id + return self._env.physics.render(height=height, width=width, camera_id=camera_id) + + +@DeveloperAPI +class ActionClip(gym.ActionWrapper): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._low = -1.0 + self._high = 1.0 + self.action_space = gym.spaces.Box( + self._low, + self._high, + self.action_space.shape, + self.action_space.dtype, + ) + + def action(self, action): + return np.clip(action, self._low, self._high) diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/dm_env_wrapper.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/dm_env_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..7aef65848ee2c888d04374c66244cea1c2086a1e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/dm_env_wrapper.py @@ -0,0 +1,97 @@ +import gymnasium as gym +import numpy as np +from gymnasium import spaces + +try: + from dm_env import specs +except ImportError: + specs = None + +from ray.rllib.utils.annotations import PublicAPI + + +def _convert_spec_to_space(spec): + if isinstance(spec, dict): + return spaces.Dict({k: _convert_spec_to_space(v) for k, v in spec.items()}) + if isinstance(spec, specs.DiscreteArray): + return spaces.Discrete(spec.num_values) + elif isinstance(spec, specs.BoundedArray): + return spaces.Box( + low=np.asscalar(spec.minimum), + high=np.asscalar(spec.maximum), + shape=spec.shape, + dtype=spec.dtype, + ) + elif isinstance(spec, specs.Array): + return spaces.Box( + low=-float("inf"), high=float("inf"), shape=spec.shape, dtype=spec.dtype + ) + + raise NotImplementedError( + ( + "Could not convert `Array` spec of type {} to Gym space. " + "Attempted to convert: {}" + ).format(type(spec), spec) + ) + + +@PublicAPI +class DMEnv(gym.Env): + """A `gym.Env` wrapper for the `dm_env` API.""" + + metadata = {"render.modes": ["rgb_array"]} + + def __init__(self, dm_env): + super(DMEnv, self).__init__() + self._env = dm_env + self._prev_obs = None + + if specs is None: + raise RuntimeError( + ( + "The `specs` module from `dm_env` was not imported. Make sure " + "`dm_env` is installed and visible in the current python " + "environment." + ) + ) + + def step(self, action): + ts = self._env.step(action) + + reward = ts.reward + if reward is None: + reward = 0.0 + + return ts.observation, reward, ts.last(), False, {"discount": ts.discount} + + def reset(self, *, seed=None, options=None): + ts = self._env.reset() + return ts.observation, {} + + def render(self, mode="rgb_array"): + if self._prev_obs is None: + raise ValueError( + "Environment not started. Make sure to reset before rendering." + ) + + if mode == "rgb_array": + return self._prev_obs + else: + raise NotImplementedError("Render mode '{}' is not supported.".format(mode)) + + @property + def action_space(self): + spec = self._env.action_spec() + return _convert_spec_to_space(spec) + + @property + def observation_space(self): + spec = self._env.observation_spec() + return _convert_spec_to_space(spec) + + @property + def reward_range(self): + spec = self._env.reward_spec() + if isinstance(spec, specs.BoundedArray): + return spec.minimum, spec.maximum + return -float("inf"), float("inf") diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/group_agents_wrapper.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/group_agents_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..bb242709a136d9b0c42e4fd33437b32a7cf80ed2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/group_agents_wrapper.py @@ -0,0 +1,158 @@ +from collections import OrderedDict +from typing import Dict, List, Optional + +import gymnasium as gym + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.typing import AgentID + +# info key for the individual rewards of an agent, for example: +# info: { +# group_1: { +# _group_rewards: [5, -1, 1], # 3 agents in this group +# } +# } +GROUP_REWARDS = "_group_rewards" + +# info key for the individual infos of an agent, for example: +# info: { +# group_1: { +# _group_infos: [{"foo": ...}, {}], # 2 agents in this group +# } +# } +GROUP_INFO = "_group_info" + + +@DeveloperAPI +class GroupAgentsWrapper(MultiAgentEnv): + """Wraps a MultiAgentEnv environment with agents grouped as specified. + + See multi_agent_env.py for the specification of groups. + + This API is experimental. + """ + + def __init__( + self, + env: MultiAgentEnv, + groups: Dict[str, List[AgentID]], + obs_space: Optional[gym.Space] = None, + act_space: Optional[gym.Space] = None, + ): + """Wrap an existing MultiAgentEnv to group agent ID together. + + See `MultiAgentEnv.with_agent_groups()` for more detailed usage info. + + Args: + env: The env to wrap and whose agent IDs to group into new agents. + groups: Mapping from group id to a list of the agent ids + of group members. If an agent id is not present in any group + value, it will be left ungrouped. The group id becomes a new agent ID + in the final environment. + obs_space: Optional observation space for the grouped + env. Must be a tuple space. If not provided, will infer this to be a + Tuple of n individual agents spaces (n=num agents in a group). + act_space: Optional action space for the grouped env. + Must be a tuple space. If not provided, will infer this to be a Tuple + of n individual agents spaces (n=num agents in a group). + """ + super().__init__() + self.env = env + self.groups = groups + self.agent_id_to_group = {} + for group_id, agent_ids in groups.items(): + for agent_id in agent_ids: + if agent_id in self.agent_id_to_group: + raise ValueError( + "Agent id {} is in multiple groups".format(agent_id) + ) + self.agent_id_to_group[agent_id] = group_id + if obs_space is not None: + self.observation_space = obs_space + if act_space is not None: + self.action_space = act_space + for group_id in groups.keys(): + self._agent_ids.add(group_id) + + def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): + obs, info = self.env.reset(seed=seed, options=options) + + return ( + self._group_items(obs), + self._group_items( + info, + agg_fn=lambda gvals: {GROUP_INFO: list(gvals.values())}, + ), + ) + + def step(self, action_dict): + # Ungroup and send actions. + action_dict = self._ungroup_items(action_dict) + obs, rewards, terminateds, truncateds, infos = self.env.step(action_dict) + + # Apply grouping transforms to the env outputs + obs = self._group_items(obs) + rewards = self._group_items(rewards, agg_fn=lambda gvals: list(gvals.values())) + # Only if all of the agents are terminated, the group is terminated as well. + terminateds = self._group_items( + terminateds, agg_fn=lambda gvals: all(gvals.values()) + ) + # If all of the agents are truncated, the group is truncated as well. + truncateds = self._group_items( + truncateds, + agg_fn=lambda gvals: all(gvals.values()), + ) + infos = self._group_items( + infos, agg_fn=lambda gvals: {GROUP_INFO: list(gvals.values())} + ) + + # Aggregate rewards, but preserve the original values in infos. + for agent_id, rew in rewards.items(): + if isinstance(rew, list): + rewards[agent_id] = sum(rew) + if agent_id not in infos: + infos[agent_id] = {} + infos[agent_id][GROUP_REWARDS] = rew + + return obs, rewards, terminateds, truncateds, infos + + def _ungroup_items(self, items): + out = {} + for agent_id, value in items.items(): + if agent_id in self.groups: + assert len(value) == len(self.groups[agent_id]), ( + agent_id, + value, + self.groups, + ) + for a, v in zip(self.groups[agent_id], value): + out[a] = v + else: + out[agent_id] = value + return out + + def _group_items(self, items, agg_fn=None): + if agg_fn is None: + agg_fn = lambda gvals: list(gvals.values()) # noqa: E731 + + grouped_items = {} + for agent_id, item in items.items(): + if agent_id in self.agent_id_to_group: + group_id = self.agent_id_to_group[agent_id] + if group_id in grouped_items: + continue # already added + group_out = OrderedDict() + for a in self.groups[group_id]: + if a in items: + group_out[a] = items[a] + else: + raise ValueError( + "Missing member of group {}: {}: {}".format( + group_id, a, items + ) + ) + grouped_items[group_id] = agg_fn(group_out) + else: + grouped_items[agent_id] = item + return grouped_items diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/multi_agent_env_compatibility.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/multi_agent_env_compatibility.py new file mode 100644 index 0000000000000000000000000000000000000000..fc8efeda08346e5486f5a51b1fc69b614a7fffab --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/multi_agent_env_compatibility.py @@ -0,0 +1,73 @@ +from typing import Optional, Tuple + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.utils.typing import MultiAgentDict + + +class MultiAgentEnvCompatibility(MultiAgentEnv): + """A wrapper converting MultiAgentEnv from old gym API to the new one. + + "Old API" refers to step() method returning (observation, reward, done, info), + and reset() only retuning the observation. + "New API" refers to step() method returning (observation, reward, terminated, + truncated, info) and reset() returning (observation, info). + + Known limitations: + - Environments that use `self.np_random` might not work as expected. + """ + + def __init__(self, old_env, render_mode: Optional[str] = None): + """A wrapper which converts old-style envs to valid modern envs. + + Some information may be lost in the conversion, so we recommend updating your + environment. + + Args: + old_env: The old MultiAgentEnv to wrap. Implemented with the old API. + render_mode: The render mode to use when rendering the environment, + passed automatically to `env.render()`. + """ + super().__init__() + + self.metadata = getattr(old_env, "metadata", {"render_modes": []}) + self.render_mode = render_mode + self.reward_range = getattr(old_env, "reward_range", None) + self.spec = getattr(old_env, "spec", None) + self.env = old_env + + self.observation_space = old_env.observation_space + self.action_space = old_env.action_space + + def reset( + self, *, seed: Optional[int] = None, options: Optional[dict] = None + ) -> Tuple[MultiAgentDict, MultiAgentDict]: + # Use old `seed()` method. + if seed is not None: + self.env.seed(seed) + # Options are ignored + + if self.render_mode == "human": + self.render() + + obs = self.env.reset() + infos = {k: {} for k in obs.keys()} + return obs, infos + + def step( + self, action + ) -> Tuple[ + MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict + ]: + obs, rewards, terminateds, infos = self.env.step(action) + + # Truncated should always be False by default. + truncateds = {k: False for k in terminateds.keys()} + + return obs, rewards, terminateds, truncateds, infos + + def render(self): + # Use the old `render()` API, where we have to pass in the mode to each call. + return self.env.render(mode=self.render_mode) + + def close(self): + self.env.close() diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/open_spiel.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/open_spiel.py new file mode 100644 index 0000000000000000000000000000000000000000..3823230fb2d4ba899dd7d06a10a314121a40d8aa --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/open_spiel.py @@ -0,0 +1,130 @@ +from typing import Optional + +import gymnasium as gym +import numpy as np + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.utils import try_import_pyspiel + +pyspiel = try_import_pyspiel(error=True) + + +class OpenSpielEnv(MultiAgentEnv): + def __init__(self, env): + super().__init__() + self.env = env + self.agents = self.possible_agents = list(range(self.env.num_players())) + # Store the open-spiel game type. + self.type = self.env.get_type() + # Stores the current open-spiel game state. + self.state = None + + self.observation_space = gym.spaces.Dict( + { + aid: gym.spaces.Box( + float("-inf"), + float("inf"), + (self.env.observation_tensor_size(),), + dtype=np.float32, + ) + for aid in self.possible_agents + } + ) + self.action_space = gym.spaces.Dict( + { + aid: gym.spaces.Discrete(self.env.num_distinct_actions()) + for aid in self.possible_agents + } + ) + + def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): + self.state = self.env.new_initial_state() + return self._get_obs(), {} + + def step(self, action): + # Before applying action(s), there could be chance nodes. + # E.g. if env has to figure out, which agent's action should get + # resolved first in a simultaneous node. + self._solve_chance_nodes() + penalties = {} + + # Sequential game: + if str(self.type.dynamics) == "Dynamics.SEQUENTIAL": + curr_player = self.state.current_player() + assert curr_player in action + try: + self.state.apply_action(action[curr_player]) + # TODO: (sven) resolve this hack by publishing legal actions + # with each step. + except pyspiel.SpielError: + self.state.apply_action(np.random.choice(self.state.legal_actions())) + penalties[curr_player] = -0.1 + + # Compile rewards dict. + rewards = dict(enumerate(self.state.returns())) + # Simultaneous game. + else: + assert self.state.current_player() == -2 + # Apparently, this works, even if one or more actions are invalid. + self.state.apply_actions([action[ag] for ag in range(self.num_agents)]) + + # Now that we have applied all actions, get the next obs. + obs = self._get_obs() + + # Compile rewards dict and add the accumulated penalties + # (for taking invalid actions). + rewards = dict(enumerate(self.state.returns())) + for ag, penalty in penalties.items(): + rewards[ag] += penalty + + # Are we done? + is_terminated = self.state.is_terminal() + terminateds = dict( + {ag: is_terminated for ag in range(self.num_agents)}, + **{"__all__": is_terminated} + ) + truncateds = dict( + {ag: False for ag in range(self.num_agents)}, **{"__all__": False} + ) + + return obs, rewards, terminateds, truncateds, {} + + def render(self, mode=None) -> None: + if mode == "human": + print(self.state) + + def _get_obs(self): + # Before calculating an observation, there could be chance nodes + # (that may have an effect on the actual observations). + # E.g. After reset, figure out initial (random) positions of the + # agents. + self._solve_chance_nodes() + + if self.state.is_terminal(): + return {} + + # Sequential game: + if str(self.type.dynamics) == "Dynamics.SEQUENTIAL": + curr_player = self.state.current_player() + return { + curr_player: np.reshape(self.state.observation_tensor(), [-1]).astype( + np.float32 + ) + } + # Simultaneous game. + else: + assert self.state.current_player() == -2 + return { + ag: np.reshape(self.state.observation_tensor(ag), [-1]).astype( + np.float32 + ) + for ag in range(self.num_agents) + } + + def _solve_chance_nodes(self): + # Chance node(s): Sample a (non-player) action and apply. + while self.state.is_chance_node(): + assert self.state.current_player() == -1 + actions, probs = zip(*self.state.chance_outcomes()) + action = np.random.choice(actions, p=probs) + self.state.apply_action(action) diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/pettingzoo_env.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/pettingzoo_env.py new file mode 100644 index 0000000000000000000000000000000000000000..5f77668cd73330d9ebfe9be67a6abac6d2222325 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/pettingzoo_env.py @@ -0,0 +1,230 @@ +from typing import Optional + +import gymnasium as gym + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.utils.annotations import PublicAPI + + +@PublicAPI +class PettingZooEnv(MultiAgentEnv): + """An interface to the PettingZoo MARL environment library. + + See: https://github.com/Farama-Foundation/PettingZoo + + Inherits from MultiAgentEnv and exposes a given AEC + (actor-environment-cycle) game from the PettingZoo project via the + MultiAgentEnv public API. + + Note that the wrapper has the following important limitation: + + Environments are positive sum games (-> Agents are expected to cooperate + to maximize reward). This isn't a hard restriction, it just that + standard algorithms aren't expected to work well in highly competitive + games. + + Also note that the earlier existing restriction of all agents having the same + observation- and action spaces has been lifted. Different agents can now have + different spaces and the entire environment's e.g. `self.action_space` is a Dict + mapping agent IDs to individual agents' spaces. Same for `self.observation_space`. + + .. testcode:: + :skipif: True + + from pettingzoo.butterfly import prison_v3 + from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv + env = PettingZooEnv(prison_v3.env()) + obs, infos = env.reset() + # only returns the observation for the agent which should be stepping + print(obs) + + .. testoutput:: + + { + 'prisoner_0': array([[[0, 0, 0], + [0, 0, 0], + [0, 0, 0], + ..., + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]], dtype=uint8) + } + + .. testcode:: + :skipif: True + + obs, rewards, terminateds, truncateds, infos = env.step({ + "prisoner_0": 1 + }) + # only returns the observation, reward, info, etc, for + # the agent who's turn is next. + print(obs) + + .. testoutput:: + + { + 'prisoner_1': array([[[0, 0, 0], + [0, 0, 0], + [0, 0, 0], + ..., + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]], dtype=uint8) + } + + .. testcode:: + :skipif: True + + print(rewards) + + .. testoutput:: + + { + 'prisoner_1': 0 + } + + .. testcode:: + :skipif: True + + print(terminateds) + + .. testoutput:: + + { + 'prisoner_1': False, '__all__': False + } + + .. testcode:: + :skipif: True + + print(truncateds) + + .. testoutput:: + + { + 'prisoner_1': False, '__all__': False + } + + .. testcode:: + :skipif: True + + print(infos) + + .. testoutput:: + + { + 'prisoner_1': {'map_tuple': (1, 0)} + } + """ + + def __init__(self, env): + super().__init__() + self.env = env + env.reset() + + self._agent_ids = set(self.env.agents) + + # If these important attributes are not set, try to infer them. + if not self.agents: + self.agents = list(self._agent_ids) + if not self.possible_agents: + self.possible_agents = self.agents.copy() + + # Set these attributes for sampling in `VectorMultiAgentEnv`s. + self.observation_spaces = { + aid: self.env.observation_space(aid) for aid in self._agent_ids + } + self.action_spaces = { + aid: self.env.action_space(aid) for aid in self._agent_ids + } + + self.observation_space = gym.spaces.Dict(self.observation_spaces) + self.action_space = gym.spaces.Dict(self.action_spaces) + + def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): + info = self.env.reset(seed=seed, options=options) + return ( + {self.env.agent_selection: self.env.observe(self.env.agent_selection)}, + info or {}, + ) + + def step(self, action): + self.env.step(action[self.env.agent_selection]) + obs_d = {} + rew_d = {} + terminated_d = {} + truncated_d = {} + info_d = {} + while self.env.agents: + obs, rew, terminated, truncated, info = self.env.last() + agent_id = self.env.agent_selection + obs_d[agent_id] = obs + rew_d[agent_id] = rew + terminated_d[agent_id] = terminated + truncated_d[agent_id] = truncated + info_d[agent_id] = info + if ( + self.env.terminations[self.env.agent_selection] + or self.env.truncations[self.env.agent_selection] + ): + self.env.step(None) + else: + break + + all_gone = not self.env.agents + terminated_d["__all__"] = all_gone and all(terminated_d.values()) + truncated_d["__all__"] = all_gone and all(truncated_d.values()) + + return obs_d, rew_d, terminated_d, truncated_d, info_d + + def close(self): + self.env.close() + + def render(self): + return self.env.render(self.render_mode) + + @property + def get_sub_environments(self): + return self.env.unwrapped + + +@PublicAPI +class ParallelPettingZooEnv(MultiAgentEnv): + def __init__(self, env): + super().__init__() + self.par_env = env + self.par_env.reset() + self._agent_ids = set(self.par_env.agents) + + # If these important attributes are not set, try to infer them. + if not self.agents: + self.agents = list(self._agent_ids) + if not self.possible_agents: + self.possible_agents = self.agents.copy() + + self.observation_space = gym.spaces.Dict( + {aid: self.par_env.observation_space(aid) for aid in self._agent_ids} + ) + self.action_space = gym.spaces.Dict( + {aid: self.par_env.action_space(aid) for aid in self._agent_ids} + ) + + def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): + obs, info = self.par_env.reset(seed=seed, options=options) + return obs, info or {} + + def step(self, action_dict): + obss, rews, terminateds, truncateds, infos = self.par_env.step(action_dict) + terminateds["__all__"] = all(terminateds.values()) + truncateds["__all__"] = all(truncateds.values()) + return obss, rews, terminateds, truncateds, infos + + def close(self): + self.par_env.close() + + def render(self): + return self.par_env.render(self.render_mode) + + @property + def get_sub_environments(self): + return self.par_env.unwrapped diff --git a/lib/python3.12/site-packages/ray/rllib/env/wrappers/unity3d_env.py b/lib/python3.12/site-packages/ray/rllib/env/wrappers/unity3d_env.py new file mode 100644 index 0000000000000000000000000000000000000000..f22498e8bd96a07a2ef17a07ec50d95481ce6a34 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/env/wrappers/unity3d_env.py @@ -0,0 +1,319 @@ +import logging +import random +import time +from typing import Callable, Optional, Tuple + +import numpy as np +from gymnasium.spaces import Box, MultiDiscrete, Tuple as TupleSpace + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import AgentID, MultiAgentDict, PolicyID + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class Unity3DEnv(MultiAgentEnv): + # Default base port when connecting directly to the Editor + _BASE_PORT_EDITOR = 5004 + # Default base port when connecting to a compiled environment + _BASE_PORT_ENVIRONMENT = 5005 + # The worker_id for each environment instance + _WORKER_ID = 0 + + def __init__( + self, + file_name: str = None, + port: Optional[int] = None, + seed: int = 0, + no_graphics: bool = False, + timeout_wait: int = 300, + episode_horizon: int = 1000, + ): + super().__init__() + + if file_name is None: + print( + "No game binary provided, will use a running Unity editor " + "instead.\nMake sure you are pressing the Play (|>) button in " + "your editor to start." + ) + + import mlagents_envs + from mlagents_envs.environment import UnityEnvironment + + # Try connecting to the Unity3D game instance. If a port is blocked + port_ = None + while True: + # Sleep for random time to allow for concurrent startup of many + # environments (num_env_runners >> 1). Otherwise, would lead to port + # conflicts sometimes. + if port_ is not None: + time.sleep(random.randint(1, 10)) + port_ = port or ( + self._BASE_PORT_ENVIRONMENT if file_name else self._BASE_PORT_EDITOR + ) + # cache the worker_id and + # increase it for the next environment + worker_id_ = Unity3DEnv._WORKER_ID if file_name else 0 + Unity3DEnv._WORKER_ID += 1 + try: + self.unity_env = UnityEnvironment( + file_name=file_name, + worker_id=worker_id_, + base_port=port_, + seed=seed, + no_graphics=no_graphics, + timeout_wait=timeout_wait, + ) + print("Created UnityEnvironment for port {}".format(port_ + worker_id_)) + except mlagents_envs.exception.UnityWorkerInUseException: + pass + else: + break + + # ML-Agents API version. + self.api_version = self.unity_env.API_VERSION.split(".") + self.api_version = [int(s) for s in self.api_version] + + # Reset entire env every this number of step calls. + self.episode_horizon = episode_horizon + # Keep track of how many times we have called `step` so far. + self.episode_timesteps = 0 + + def step( + self, action_dict: MultiAgentDict + ) -> Tuple[ + MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict + ]: + from mlagents_envs.base_env import ActionTuple + + # Set only the required actions (from the DecisionSteps) in Unity3D. + all_agents = [] + for behavior_name in self.unity_env.behavior_specs: + # New ML-Agents API: Set all agents actions at the same time + # via an ActionTuple. Since API v1.4.0. + if self.api_version[0] > 1 or ( + self.api_version[0] == 1 and self.api_version[1] >= 4 + ): + actions = [] + for agent_id in self.unity_env.get_steps(behavior_name)[0].agent_id: + key = behavior_name + "_{}".format(agent_id) + all_agents.append(key) + actions.append(action_dict[key]) + if actions: + if actions[0].dtype == np.float32: + action_tuple = ActionTuple(continuous=np.array(actions)) + else: + action_tuple = ActionTuple(discrete=np.array(actions)) + self.unity_env.set_actions(behavior_name, action_tuple) + # Old behavior: Do not use an ActionTuple and set each agent's + # action individually. + else: + for agent_id in self.unity_env.get_steps(behavior_name)[ + 0 + ].agent_id_to_index.keys(): + key = behavior_name + "_{}".format(agent_id) + all_agents.append(key) + self.unity_env.set_action_for_agent( + behavior_name, agent_id, action_dict[key] + ) + # Do the step. + self.unity_env.step() + + obs, rewards, terminateds, truncateds, infos = self._get_step_results() + + # Global horizon reached? -> Return __all__ truncated=True, so user + # can reset. Set all agents' individual `truncated` to True as well. + self.episode_timesteps += 1 + if self.episode_timesteps > self.episode_horizon: + return ( + obs, + rewards, + terminateds, + dict({"__all__": True}, **{agent_id: True for agent_id in all_agents}), + infos, + ) + + return obs, rewards, terminateds, truncateds, infos + + def reset( + self, *, seed=None, options=None + ) -> Tuple[MultiAgentDict, MultiAgentDict]: + """Resets the entire Unity3D scene (a single multi-agent episode).""" + self.episode_timesteps = 0 + self.unity_env.reset() + obs, _, _, _, infos = self._get_step_results() + return obs, infos + + def _get_step_results(self): + obs = {} + rewards = {} + infos = {} + for behavior_name in self.unity_env.behavior_specs: + decision_steps, terminal_steps = self.unity_env.get_steps(behavior_name) + # Important: Only update those sub-envs that are currently + # available within _env_state. + # Loop through all envs ("agents") and fill in, whatever + # information we have. + for agent_id, idx in decision_steps.agent_id_to_index.items(): + key = behavior_name + "_{}".format(agent_id) + os = tuple(o[idx] for o in decision_steps.obs) + os = os[0] if len(os) == 1 else os + obs[key] = os + rewards[key] = ( + decision_steps.reward[idx] + decision_steps.group_reward[idx] + ) + for agent_id, idx in terminal_steps.agent_id_to_index.items(): + key = behavior_name + "_{}".format(agent_id) + # Only overwrite rewards (last reward in episode), b/c obs + # here is the last obs (which doesn't matter anyways). + # Unless key does not exist in obs. + if key not in obs: + os = tuple(o[idx] for o in terminal_steps.obs) + obs[key] = os = os[0] if len(os) == 1 else os + rewards[key] = ( + terminal_steps.reward[idx] + terminal_steps.group_reward[idx] + ) + + # Only use dones if all agents are done, then we should do a reset. + return obs, rewards, {"__all__": False}, {"__all__": False}, infos + + @staticmethod + def get_policy_configs_for_game( + game_name: str, + ) -> Tuple[dict, Callable[[AgentID], PolicyID]]: + + # The RLlib server must know about the Spaces that the Client will be + # using inside Unity3D, up-front. + obs_spaces = { + # 3DBall. + "3DBall": Box(float("-inf"), float("inf"), (8,)), + # 3DBallHard. + "3DBallHard": Box(float("-inf"), float("inf"), (45,)), + # GridFoodCollector + "GridFoodCollector": Box(float("-inf"), float("inf"), (40, 40, 6)), + # Pyramids. + "Pyramids": TupleSpace( + [ + Box(float("-inf"), float("inf"), (56,)), + Box(float("-inf"), float("inf"), (56,)), + Box(float("-inf"), float("inf"), (56,)), + Box(float("-inf"), float("inf"), (4,)), + ] + ), + # SoccerTwos. + "SoccerPlayer": TupleSpace( + [ + Box(-1.0, 1.0, (264,)), + Box(-1.0, 1.0, (72,)), + ] + ), + # SoccerStrikersVsGoalie. + "Goalie": Box(float("-inf"), float("inf"), (738,)), + "Striker": TupleSpace( + [ + Box(float("-inf"), float("inf"), (231,)), + Box(float("-inf"), float("inf"), (63,)), + ] + ), + # Sorter. + "Sorter": TupleSpace( + [ + Box( + float("-inf"), + float("inf"), + ( + 20, + 23, + ), + ), + Box(float("-inf"), float("inf"), (10,)), + Box(float("-inf"), float("inf"), (8,)), + ] + ), + # Tennis. + "Tennis": Box(float("-inf"), float("inf"), (27,)), + # VisualHallway. + "VisualHallway": Box(float("-inf"), float("inf"), (84, 84, 3)), + # Walker. + "Walker": Box(float("-inf"), float("inf"), (212,)), + # FoodCollector. + "FoodCollector": TupleSpace( + [ + Box(float("-inf"), float("inf"), (49,)), + Box(float("-inf"), float("inf"), (4,)), + ] + ), + } + action_spaces = { + # 3DBall. + "3DBall": Box(-1.0, 1.0, (2,), dtype=np.float32), + # 3DBallHard. + "3DBallHard": Box(-1.0, 1.0, (2,), dtype=np.float32), + # GridFoodCollector. + "GridFoodCollector": MultiDiscrete([3, 3, 3, 2]), + # Pyramids. + "Pyramids": MultiDiscrete([5]), + # SoccerStrikersVsGoalie. + "Goalie": MultiDiscrete([3, 3, 3]), + "Striker": MultiDiscrete([3, 3, 3]), + # SoccerTwos. + "SoccerPlayer": MultiDiscrete([3, 3, 3]), + # Sorter. + "Sorter": MultiDiscrete([3, 3, 3]), + # Tennis. + "Tennis": Box(-1.0, 1.0, (3,)), + # VisualHallway. + "VisualHallway": MultiDiscrete([5]), + # Walker. + "Walker": Box(-1.0, 1.0, (39,)), + # FoodCollector. + "FoodCollector": MultiDiscrete([3, 3, 3, 2]), + } + + # Policies (Unity: "behaviors") and agent-to-policy mapping fns. + if game_name == "SoccerStrikersVsGoalie": + policies = { + "Goalie": PolicySpec( + observation_space=obs_spaces["Goalie"], + action_space=action_spaces["Goalie"], + ), + "Striker": PolicySpec( + observation_space=obs_spaces["Striker"], + action_space=action_spaces["Striker"], + ), + } + + def policy_mapping_fn(agent_id, episode, worker, **kwargs): + return "Striker" if "Striker" in agent_id else "Goalie" + + elif game_name == "SoccerTwos": + policies = { + "PurplePlayer": PolicySpec( + observation_space=obs_spaces["SoccerPlayer"], + action_space=action_spaces["SoccerPlayer"], + ), + "BluePlayer": PolicySpec( + observation_space=obs_spaces["SoccerPlayer"], + action_space=action_spaces["SoccerPlayer"], + ), + } + + def policy_mapping_fn(agent_id, episode, worker, **kwargs): + return "BluePlayer" if "1_" in agent_id else "PurplePlayer" + + else: + policies = { + game_name: PolicySpec( + observation_space=obs_spaces[game_name], + action_space=action_spaces[game_name], + ), + } + + def policy_mapping_fn(agent_id, episode, worker, **kwargs): + return game_name + + return policies, policy_mapping_fn diff --git a/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc89ab128b3010e131e91d078000da4593d700e0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/centralized_critic.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/centralized_critic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bfeeb5caf1cbd43f5d1f0650bcefeaa53dce6ed Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/centralized_critic.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/compute_adapted_gae_on_postprocess_trajectory.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/compute_adapted_gae_on_postprocess_trajectory.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6238b66158dfa991f2b2b02663ca679ef74fe5b8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/compute_adapted_gae_on_postprocess_trajectory.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/quadx_waypoints.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/quadx_waypoints.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff87b7a8cc6b0dd36ee1e274027948a728bbc95c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/quadx_waypoints.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/replay_buffer_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/replay_buffer_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d73898735a88e55d934bfef958d9dc950431be5a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/__pycache__/replay_buffer_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/actions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c97ee0f408309bf615afc7b426a2c057ef0e185e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/autoregressive_actions.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/autoregressive_actions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b5f170511eb8c61e4fc1fbcd40812c700cf424b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/autoregressive_actions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/custom_action_distribution.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/custom_action_distribution.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd95744b3bcfa70c0d8ba35412bd9a1426f1599f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/custom_action_distribution.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/nested_action_spaces.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/nested_action_spaces.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ac56e771285055bf023e956ab741f4bf5ff374a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/actions/__pycache__/nested_action_spaces.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/autoregressive_actions.py b/lib/python3.12/site-packages/ray/rllib/examples/actions/autoregressive_actions.py new file mode 100644 index 0000000000000000000000000000000000000000..370223c593b13e8b0feade804123e141a35f6c3d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/actions/autoregressive_actions.py @@ -0,0 +1,107 @@ +"""Example on how to define and run with an RLModule with a dependent action space. + +This examples: + - Shows how to write a custom RLModule outputting autoregressive actions. + The RLModule class used here implements a prior distribution for the first couple + of actions and then uses the sampled actions to compute the parameters for and + sample from a posterior distribution. + - Shows how to configure a PPO algorithm to use the custom RLModule. + - Stops the training after 100k steps or when the mean episode return + exceeds -0.012 in evaluation, i.e. if the agent has learned to + synchronize its actions. + +For details on the environment used, take a look at the `CorrelatedActionsEnv` +class. To receive an episode return over 100, the agent must learn how to synchronize +its actions. + + +How to run this script +---------------------- +`python [script file name].py --num-env-runners 2` + +Control the number of `EnvRunner`s with the `--num-env-runners` flag. This +will increase the sampling speed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should reach an episode return of better than -0.5 quickly through a simple PPO +policy. The logic behind beating the env is roughly: + +OBS: optimal a1: r1: optimal a2: r2: +-1 2 0 -1.0 0 +-0.5 1/2 -0.5 -0.5/-1.5 0 +0 1 0 -1.0 0 +0.5 0/1 -0.5 -0.5/-1.5 0 +1 0 0 -1.0 0 + +Meaning, most of the time, you would receive a reward better than -0.5, but worse than +0.0. + ++--------------------------------------+------------+--------+------------------+ +| Trial name | status | iter | total time (s) | +| | | | | +|--------------------------------------+------------+--------+------------------+ +| PPO_CorrelatedActionsEnv_6660d_00000 | TERMINATED | 76 | 132.438 | ++--------------------------------------+------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| episode_return_mean | num_env_steps_sample | ...env_steps_sampled | +| | d_lifetime | _lifetime_throughput | +|------------------------+------------------------+------------------------| +| -0.43 | 152000 | 1283.48 | ++------------------------+------------------------+------------------------+ +""" + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.envs.classes.correlated_actions_env import CorrelatedActionsEnv +from ray.rllib.examples.rl_modules.classes.autoregressive_actions_rlm import ( + AutoregressiveActionsRLM, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_iters=1000, + default_timesteps=2000000, + default_reward=-0.45, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + if args.algo != "PPO": + raise ValueError( + "This example script only runs with PPO! Set --algo=PPO on the command " + "line." + ) + + base_config = ( + PPOConfig() + .environment(CorrelatedActionsEnv) + .training( + train_batch_size_per_learner=2000, + num_epochs=12, + minibatch_size=256, + entropy_coeff=0.005, + lr=0.0003, + ) + # Specify the RLModule class to be used. + .rl_module( + rl_module_spec=RLModuleSpec(module_class=AutoregressiveActionsRLM), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/custom_action_distribution.py b/lib/python3.12/site-packages/ray/rllib/examples/actions/custom_action_distribution.py new file mode 100644 index 0000000000000000000000000000000000000000..6f96a8cd49e0128ac0eda2127e841adeb818d9c2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/actions/custom_action_distribution.py @@ -0,0 +1,116 @@ +"""Example on how to define and run an experiment with a custom action distribution. + +The example uses an additional `temperature` parameter on top of the built-in +`TorchCategorical` class. Incoming logits (outputs from the RLModule) are divided by +this temperature before creating the underlying +torch.distributions.categorical.Categorical object. + +This examples: + - Shows how to write a custom RLlib action distribution class accepting an + additional parameter in its constructor. + - demonstrates how you can subclass the TorchRLModule base class and write your + own architecture by overriding the `setup()` method. + - shows how to set the attribute `self.action_dist_cls` in that same `setup()` + method. For an alternative way of defining action distribution classes for your + RLModules, see the `setup()` method implementation in the imported + `CustomActionDistributionRLModule` class. + - shows how you then configure an RLlib Algorithm such that it uses your custom + RLModule (instead of a default RLModule). + + +How to run this script +---------------------- +`python [script file name].py --temperature=0.8` + +Use the `--temperature` setting to change the temperature. Higher values (>>1.0) lead +to almost random behavior, lower values (<<1.0) lead to always-greedy behavior. Note +though, that both extremes hurt learning performance. + +Control the number of `EnvRunner`s with the `--num-env-runners` flag. This +will increase the sampling speed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +With a --temperature setting of 0.75, learning seems to be particularly easy with the +given other parameters: + ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_1bbe0_00000 | TERMINATED | 127.0.0.1:81594 | 22 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | episode_return_mean | num_env_steps_sample | +| | | d_lifetime | +|------------------+------------------------+------------------------| +| 17.6368 | 450.54 | 88000 | ++------------------+------------------------+------------------------+ +""" + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.rl_modules.classes.custom_action_distribution_rlm import ( + CustomActionDistributionRLModule, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_timesteps=200000, + default_reward=450.0, +) +parser.add_argument( + "--temperature", + type=float, + default=2.0, + help="The action distribution temperature to apply to the raw model logits. " + "Logits are first divided by the temperature, then an underlying torch.Categorical " + "distribution is created from those altered logits and used for sampling actions. " + "Set this to <<1.0 to approximate greedy behavior and to >>1.0 to approximate " + "random behavior.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + if args.algo != "PPO": + raise ValueError( + "This example script only runs with PPO! Set --algo=PPO on the command " + "line." + ) + + base_config = ( + PPOConfig() + .environment("CartPole-v1") + .training( + lr=0.0003, + num_epochs=6, + vf_loss_coeff=0.01, + ) + # Specify the RLModule class to be used. + .rl_module( + rl_module_spec=RLModuleSpec( + module_class=CustomActionDistributionRLModule, + model_config={ + "hidden_dim": 128, + "action_dist_temperature": args.temperature, + }, + ), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/actions/nested_action_spaces.py b/lib/python3.12/site-packages/ray/rllib/examples/actions/nested_action_spaces.py new file mode 100644 index 0000000000000000000000000000000000000000..821c57ff682890c3060dd28b2b0348bc4b2421ee --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/actions/nested_action_spaces.py @@ -0,0 +1,83 @@ +from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple + +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.examples.envs.classes.multi_agent import ( + MultiAgentNestedSpaceRepeatAfterMeEnv, +) +from ray.rllib.examples.envs.classes.nested_space_repeat_after_me_env import ( + NestedSpaceRepeatAfterMeEnv, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +# Read in common example script command line arguments. +parser = add_rllib_example_script_args(default_timesteps=200000, default_reward=-500.0) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Define env-to-module-connector pipeline for the new stack. + def _env_to_module_pipeline(env, spaces, device): + return FlattenObservations(multi_agent=args.num_agents > 0) + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "env", + lambda c: MultiAgentNestedSpaceRepeatAfterMeEnv( + config=dict(c, **{"num_agents": args.num_agents}) + ), + ) + else: + register_env("env", lambda c: NestedSpaceRepeatAfterMeEnv(c)) + + # Define the AlgorithmConfig used. + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + "env", + env_config={ + "space": Dict( + { + "a": Tuple( + [Dict({"d": Box(-15.0, 3.0, ()), "e": Discrete(3)})] + ), + "b": Box(-10.0, 10.0, (2,)), + "c": MultiDiscrete([3, 3]), + "d": Discrete(2), + } + ), + "episode_len": 100, + }, + ) + .env_runners(env_to_module_connector=_env_to_module_pipeline) + # No history in Env (bandit problem). + .training( + gamma=0.0, + lr=0.0005, + ) + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + # Fix some PPO-specific settings. + if args.algo == "PPO": + base_config.training( + # We don't want high entropy in this Env. + entropy_coeff=0.00005, + num_epochs=4, + vf_loss_coeff=0.01, + ) + + # Run everything as configured. + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..022f12bb1274d8363277685bff63f652ea892c24 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/appo_custom_algorithm_w_shared_data_actor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/appo_custom_algorithm_w_shared_data_actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84f32ddfd0dfcf96a683c97f3e44a5841958aee4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/appo_custom_algorithm_w_shared_data_actor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/maml_lr_supervised_learning.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/maml_lr_supervised_learning.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0083fc2f3605d9c493abb4148c69507d76911a29 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/maml_lr_supervised_learning.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/vpg_custom_algorithm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/vpg_custom_algorithm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7f7dd7cd6e45a2eaf0950a2bed0b8be16f3e7dc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/__pycache__/vpg_custom_algorithm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/appo_custom_algorithm_w_shared_data_actor.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/appo_custom_algorithm_w_shared_data_actor.py new file mode 100644 index 0000000000000000000000000000000000000000..ecb6740700b2b3cfeaee2bc7f19ae88b3ba737cc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/appo_custom_algorithm_w_shared_data_actor.py @@ -0,0 +1,183 @@ +"""Example of how to write a custom APPO that uses a global shared data actor. + +The actor is custom code and its remote APIs can be designed as the user requires. +It is created inside the Algorithm's `setup` method and then shared through its +reference with all of the Algorithm's other actors, like EnvRunners, Learners, and +aggregator actors. + +During sampling and through using callbacks, each EnvRunner assigns a unique ID +to each sampled episode chunk, then sends manipulated reward data for each sampled +episode chunk to the shared data actor. In particular, the manipulation consists of +each individual reward being multiplied by the EnvRunner's index (from 1 to ...). +Note that the actual reward in the episode is not altered and thus the metrics +reporting continues to show the original reward. + +In the learner connector, which creates the train batch from episode data, a custom +connector piece then gets the manipulated rewards from the shared data actor using +the episode chunk's unique ID (see above) and uses the manipulated reward for training. +Note that because of this, different EnvRunners provide different reward signals, which +should make it slightly harder for the value function to learn consistently. +Nevertheless, because the default config here only uses 2 EnvRunners, each multiplying +their rewards by 1 and 2, respectively, this effect is negligible here and the example +should learn how to solve the CartPole-1 env either way. + +This example shows: + + - how to write a custom, global shared data actor class with a custom remote API. + - how an instance of this shared data actor is created upon algorithm + initialization. + - how to distribute the actor reference of the shared actor to all other actors + in the Algorithm, for example EnvRunners, AggregatorActors, and Learners + - how to subclass an existing algorithm class (APPO) to implement a custom + Algorithm, overriding the `setup` method to control, which additional actors + should be created (and shared) by the algo, the `get_state/set_state` methods + to include the state of the new actor. + - how - through custom callbacks - the new actor can be written to and queried + from anywhere within the algorithm, for example its EnvRunner actors or Learners. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +The experiment should work regardless of whether you are using aggregator +actors or not. By default, the experiment provides one agg. actor per Learner, +but you can set `--num-aggregator-actors-per-learner=0` to have the learner +connector pipeline work directly inside the Learner actor(s). + ++-------------------------------------------------+------------+--------+ +| Trial name | status | iter | +| | | | +|-------------------------------------------------+------------+--------+ +| APPOWithSharedDataActor_CartPole-v1_4e860_00000 | TERMINATED | 7 | ++-------------------------------------------------+------------+--------+ ++------------------+------------------------+ +| total time (s) | episode_return_mean | +| | | +|------------------+------------------------+ +| 70.0315 | 468.42 | ++------------------+------------------------+ +""" + +import uuid + +import ray +from ray.rllib.algorithms.appo import APPOConfig +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core import Columns +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.examples.algorithms.classes.appo_w_shared_data_actor import ( + APPOWithSharedDataActor, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_reward=450.0, + default_iters=200, + default_timesteps=2000000, +) +parser.set_defaults( + num_aggregator_actors_per_learner=1, +) + +SPECIAL_REWARDS_KEY = "special_(double)_rewards" +ENV_RUNNER_IDX_KEY = "env_runner_index" +UNIQUE_EPISODE_CHUNK_KEY = "unique_eps_chunk" + + +# Define 2 simple EnvRunner-based callbacks: + + +def on_episode_step(*, episode, env_runner, **kwargs): + # Multiplies the received reward by the env runner index. + if SPECIAL_REWARDS_KEY not in episode.custom_data: + episode.custom_data[SPECIAL_REWARDS_KEY] = [] + episode.custom_data[SPECIAL_REWARDS_KEY].append( + episode.get_rewards(-1) * env_runner.worker_index + ) + + +def on_sample_end(*, samples, env_runner, **kwargs): + # Sends the (manipulated) reward sequence to the shared data actor for "pickup" by + # a Learner. Alternatively, one could also just store the information in the + # `custom_data` property. + for episode in samples: + # Provide a unique key for both episode AND record in the shared + # data actor. + unique_key = str(uuid.uuid4()) + + # Store the EnvRunner index and unique key in the episode. + episode.custom_data[ENV_RUNNER_IDX_KEY] = env_runner.worker_index + episode.custom_data[UNIQUE_EPISODE_CHUNK_KEY] = unique_key + + # Get the manipulated rewards from the episode .. + special_rewards = episode.custom_data.pop(SPECIAL_REWARDS_KEY) + # .. and send them under the unique key to the shared data actor. + env_runner._shared_data_actor.put.remote( + key=unique_key, + value=special_rewards, + ) + + +class ManipulatedRewardConnector(ConnectorV2): + def __call__(self, *, episodes, batch, metrics, **kwargs): + if not isinstance(episodes[0], SingleAgentEpisode): + raise ValueError("This connector only works on `SingleAgentEpisodes`.") + # Get the manipulated rewards from the shared actor and add them to the train + # batch. + for sa_episode in self.single_agent_episode_iterator(episodes): + unique_key = sa_episode.custom_data[UNIQUE_EPISODE_CHUNK_KEY] + special_rewards = ray.get( + self._shared_data_actor.get.remote(unique_key, delete=True) + ) + if special_rewards is None: + continue + + assert int(special_rewards[0]) == sa_episode.custom_data[ENV_RUNNER_IDX_KEY] + + # Add one more fake reward, b/c all episodes will be extended + # (in PPO-style algos) by one artificial timestep for GAE/v-trace + # computation purposes. + special_rewards += [0.0] + self.add_n_batch_items( + batch=batch, + column=Columns.REWARDS, + items_to_add=special_rewards[-len(sa_episode) :], + num_items=len(sa_episode), + single_agent_episode=sa_episode, + ) + + return batch + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + APPOConfig(algo_class=APPOWithSharedDataActor) + .environment("CartPole-v1") + .callbacks( + on_episode_step=on_episode_step, + on_sample_end=on_sample_end, + ) + .training( + learner_connector=(lambda obs_sp, act_sp: ManipulatedRewardConnector()), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e17450a366e1984b8d5345db1d8a9dc2e66f644 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/appo_w_shared_data_actor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/appo_w_shared_data_actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ecd761f061617167a4f7311a75ba9a425307b079 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/appo_w_shared_data_actor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_differentiable_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_differentiable_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4072a8cd6bfb471cab2d010d7d2b50999df9e343 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_differentiable_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_differentiable_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_differentiable_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d351a3e464867bf906f4788984c7f423c86b0d4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_differentiable_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_meta_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_meta_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f93b260d964496c413034db66eeb1d2804a6071 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/maml_lr_meta_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/vpg.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/vpg.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d0c5ea3c2844f920343e61eb24bf7830e41d171 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/__pycache__/vpg.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/appo_w_shared_data_actor.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/appo_w_shared_data_actor.py new file mode 100644 index 0000000000000000000000000000000000000000..1b2a1efef43854aac5e51172eabbfad817aeddaa --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/appo_w_shared_data_actor.py @@ -0,0 +1,88 @@ +from typing import List + +import ray +from ray.rllib.algorithms import AlgorithmConfig +from ray.rllib.algorithms.appo import APPO +from ray.rllib.env.env_runner_group import EnvRunnerGroup + + +@ray.remote +class SharedDataActor: + """Simple example of an actor that's accessible from all other actors of an algo. + + Exposes remote APIs `put` and `get` to other actors for storing and retrieving + arbitrary data. + """ + + def __init__(self): + self.storage = {} + + def get(self, key, delete: bool = False): + value = self.storage.get(key) + if delete and key in self.storage: + del self.storage[key] + return value + + def put(self, key, value): + self.storage[key] = value + + def get_state(self): + return self.storage + + def set_state(self, state): + self.storage = state + + +class APPOWithSharedDataActor(APPO): + def setup(self, config: AlgorithmConfig): + # Call to parent `setup`. + super().setup(config) + + # Create shared data actor. + self.shared_data_actor = SharedDataActor.remote() + + # Share the actor with all other relevant actors. + + def _share(actor, shared_act=self.shared_data_actor): + actor._shared_data_actor = shared_act + # Also add shared actor reference to all the learner connector pieces, + # if applicable. + if hasattr(actor, "_learner_connector") and actor._learner_connector: + for conn in actor._learner_connector: + conn._shared_data_actor = shared_act + + self.env_runner_group.foreach_env_runner(func=_share) + if self.eval_env_runner_group: + self.eval_env_runner_group.foreach_env_runner(func=_share) + self.learner_group.foreach_learner(func=_share) + if self._aggregator_actor_manager: + self._aggregator_actor_manager.foreach_actor(func=_share) + + def get_state(self, *args, **kwargs): + state = super().get_state(*args, **kwargs) + # Add shared actor's state. + state["shared_data_actor"] = ray.get(self.shared_data_actor.get_state.remote()) + return state + + def set_state(self, state, *args, **kwargs): + super().set_state(state, *args, **kwargs) + # Set shared actor's state. + if "shared_data_actor" in state: + self.shared_data_actor.set_state.remote(state["shared_data_actor"]) + + def restore_env_runners(self, env_runner_group: EnvRunnerGroup) -> List[int]: + restored = super().restore_env_runners(env_runner_group) + + # For the restored EnvRunners, send them the latest shared, global state + # from the `SharedDataActor`. + for restored_idx in restored: + state_ref = self.shared_data_actor.get.remote( + key=f"EnvRunner_{restored_idx}" + ) + env_runner_group.foreach_env_runner( + lambda env_runner, state=state_ref: env_runner._global_state, + remote_worker_ids=[restored_idx], + timeout_seconds=0.0, + ) + + return restored diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_differentiable_learner.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_differentiable_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..f19d755e7df92d0538abd632a45ba261a99cb4a6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_differentiable_learner.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING, Any, Dict + +from ray.rllib.core.learner.torch.torch_differentiable_learner import ( + TorchDifferentiableLearner, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModuleID, TensorType + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +torch, nn = try_import_torch() + + +class MAMLTorchDifferentiableLearner(TorchDifferentiableLearner): + """A `TorchDifferentiableLearner` to perform MAML learning. + + This `TorchDifferentiableLearner` + - defines a funcitonal MSE loss for learning simple (here non-linear) + prediction. + """ + + @override(TorchDifferentiableLearner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + """Defines a simple MSE prediction loss for continuous task.""" + + return nn.functional.mse_loss(fwd_out["y_pred"], batch["y"]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_differentiable_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_differentiable_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..4fa91895f0ded2da96e3abe846b89361b473c49d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_differentiable_rlm.py @@ -0,0 +1,39 @@ +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class DifferentiableTorchRLModule(TorchRLModule): + """Differentiable neural network to learn sinusoid curves. + + This `TorchRLModule`: + - defines a simple neural network to learn sinusoid curves with two + feed forward layern and ReLU activations, + - defines a differentiable `forward` call by overriding the `_forward` + method (which is implicitly used by the module's `forward` method); this + enables `torch.func.functional_call?` to work. + """ + + def setup(self): + """Sets up a simple neural network + + The network contains two hidden layers and ReLU activations. Note, + input and output are single dimensional b/c the sinusoid curve is. + """ + self.net = nn.Sequential( + nn.Linear(1, 40), nn.ReLU(), nn.Linear(40, 40), nn.ReLU(), nn.Linear(40, 1) + ) + + def _forward(self, batch, **kwargs): + """Defines method to be called for general forward path. + + Note, it is important that the `RLModule.forward` method contains the + logic to be used for training forward pass b/c otherwise the functional + call via `torch.func.functional_call` will not work. See for reference + https://pytorch.org/docs/stable/generated/torch.func.functional_call.html. + """ + outs = {} + outs["y_pred"] = self.net(batch[Columns.OBS]) + return outs diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_meta_learner.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_meta_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..35fae46e9397a6f288845628b2fc55d067222b0e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/maml_lr_meta_learner.py @@ -0,0 +1,38 @@ +from typing import TYPE_CHECKING, Any, Dict, List + +from ray.rllib.core.learner.torch.torch_meta_learner import TorchMetaLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModuleID, TensorType + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +torch, nn = try_import_torch() + + +class MAMLTorchMetaLearner(TorchMetaLearner): + """A `TorchMetaLearner` to perform MAML learning. + + This `TorchMetaLearner` + - defines a MSE loss for learning simple (here non-linear) prediction. + """ + + @override(TorchMetaLearner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + others_loss_per_module: List[Dict[ModuleID, TensorType]] = None, + ) -> TensorType: + """Defines a simple MSE prediction loss for continuous task. + + Note, MAML does not need the losses from the registered differentiable + learners (contained in `others_loss_per_module`) b/c it computes a test + loss on an unseen data batch. + """ + # Use a simple MSE loss for the meta learning task. + return torch.nn.functional.mse_loss(fwd_out["y_pred"], batch["y"]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/vpg.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/vpg.py new file mode 100644 index 0000000000000000000000000000000000000000..1eef81c7794339c112db7d7040104241df0ddcde --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/classes/vpg.py @@ -0,0 +1,175 @@ +import tree # pip install dm_tree +from typing_extensions import Self + +from ray.rllib.algorithms import Algorithm +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + ENV_RUNNER_SAMPLING_TIMER, + LEARNER_RESULTS, + LEARNER_UPDATE_TIMER, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + SYNCH_WORKER_WEIGHTS_TIMER, + TIMERS, +) + + +class VPGConfig(AlgorithmConfig): + """A simple VPG (vanilla policy gradient) algorithm w/o value function support. + + Use for testing purposes only! + + This Algorithm should use the VPGTorchLearner and VPGTorchRLModule + """ + + # A test setting to activate metrics on mean weights. + report_mean_weights: bool = True + + def __init__(self, algo_class=None): + super().__init__(algo_class=algo_class or VPG) + + # VPG specific settings. + self.num_episodes_per_train_batch = 10 + # Note that we don't have to set this here, because we tell the EnvRunners + # explicitly to sample entire episodes. However, for good measure, we change + # this setting here either way. + self.batch_mode = "complete_episodes" + + # VPG specific defaults (from AlgorithmConfig). + self.num_env_runners = 1 + + @override(AlgorithmConfig) + def training(self, *, num_episodes_per_train_batch=NotProvided, **kwargs) -> Self: + """Sets the training related configuration. + + Args: + num_episodes_per_train_batch: The number of complete episodes per train + batch. VPG requires entire episodes to be sampled from the EnvRunners. + For environments with varying episode lengths, this leads to varying + batch sizes (in timesteps) as well possibly causing slight learning + instabilities. However, for simplicity reasons, we stick to collecting + always exactly n episodes per training update. + + Returns: + This updated AlgorithmConfig object. + """ + # Pass kwargs onto super's `training()` method. + super().training(**kwargs) + + if num_episodes_per_train_batch is not NotProvided: + self.num_episodes_per_train_batch = num_episodes_per_train_batch + + return self + + @override(AlgorithmConfig) + def get_default_rl_module_spec(self): + if self.framework_str == "torch": + from ray.rllib.examples.rl_modules.classes.vpg_torch_rlm import ( + VPGTorchRLModule, + ) + + spec = RLModuleSpec( + module_class=VPGTorchRLModule, + model_config={"hidden_dim": 64}, + ) + else: + raise ValueError(f"Unsupported framework: {self.framework_str}") + + return spec + + @override(AlgorithmConfig) + def get_default_learner_class(self): + if self.framework_str == "torch": + from ray.rllib.examples.learners.classes.vpg_torch_learner import ( + VPGTorchLearner, + ) + + return VPGTorchLearner + else: + raise ValueError(f"Unsupported framework: {self.framework_str}") + + +class VPG(Algorithm): + @classmethod + @override(Algorithm) + def get_default_config(cls) -> VPGConfig: + return VPGConfig() + + @override(Algorithm) + def training_step(self) -> None: + """Override of the training_step method of `Algorithm`. + + Runs the following steps per call: + - Sample B timesteps (B=train batch size). Note that we don't sample complete + episodes due to simplicity. For an actual VPG algo, due to the loss computation, + you should always sample only completed episodes. + - Send the collected episodes to the VPG LearnerGroup for model updating. + - Sync the weights from LearnerGroup to all EnvRunners. + """ + # Sample. + with self.metrics.log_time((TIMERS, ENV_RUNNER_SAMPLING_TIMER)): + episodes, env_runner_results = self._sample_episodes() + # Merge results from n parallel sample calls into self's metrics logger. + self.metrics.aggregate(env_runner_results, key=ENV_RUNNER_RESULTS) + + # Just for demonstration purposes, log the number of time steps sampled in this + # `training_step` round. + # Mean over a window of 100: + self.metrics.log_value( + "episode_timesteps_sampled_mean_win100", + sum(map(len, episodes)), + reduce="mean", + window=100, + ) + # Exponential Moving Average (EMA) with coeff=0.1: + self.metrics.log_value( + "episode_timesteps_sampled_ema", + sum(map(len, episodes)), + ema_coeff=0.1, # <- weight of new value; weight of old avg=1.0-ema_coeff + ) + + # Update model. + with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)): + learner_results = self.learner_group.update( + episodes=episodes, + timesteps={ + NUM_ENV_STEPS_SAMPLED_LIFETIME: ( + self.metrics.peek( + (ENV_RUNNER_RESULTS, NUM_ENV_STEPS_SAMPLED_LIFETIME) + ) + ), + }, + ) + # Merge results from m parallel update calls into self's metrics logger. + self.metrics.aggregate(learner_results, key=LEARNER_RESULTS) + + # Sync weights. + with self.metrics.log_time((TIMERS, SYNCH_WORKER_WEIGHTS_TIMER)): + self.env_runner_group.sync_weights( + from_worker_or_learner_group=self.learner_group, + inference_only=True, + ) + + def _sample_episodes(self): + # How many episodes to sample from each EnvRunner? + num_episodes_per_env_runner = self.config.num_episodes_per_train_batch // ( + self.config.num_env_runners or 1 + ) + # Send parallel remote requests to sample and get the metrics. + sampled_data = self.env_runner_group.foreach_env_runner( + # Return tuple of [episodes], [metrics] from each EnvRunner. + lambda env_runner: ( + env_runner.sample(num_episodes=num_episodes_per_env_runner), + env_runner.get_metrics(), + ), + # Loop over remote EnvRunners' `sample()` method in parallel or use the + # local EnvRunner if there aren't any remote ones. + local_env_runner=self.env_runner_group.num_remote_workers() <= 0, + ) + # Return one list of episodes and a list of metrics dicts (one per EnvRunner). + episodes = tree.flatten([s[0] for s in sampled_data]) + stats_dicts = [s[1] for s in sampled_data] + + return episodes, stats_dicts diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/maml_lr_supervised_learning.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/maml_lr_supervised_learning.py new file mode 100644 index 0000000000000000000000000000000000000000..f519d8c14672d943b23aaee5672e0aa202a6712c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/maml_lr_supervised_learning.py @@ -0,0 +1,443 @@ +"""Example of how to use `TorchMetaLearner` and `DifferentiableLearner` for MAML. + +Meta-learning, or “learning to learn,” trains models to quickly adapt to new tasks +using only a few examples. One prominent method is Model-Agnostic Meta-Learning +(MAML), which is compatible with any model trained via gradient descent. MAML has +been successfully applied across domains such as classification, regression, and +reinforcement learning. + +In this MAML example, the goal is to train a model that can adapt to an infinite +number of tasks, where each task corresponds to a sinusoidal function with randomly +sampled amplitude and phase. Because each new task introduces a shift in data +distribution, traditional learning algorithms would fail to generalize — they’d +overfit to the training task and struggle on unseen ones. Meta-learning addresses +this by optimizing the model parameters such that they can be fine-tuned rapidly +for any new task. + +During training, a DifferentiableLearner performs an inner-loop update using the +training error for each task. The outer-loop TorchMetaLearner then evaluates the +model’s performance on held-out data (the task's test set) and updates the meta- +parameters so that they lead to better generalization across all tasks. This bi- +level optimization ensures that gradients across tasks remain close, enabling +fast adaptation. + +At inference time, the trained model can adapt to a new task using just a small +batch of examples — performing few-shot learning to adjust quickly and accurately. + +This example shows: +- how to implement MAML with RLlib in just a few lines of code. +- how to define a `TorchDifferentiableLearner` to register a custom train loss + function. +- how to define a `TorchMetaLearner` class to implement a custom meta (test) train + loss function. +- how to configure both learners top be used with each others via the + `DifferentiableAlgorithmConfig` and `DifferentiableLearnerConfig`. +- how to update the `RLModule` in a meta-learning fashion. +- how to fine-tune an `RLModule` with gradient descent within a few iterations with + only using the meta (test) loss. + +See :py:class:`~ray.rllib.examples.learners.classes.lr_meta_learner.LRTorchMetaLearner` # noqa +class for details on how to override the main `TorchMetaLearner`. And see +:py:class:`~ray.rllib.examples.learners.classes.lr_differentiable_learner.LRTorchDifferentiableLearner` # noqa +class for an example of how to override the main `TorchDifferentiableLearner`. + +Note, the meta-learner needs a long-enough training (`default_iters`=~70,000) to learn +to adapt quickly to new tasks. + +How to run this script +---------------------- +`python [script file name].py --iters=70000 --meta-train-batch-size=5 --fine-tune-batch-size=5` + +Use the `--meta-train-batch-size` to set the training/testing batch size in meta-learning and +the `--fine-tune-batch-size` to adjust the number of samples used in all updates during +few-shot learning. + +To suppress plotting (plotting is the default) use `--no-plot` and for taking a longer +look at the plot increase the seconds for which plotting is paused at the end of the +script by `--pause-plot-secs`. + +Results to expect +----------------- +You should expect to see sometimes alternating test losses ("Total Loss") due to new +(unseen) tasks during meta learning. In few-shot learning after the meta-learning the +(few shot) loss should decrease almost monotonically. In the plot you can expect to see +a decent adaption to the new task after fine-tuning updates of the `RLModule` weights. + +With `--iters=70_000`, `--meta-train-batch-size=5`, `--fine-tune-batch-size=5`, +`--fine-tune-lr=0.01`, `--fine-tune-iters=10`, `--meta-lr=0.001`, `--noise-std=0.0`, +and no seed defined. +------------------------- + +Iteration: 68000 +Total loss: 0.013758559711277485 +------------------------- + +Iteration: 69000 +Total loss: 0.7246640920639038 +------------------------- + +Iteration: 70000 +Total loss: 3.091259002685547 + +Few shot loss: 2.754437208175659 +Few shot loss: 2.7399725914001465 +Few shot loss: 2.499554395675659 +Few shot loss: 2.1763901710510254 +Few shot loss: 1.793503999710083 +Few shot loss: 1.4362313747406006 +Few shot loss: 1.083552598953247 +Few shot loss: 0.7845061421394348 +Few shot loss: 0.5579453110694885 +Few shot loss: 0.4087105393409729 +""" +import gymnasium as gym +import matplotlib.pyplot as plt +import numpy as np + +from ray.rllib.algorithms.algorithm_config import DifferentiableAlgorithmConfig +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.differentiable_learner_config import ( + DifferentiableLearnerConfig, +) +from ray.rllib.core.learner.training_data import TrainingData +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.algorithms.classes.maml_lr_differentiable_learner import ( + MAMLTorchDifferentiableLearner, +) +from ray.rllib.examples.algorithms.classes.maml_lr_differentiable_rlm import ( + DifferentiableTorchRLModule, +) +from ray.rllib.examples.algorithms.classes.maml_lr_meta_learner import ( + MAMLTorchMetaLearner, +) +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +# Import torch. +torch, _ = try_import_torch() + +# Implement generation of data from sinusoid curves. +def generate_sinusoid_task(batch_size, noise_std=0.1, return_params=False): + """Generate a sinusoid task with random amplitude and phase. + + Args: + batch_size: The number of data points to be generated. + noise_std: An optional standard deviation to be used in the sinusoid + data generation. Defines a linear error term added to the sine + curve. + return_params: If the sampled amplitude and phase should be returned. + + Returns: + Torch tensors with the support data and the labels of a sinusoid + curve. + """ + # Sample the amplitude and the phase for a task. + amplitude = np.random.uniform(0.1, 5.0) + phase = np.random.uniform(0.0, np.pi) + + # Sample the support. + x = np.random.uniform(-5.0, 5.0, (batch_size, 1)) + + # Generate the labels. + y = amplitude * np.sin(x - phase) + + # Add noise. + y += noise_std * np.random.random((batch_size, 1)) + + # If sampled parameters should be returned. + if return_params: + # Return torch tensors. + return ( + torch.tensor(x, dtype=torch.float32), + torch.tensor(y, dtype=torch.float32), + amplitude, + phase, + ) + # Otherwise, return only the sampled data. + else: + return ( + torch.tensor(x, dtype=torch.float32), + torch.tensor(y, dtype=torch.float32), + ) + + +def sample_task(batch_size=10, noise_std=0.1, training_data=False, return_params=False): + """Samples training batches for meta learner and differentiable learner. + + Args: + batch_size: The batch size for both meta learning and task learning. + noise_std: An optional standard deviation to be used in the sinusoid + data generation. Defines a linear error term added to the sine + curve. + training_data: Whether data should be returned as `TrainingData`. + Otherwise, a `MultiAgentBatch` is returned. Default is `False`. + return_params: If the sampled amplitude and phase should be returned. + + Returns: + A tuple with training batches for the meta learner and the differentiable + learner. If `training_data` is `True`, the data is wrapped into + `TrainingData`, otherwise both batches are `MultiAgentBatch`es. + + """ + # Generate training data for meta learner and differentiable learner. + train_batch = {} + generated_data = generate_sinusoid_task( + batch_size * 2, noise_std=noise_std, return_params=return_params + ) + train_batch[Columns.OBS], train_batch["y"] = generated_data[:2] + + # Convert to `MultiAgentBatch`. + meta_train_batch = MultiAgentBatch( + env_steps=batch_size, + policy_batches={ + DEFAULT_MODULE_ID: SampleBatch( + {k: train_batch[k][:batch_size] for k in train_batch} + ) + }, + ) + task_train_batch = MultiAgentBatch( + env_steps=batch_size, + policy_batches={ + DEFAULT_MODULE_ID: SampleBatch( + {k: train_batch[k][batch_size:] for k in train_batch} + ) + }, + ) + # If necessary convert to `TrainingData`. + if training_data: + meta_train_batch = TrainingData( + batch=meta_train_batch, + ) + task_train_batch = TrainingData( + batch=task_train_batch, + ) + + # If amplitude and phase should be returned add them to the return tuple. + if return_params: + return meta_train_batch, task_train_batch, *generated_data[2:] + # Otherwise return solely train data. + else: + return meta_train_batch, task_train_batch + + +# Define arguments. +parser = add_rllib_example_script_args(default_iters=70_000) + +parser.add_argument( + "--meta-train-batch-size", + type=int, + default=5, + help="The number of samples per train and test update (meta-learning).", +) +parser.add_argument( + "--meta-lr", + type=float, + default=0.001, + help="The learning rate to be used for meta learning (in the `MetaLearner`).", +) +parser.add_argument( + "--fine-tune-batch-size", + type=int, + default=10, + help="The number of samples for the fine-tuning updates.", +) +parser.add_argument( + "--noise-std", + type=float, + default=0.0, + help="The standard deviation for noise added to the single tasks.", +) +parser.add_argument( + "--seed", + type=int, + default=None, + help="An optional random seed. If not set, the experiment is not reproducable.", +) +parser.add_argument( + "--fine-tune-iters", + type=int, + default=10, + help="The number of updates in fine-tuning.", +) +parser.add_argument( + "--fine-tune-lr", + type=float, + default=0.01, + help="The learning rate to be used in fine-tuning the model in the test phase.", +) +parser.add_argument( + "--no-plot", + action="store_true", + help=( + "If plotting should suppressed. Otherwise user action is needed to close " + "the plot early." + ), +) +parser.add_argument( + "--pause-plot-secs", + type=int, + default=1000, + help=( + "The number of seconds to keep the plot open. Note the plot can always be " + "closed by the user when open." + ), +) + +# Parse the arguments. +args = parser.parse_args() + +# If a random seed is provided set it for torch and numpy. +if args.seed: + torch.random.manual_seed(args.seed) + np.random.seed(args.seed) + + +if __name__ == "__main__": + # Define the `RLModule`. + module_spec = RLModuleSpec( + module_class=DifferentiableTorchRLModule, + # Note, the spaces are needed by default but are not used. + observation_space=gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32), + action_space=gym.spaces.Box(-np.inf, np.inf, (1,), dtype=np.float32), + ) + # `Learner`s work on `MultiRLModule`s. + multi_module_spec = MultiRLModuleSpec( + rl_module_specs={DEFAULT_MODULE_ID: module_spec} + ) + + # Build the `MultiRLModule`. + module = multi_module_spec.build() + + # Configure the `DifferentiableLearner`. + diff_learner_config = DifferentiableLearnerConfig( + learner_class=MAMLTorchDifferentiableLearner, + minibatch_size=args.meta_train_batch_size, + lr=0.01, + ) + + # Configure the `TorchMetaLearner` via the `DifferentiableAlgorithmConfig`. + config = ( + DifferentiableAlgorithmConfig() + .learners( + # Add the `DifferentiableLearnerConfig`s. + differentiable_learner_configs=[diff_learner_config], + num_gpus_per_learner=args.num_gpus_per_learner or 0, + ) + .training( + lr=args.meta_lr, + train_batch_size=args.meta_train_batch_size, + # Use the full batch in a single update. + minibatch_size=args.meta_train_batch_size, + ) + ) + + # Initialize the `TorchMetaLearner`. + meta_learner = MAMLTorchMetaLearner(config=config, module_spec=module_spec) + # Build the `TorchMetaLearner`. + meta_learner.build() + + for i in range(args.stop_iters): + # Sample the training data. + meta_training_data, task_training_data = sample_task( + args.meta_train_batch_size, noise_std=args.noise_std, training_data=True + ) + + # Update the module. + outs = meta_learner.update( + training_data=meta_training_data, + num_epochs=1, + others_training_data=[task_training_data], + ) + iter = i + 1 + if iter % 1000 == 0: + total_loss = outs["default_policy"]["total_loss"].peek() + print("-------------------------\n") + print(f"Iteration: {iter}") + print(f"Total loss: {total_loss}") + + # Generate test data. + test_batch, _, amplitude, phase = sample_task( + batch_size=args.fine_tune_batch_size, + noise_std=args.noise_std, + return_params=True, + ) + + if config.num_gpus_per_learner > 0: + test_batch = meta_learner._convert_batch_type(test_batch) + + # Run inference and plot results. + with torch.no_grad(): + # Generate a grid for the support. + x_grid = torch.tensor( + np.arange(-5.0, 5.0, 0.02), dtype=torch.float32, device=meta_learner._device + ).view(-1, 1) + # Get label prediction from the model trained by MAML. + y_pred = meta_learner.module[DEFAULT_MODULE_ID]({Columns.OBS: x_grid})["y_pred"] + + # Plot the results if requested. + if not args.no_plot: + # Sort the data by the support. + x_order = np.argsort(test_batch[DEFAULT_MODULE_ID][Columns.OBS].numpy()[:, 0]) + x_sorted = test_batch[DEFAULT_MODULE_ID][Columns.OBS].numpy()[:, 0][x_order] + y_sorted = test_batch[DEFAULT_MODULE_ID]["y"][:, 0][x_order] + + # Plot the data. + def sinusoid(t): + return amplitude * np.sin(t - phase) + + plt.ion() + plt.figure(figsize=(5, 3)) + # Plot the true sinusoid curve. + plt.plot(x_grid, sinusoid(x_grid), "r", label="Ground Truth") + # Add the sampled support values. + plt.plot(x_sorted, y_sorted, "^", color="purple") + # Add the prediction made by the model after MAML training. + plt.plot(x_grid, y_pred, ":", label="Prediction", color="#90EE90") + plt.title(f"MAML Results from {args.fine_tune_iters} fine-tuning steps.") + + # Fine-tune with the meta loss for just a few steps. + optim = meta_learner.get_optimizers_for_module(DEFAULT_MODULE_ID)[0][1] + # Set the learning rate to a larger value. + for g in optim.param_groups: + g["lr"] = args.fine_tune_lr + # Now run the fine-tune iterations and update the model via the meta-learner loss. + for i in range(args.fine_tune_iters): + # Forward pass. + fwd_out = { + DEFAULT_MODULE_ID: meta_learner.module[DEFAULT_MODULE_ID]( + test_batch[DEFAULT_MODULE_ID] + ) + } + # Compute the MSE prediction loss. + loss_per_module = meta_learner.compute_losses(fwd_out=fwd_out, batch=test_batch) + # Optimize parameters. + optim.zero_grad(set_to_none=True) + loss_per_module[DEFAULT_MODULE_ID].backward() + optim.step() + # Show the loss for few-shot learning (fine-tuning). + print(f"Few shot loss: {loss_per_module[DEFAULT_MODULE_ID].item()}") + + # Run the model again after fine-tuning. + with torch.no_grad(): + y_pred_fine_tuned = meta_learner.module[DEFAULT_MODULE_ID]( + {Columns.OBS: x_grid} + )["y_pred"] + + if not args.no_plot: + # Plot the predictions of the fine-tuned model. + plt.plot( + x_grid, + y_pred_fine_tuned, + "-.", + label="Tuned Prediction", + color="green", + mfc="gray", + ) + plt.legend() + plt.show() + + # Pause the plot until the user closes it. + plt.pause(args.pause_plot_secs) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/algorithms/vpg_custom_algorithm.py b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/vpg_custom_algorithm.py new file mode 100644 index 0000000000000000000000000000000000000000..feb3084f5286aad4ec9a9c5b6a4a112c83a13298 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/algorithms/vpg_custom_algorithm.py @@ -0,0 +1,115 @@ +"""Example of how to write a custom Algorithm. + +This is an end-to-end example for how to implement a custom Algorithm, including +a matching AlgorithmConfig class and Learner class. There is no particular RLModule API +needed for this algorithm, which means that any TorchRLModule returning actions +or action distribution parameters suffices. + +The RK algorithm implemented here is "vanilla policy gradient" (VPG) in its simplest +form, without a value function baseline. + +See the actual VPG algorithm class here: +https://github.com/ray-project/ray/blob/master/rllib/examples/algorithms/classes/vpg.py + +The Learner class the algorithm uses by default (if the user doesn't specify a custom +Learner): +https://github.com/ray-project/ray/blob/master/rllib/examples/learners/classes/vpg_torch_learner.py # noqa + +And the RLModule class the algorithm uses by default (if the user doesn't specify a +custom RLModule): +https://github.com/ray-project/ray/blob/master/rllib/examples/rl_modules/classes/vpg_torch_rlm.py # noqa + +This example shows: + - how to subclass the AlgorithmConfig base class to implement a custom algorithm's. + config class. + - how to subclass the Algorithm base class to implement a custom Algorithm, + including its `training_step` method. + - how to subclass the TorchLearner base class to implement a custom Learner with + loss function, overriding `compute_loss_for_module` and + `after_gradient_based_update`. + - how to define a default RLModule used by the algorithm in case the user + doesn't bring their own custom RLModule. The VPG algorithm doesn't require any + specific RLModule APIs, so any RLModule returning actions or action distribution + inputs suffices. + +We compute a plain policy gradient loss without value function baseline. +The experiment shows that even with such a simple setup, our custom algorithm is still +able to successfully learn CartPole-v1. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +With some fine-tuning of the learning rate, the batch size, and maybe the +number of env runners and number of envs per env runner, you should see decent +learning behavior on the CartPole-v1 environment: + ++-----------------------------+------------+--------+------------------+ +| Trial name | status | iter | total time (s) | +| | | | | +|-----------------------------+------------+--------+------------------+ +| VPG_CartPole-v1_2973e_00000 | TERMINATED | 451 | 59.5184 | ++-----------------------------+------------+--------+------------------+ ++-----------------------+------------------------+------------------------+ +| episode_return_mean | num_env_steps_sample | ...env_steps_sampled | +| | d_lifetime | _lifetime_throughput | +|-----------------------+------------------------+------------------------| +| 250.52 | 415787 | 7428.98 | ++-----------------------+------------------------+------------------------+ +""" + +from ray.rllib.examples.algorithms.classes.vpg import VPGConfig +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_reward=250.0, + default_iters=1000, + default_timesteps=1_000_000, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + VPGConfig() + .environment("CartPole-v1") + .training( + # The only VPG-specific setting. How many episodes per train batch? + num_episodes_per_train_batch=10, + # Set other config parameters. + lr=0.0005, + # Note that you don't have to set any specific Learner class, because + # our custom Algorithm already defines the default Learner class to use + # through its `get_default_learner_class` method, which returns + # `VPGTorchLearner`. + # learner_class=VPGTorchLearner, + ) + # Increase the number of EnvRunners (default is 1 for VPG) + # or the number of envs per EnvRunner. + .env_runners(num_env_runners=2, num_envs_per_env_runner=1) + # Plug in your own RLModule class. VPG doesn't require any specific + # RLModule APIs, so any RLModule returning `actions` or `action_dist_inputs` + # from the forward methods works ok. + # .rl_module( + # rl_module_spec=RLModuleSpec(module_class=...), + # ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f5b3516e82144b98dec095177ba3c63982f64cf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/cartpole_dqn_export.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/cartpole_dqn_export.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d05d58b850a1cac3e53a405ce9f6cf9645bb9a89 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/cartpole_dqn_export.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/change_config_during_training.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/change_config_during_training.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f0cfcd3c3c6c0a6e2d4cc345543f8dda595c6eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/change_config_during_training.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/checkpoint_by_custom_criteria.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/checkpoint_by_custom_criteria.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64abbc7015076cb32fe5b094524bcd474931bd6c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/checkpoint_by_custom_criteria.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/continue_training_from_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/continue_training_from_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5093183f4fac9e6bb7fd882b7e562652230f89fc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/continue_training_from_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/onnx_torch_lstm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/onnx_torch_lstm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23e020af9dfdaa5129a5eb84b36c2016da472db2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/onnx_torch_lstm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/restore_1_of_n_agents_from_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/restore_1_of_n_agents_from_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e07dfef050d857026066832d28385d674b65c0af Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/__pycache__/restore_1_of_n_agents_from_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/cartpole_dqn_export.py b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/cartpole_dqn_export.py new file mode 100644 index 0000000000000000000000000000000000000000..dc0f4c79265abc4982e55fbb13b2817d8905bfad --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/cartpole_dqn_export.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python + +# @OldAPIStack + +import os + +import numpy as np + +import ray +import ray._common +from ray.rllib.policy.policy import Policy +from ray.rllib.utils.framework import try_import_tf +from ray.tune.registry import get_trainable_cls + +tf1, tf, tfv = try_import_tf() + +ray.init() + + +def train_and_export_policy_and_model(algo_name, num_steps, model_dir, ckpt_dir): + cls = get_trainable_cls(algo_name) + config = cls.get_default_config() + config.api_stack( + enable_rl_module_and_learner=False, + enable_env_runner_and_connector_v2=False, + ) + # This Example is only for tf. + config.framework("tf") + # Set exporting native (DL-framework) model files to True. + config.export_native_model_files = True + config.env = "CartPole-v1" + alg = config.build() + for _ in range(num_steps): + alg.train() + + # Export Policy checkpoint. + alg.export_policy_checkpoint(ckpt_dir) + # Export tensorflow keras Model for online serving + alg.export_policy_model(model_dir) + + +def restore_saved_model(export_dir): + signature_key = ( + tf1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY + ) + g = tf1.Graph() + with g.as_default(): + with tf1.Session(graph=g) as sess: + meta_graph_def = tf1.saved_model.load( + sess, [tf1.saved_model.tag_constants.SERVING], export_dir + ) + print("Model restored!") + print("Signature Def Information:") + print(meta_graph_def.signature_def[signature_key]) + print("You can inspect the model using TensorFlow SavedModel CLI.") + print("https://www.tensorflow.org/guide/saved_model") + + +def restore_policy_from_checkpoint(export_dir): + # Load the model from the checkpoint. + policy = Policy.from_checkpoint(export_dir) + # Perform a dummy (CartPole) forward pass. + test_obs = np.array([0.1, 0.2, 0.3, 0.4]) + results = policy.compute_single_action(test_obs) + # Check results for correctness. + assert len(results) == 3 + assert results[0].shape == () # pure single action (int) + assert results[1] == [] # RNN states + assert results[2]["action_dist_inputs"].shape == (2,) # categorical inputs + + +if __name__ == "__main__": + algo = "PPO" + model_dir = os.path.join(ray._common.utils.get_user_temp_dir(), "model_export_dir") + ckpt_dir = os.path.join(ray._common.utils.get_user_temp_dir(), "ckpt_export_dir") + num_steps = 1 + train_and_export_policy_and_model(algo, num_steps, model_dir, ckpt_dir) + restore_saved_model(model_dir) + restore_policy_from_checkpoint(ckpt_dir) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/change_config_during_training.py b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/change_config_during_training.py new file mode 100644 index 0000000000000000000000000000000000000000..e2ff77b2111e343d5d79072e7b3b7cf691daf138 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/change_config_during_training.py @@ -0,0 +1,245 @@ +"""Example showing how to continue training an Algorithm with a changed config. + +Use the setup shown in this script if you want to continue a prior experiment, but +would also like to change some of the config values you originally used. + +This example: + - runs a single- or multi-agent CartPole experiment (for multi-agent, we use + different learning rates) thereby checkpointing the state of the Algorithm every n + iterations. The config used is hereafter called "1st config". + - stops the experiment due to some episode return being achieved. + - just for testing purposes, restores the entire algorithm from the latest + checkpoint and checks, whether the state of the restored algo exactly match the + state of the previously saved one. + - then changes the original config used (learning rate and other settings) and + continues training with the restored algorithm and the changed config until a + final episode return is reached. The new config is hereafter called "2nd config". + + +How to run this script +---------------------- +`python [script file name].py --num-agents=[0 or 2] +--stop-reward-first-config=[return at which the algo on 1st config should stop training] +--stop-reward=[the final return to achieve after restoration from the checkpoint with +the 2nd config] +` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +First, you should see the initial tune.Tuner do it's thing: + +Trial status: 1 RUNNING +Current time: 2024-06-03 12:03:39. Total running time: 30s +Logical resource usage: 3.0/12 CPUs, 0/0 GPUs +╭──────────────────────────────────────────────────────────────────────── +│ Trial name status iter total time (s) +├──────────────────────────────────────────────────────────────────────── +│ PPO_CartPole-v1_7b1eb_00000 RUNNING 6 16.265 +╰──────────────────────────────────────────────────────────────────────── +───────────────────────────────────────────────────────────────────────╮ +..._sampled_lifetime ..._trained_lifetime ...episodes_lifetime │ +───────────────────────────────────────────────────────────────────────┤ + 24000 24000 340 │ +───────────────────────────────────────────────────────────────────────╯ +... + +The experiment stops at an average episode return of `--stop-reward-first-config`. + +After the validation of the last checkpoint, a new experiment is started from +scratch, but with the RLlib callback restoring the Algorithm right after +initialization using the previous checkpoint. This new experiment then runs +until `--stop-reward` is reached. + +Trial status: 1 RUNNING +Current time: 2024-06-03 12:05:00. Total running time: 1min 0s +Logical resource usage: 3.0/12 CPUs, 0/0 GPUs +╭──────────────────────────────────────────────────────────────────────── +│ Trial name status iter total time (s) +├──────────────────────────────────────────────────────────────────────── +│ PPO_CartPole-v1_7b1eb_00000 RUNNING 23 14.8372 +╰──────────────────────────────────────────────────────────────────────── +───────────────────────────────────────────────────────────────────────╮ +..._sampled_lifetime ..._trained_lifetime ...episodes_lifetime │ +───────────────────────────────────────────────────────────────────────┤ + 109078 109078 531 │ +───────────────────────────────────────────────────────────────────────╯ + +And if you are using the `--as-test` option, you should see a finel message: + +``` +`env_runners/episode_return_mean` of 450.0 reached! ok +``` +""" +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + LEARNER_RESULTS, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + check, + run_rllib_example_script_experiment, +) +from ray.tune.registry import register_env + +parser = add_rllib_example_script_args( + default_reward=450.0, default_timesteps=10000000, default_iters=2000 +) +parser.add_argument( + "--stop-reward-first-config", + type=float, + default=150.0, + help="Mean episode return after which the Algorithm on the first config should " + "stop training.", +) +# By default, set `args.checkpoint_freq` to 1 and `args.checkpoint_at_end` to True. +parser.set_defaults( + checkpoint_freq=1, + checkpoint_at_end=True, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + register_env( + "ma_cart", lambda cfg: MultiAgentCartPole({"num_agents": args.num_agents}) + ) + + # Simple generic config. + base_config = ( + PPOConfig() + .environment("CartPole-v1" if args.num_agents == 0 else "ma_cart") + .env_runners(create_env_on_local_worker=True) + .training(lr=0.0001) + # TODO (sven): Tune throws a weird error inside the "log json" callback + # when running with this option. The `perf` key in the result dict contains + # binary data (instead of just 2 float values for mem and cpu usage). + # .experimental(_use_msgpack_checkpoints=True) + ) + + # Setup multi-agent, if required. + if args.num_agents > 0: + base_config.multi_agent( + policies={ + f"p{aid}": PolicySpec( + config=AlgorithmConfig.overrides( + lr=5e-5 + * (aid + 1), # agent 1 has double the learning rate as 0. + ) + ) + for aid in range(args.num_agents) + }, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + # Define some stopping criterion. Note that this criterion is an avg episode return + # to be reached. + metric = f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}" + stop = {metric: args.stop_reward_first_config} + + tuner_results = run_rllib_example_script_experiment( + base_config, + args, + stop=stop, + keep_ray_up=True, + ) + + # Perform a very quick test to make sure our algo (upon restoration) did not lose + # its ability to perform well in the env. + # - Extract the best checkpoint. + best_result = tuner_results.get_best_result(metric=metric, mode="max") + assert ( + best_result.metrics[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + >= args.stop_reward_first_config + ) + best_checkpoint_path = best_result.checkpoint.path + + # Rebuild the algorithm (just for testing purposes). + test_algo = base_config.build() + # Load algo's state from the best checkpoint. + test_algo.restore_from_path(best_checkpoint_path) + # Perform some checks on the restored state. + assert test_algo.training_iteration > 0 + # Evaluate on the restored algorithm. + test_eval_results = test_algo.evaluate() + assert ( + test_eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + >= args.stop_reward_first_config + ), test_eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + # Train one iteration to make sure, the performance does not collapse (e.g. due + # to the optimizer weights not having been restored properly). + test_results = test_algo.train() + assert ( + test_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + >= args.stop_reward_first_config + ), test_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + # Stop the test algorithm again. + test_algo.stop() + + # Make sure the algorithm gets restored from a checkpoint right after + # initialization. Note that this includes all subcomponents of the algorithm, + # including the optimizer states in the LearnerGroup/Learner actors. + def on_algorithm_init(algorithm, **kwargs): + module_p0 = algorithm.get_module("p0") + weight_before = convert_to_numpy(next(iter(module_p0.parameters()))) + + algorithm.restore_from_path(best_checkpoint_path) + + # Make sure weights were restored (changed). + weight_after = convert_to_numpy(next(iter(module_p0.parameters()))) + check(weight_before, weight_after, false=True) + + # Change the config. + ( + base_config + # Make sure the algorithm gets restored upon initialization. + .callbacks(on_algorithm_init=on_algorithm_init) + # Change training parameters considerably. + .training( + lr=0.0003, + train_batch_size=5000, + grad_clip=100.0, + gamma=0.996, + num_epochs=6, + vf_loss_coeff=0.01, + ) + # Make multi-CPU/GPU. + .learners(num_learners=2) + # Use more env runners and more envs per env runner. + .env_runners(num_env_runners=3, num_envs_per_env_runner=5) + ) + + # Update the stopping criterium to the final target return per episode. + stop = {metric: args.stop_reward} + + # Run a new experiment with the (RLlib) callback `on_algorithm_init` restoring + # from the best checkpoint. + # Note that the new experiment starts again from iteration=0 (unlike when you + # use `tune.Tuner.restore()` after a crash or interrupted trial). + tuner_results = run_rllib_example_script_experiment(base_config, args, stop=stop) + + # Assert that we have continued training with a different learning rate. + assert ( + tuner_results[0].metrics[LEARNER_RESULTS][DEFAULT_MODULE_ID][ + "default_optimizer_learning_rate" + ] + == base_config.lr + == 0.0003 + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/checkpoint_by_custom_criteria.py b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/checkpoint_by_custom_criteria.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8d30818565b37110db97b04c75697af15d6cba --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/checkpoint_by_custom_criteria.py @@ -0,0 +1,148 @@ +"""Example extracting a checkpoint from n trials using one or more custom criteria. + +This example: + - runs a CartPole experiment with three different learning rates (three tune + "trials"). During the experiment, for each trial, we create a checkpoint at each + iteration. + - at the end of the experiment, we compare the trials and pick the one that + performed best, based on the criterion: Lowest episode count per single iteration + (for CartPole, a low episode count means the episodes are very long and thus the + reward is also very high). + - from that best trial (with the lowest episode count), we then pick those + checkpoints that a) have the lowest policy loss (good) and b) have the highest value + function loss (bad). + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see the performance of the three different learning +rates used here: + ++-----------------------------+------------+-----------------+--------+--------+ +| Trial name | status | loc | lr | iter | +|-----------------------------+------------+-----------------+--------+--------+ +| PPO_CartPole-v1_d7dbe_00000 | TERMINATED | 127.0.0.1:98487 | 0.01 | 17 | +| PPO_CartPole-v1_d7dbe_00001 | TERMINATED | 127.0.0.1:98488 | 0.001 | 8 | +| PPO_CartPole-v1_d7dbe_00002 | TERMINATED | 127.0.0.1:98489 | 0.0001 | 9 | ++-----------------------------+------------+-----------------+--------+--------+ + ++------------------+-------+----------+----------------------+----------------------+ +| total time (s) | ts | reward | episode_reward_max | episode_reward_min | +|------------------+-------+----------+----------------------+----------------------+ +| 28.1068 | 39797 | 151.11 | 500 | 12 | +| 13.304 | 18728 | 158.91 | 500 | 15 | +| 14.8848 | 21069 | 167.36 | 500 | 13 | ++------------------+-------+----------+----------------------+----------------------+ + ++--------------------+ +| episode_len_mean | +|--------------------| +| 151.11 | +| 158.91 | +| 167.36 | ++--------------------+ +""" + +from ray import tune +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + LEARNER_RESULTS, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +parser = add_rllib_example_script_args( + default_reward=450.0, + default_timesteps=100000, + default_iters=200, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Force-set `args.checkpoint_freq` to 1. + args.checkpoint_freq = 1 + + # Simple generic config. + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("CartPole-v1") + # Run 3 trials, each w/ a different learning rate. + .training(lr=tune.grid_search([0.01, 0.001, 0.0001]), train_batch_size=2341) + ) + # Run tune for some iterations and generate checkpoints. + results = run_rllib_example_script_experiment(base_config, args) + + # Get the best of the 3 trials by using some metric. + # NOTE: Choosing the min `episodes_this_iter` automatically picks the trial + # with the best performance (over the entire run (scope="all")): + # The fewer episodes, the longer each episode lasted, the more reward we + # got each episode. + # Setting scope to "last", "last-5-avg", or "last-10-avg" will only compare + # (using `mode=min|max`) the average values of the last 1, 5, or 10 + # iterations with each other, respectively. + # Setting scope to "avg" will compare (using `mode`=min|max) the average + # values over the entire run. + metric = "env_runners/num_episodes" + # notice here `scope` is `all`, meaning for each trial, + # all results (not just the last one) will be examined. + best_result = results.get_best_result(metric=metric, mode="min", scope="all") + value_best_metric = best_result.metrics_dataframe[metric].min() + best_return_best = best_result.metrics_dataframe[ + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}" + ].max() + print( + f"Best trial was the one with lr={best_result.metrics['config']['lr']}. " + f"Reached lowest episode count ({value_best_metric}) in a single iteration and " + f"an average return of {best_return_best}." + ) + + # Confirm, we picked the right trial. + + assert ( + value_best_metric + == results.get_dataframe(filter_metric=metric, filter_mode="min")[metric].min() + ) + + # Get the best checkpoints from the trial, based on different metrics. + # Checkpoint with the lowest policy loss value: + if not args.old_api_stack: + policy_loss_key = f"{LEARNER_RESULTS}/{DEFAULT_MODULE_ID}/policy_loss" + else: + policy_loss_key = "info/learner/default_policy/learner_stats/policy_loss" + best_result = results.get_best_result(metric=policy_loss_key, mode="min") + ckpt = best_result.checkpoint + lowest_policy_loss = best_result.metrics_dataframe[policy_loss_key].min() + print(f"Checkpoint w/ lowest policy loss ({lowest_policy_loss}): {ckpt}") + + # Checkpoint with the highest value-function loss: + if not args.old_api_stack: + vf_loss_key = f"{LEARNER_RESULTS}/{DEFAULT_MODULE_ID}/vf_loss" + else: + vf_loss_key = "info/learner/default_policy/learner_stats/vf_loss" + best_result = results.get_best_result(metric=vf_loss_key, mode="max") + ckpt = best_result.checkpoint + highest_value_fn_loss = best_result.metrics_dataframe[vf_loss_key].max() + print(f"Checkpoint w/ highest value function loss: {ckpt}") + print(f"Highest value function loss: {highest_value_fn_loss}") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/continue_training_from_checkpoint.py b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/continue_training_from_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..1f175abd80f2705ab0643ec76e555c53e2b4b3b3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/continue_training_from_checkpoint.py @@ -0,0 +1,266 @@ +"""Example showing how to restore an Algorithm from a checkpoint and resume training. + +Use the setup shown in this script if your experiments tend to crash after some time, +and you would therefore like to make your setup more robust and fault-tolerant. + +This example: + - runs a single- or multi-agent CartPole experiment (for multi-agent, we use + different learning rates) thereby checkpointing the state of the Algorithm every n + iterations. + - stops the experiment due to an expected crash in the algorithm's main process + after a certain number of iterations. + - just for testing purposes, restores the entire algorithm from the latest + checkpoint and checks, whether the state of the restored algo exactly match the + state of the crashed one. + - then continues training with the restored algorithm until the desired final + episode return is reached. + + +How to run this script +---------------------- +`python [script file name].py --num-agents=[0 or 2] +--stop-reward-crash=[the episode return after which the algo should crash] +--stop-reward=[the final episode return to achieve after(!) restoration from the +checkpoint] +` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +First, you should see the initial tune.Tuner do it's thing: + +Trial status: 1 RUNNING +Current time: 2024-06-03 12:03:39. Total running time: 30s +Logical resource usage: 3.0/12 CPUs, 0/0 GPUs +╭──────────────────────────────────────────────────────────────────────── +│ Trial name status iter total time (s) +├──────────────────────────────────────────────────────────────────────── +│ PPO_CartPole-v1_7b1eb_00000 RUNNING 6 15.362 +╰──────────────────────────────────────────────────────────────────────── +───────────────────────────────────────────────────────────────────────╮ +..._sampled_lifetime ..._trained_lifetime ...episodes_lifetime │ +───────────────────────────────────────────────────────────────────────┤ + 24000 24000 340 │ +───────────────────────────────────────────────────────────────────────╯ +... + +then, you should see the experiment crashing as soon as the `--stop-reward-crash` +has been reached: + +```RuntimeError: Intended crash after reaching trigger return.``` + +At some point, the experiment should resume exactly where it left off (using +the checkpoint and restored Tuner): + +Trial status: 1 RUNNING +Current time: 2024-06-03 12:05:00. Total running time: 1min 0s +Logical resource usage: 3.0/12 CPUs, 0/0 GPUs +╭──────────────────────────────────────────────────────────────────────── +│ Trial name status iter total time (s) +├──────────────────────────────────────────────────────────────────────── +│ PPO_CartPole-v1_7b1eb_00000 RUNNING 27 66.1451 +╰──────────────────────────────────────────────────────────────────────── +───────────────────────────────────────────────────────────────────────╮ +..._sampled_lifetime ..._trained_lifetime ...episodes_lifetime │ +───────────────────────────────────────────────────────────────────────┤ + 108000 108000 531 │ +───────────────────────────────────────────────────────────────────────╯ + +And if you are using the `--as-test` option, you should see a finel message: + +``` +`env_runners/episode_return_mean` of 500.0 reached! ok +``` +""" +import re +import time + +from ray import tune +from ray.air.integrations.wandb import WandbLoggerCallback +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + check_learning_achieved, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args( + default_reward=500.0, default_timesteps=10000000, default_iters=2000 +) +parser.add_argument( + "--stop-reward-crash", + type=float, + default=200.0, + help="Mean episode return after which the Algorithm should crash.", +) +# By default, set `args.checkpoint_freq` to 1 and `args.checkpoint_at_end` to True. +parser.set_defaults( + checkpoint_freq=1, + checkpoint_at_end=True, +) + + +class CrashAfterNIters(RLlibCallback): + """Callback that makes the algo crash after a certain avg. return is reached.""" + + def __init__(self): + super().__init__() + # We have to delay crashing by one iteration just so the checkpoint still + # gets created by Tune after(!) we have reached the trigger avg. return. + self._should_crash = False + + def on_train_result(self, *, algorithm, metrics_logger, result, **kwargs): + # We had already reached the mean-return to crash, the last checkpoint written + # (the one from the previous iteration) should yield that exact avg. return. + if self._should_crash: + raise RuntimeError("Intended crash after reaching trigger return.") + # Reached crashing criterion, crash on next iteration. + elif result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= args.stop_reward_crash: + print( + "Reached trigger return of " + f"{result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}" + ) + self._should_crash = True + + +if __name__ == "__main__": + args = parser.parse_args() + + register_env( + "ma_cart", lambda cfg: MultiAgentCartPole({"num_agents": args.num_agents}) + ) + + # Simple generic config. + config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("CartPole-v1" if args.num_agents == 0 else "ma_cart") + .env_runners(create_env_on_local_worker=True) + .training(lr=0.0001) + .callbacks(CrashAfterNIters) + ) + + # Tune config. + # Need a WandB callback? + tune_callbacks = [] + if args.wandb_key: + project = args.wandb_project or ( + args.algo.lower() + "-" + re.sub("\\W+", "-", str(config.env).lower()) + ) + tune_callbacks.append( + WandbLoggerCallback( + api_key=args.wandb_key, + project=args.wandb_project, + upload_checkpoints=False, + **({"name": args.wandb_run_name} if args.wandb_run_name else {}), + ) + ) + + # Setup multi-agent, if required. + if args.num_agents > 0: + config.multi_agent( + policies={ + f"p{aid}": PolicySpec( + config=AlgorithmConfig.overrides( + lr=5e-5 + * (aid + 1), # agent 1 has double the learning rate as 0. + ) + ) + for aid in range(args.num_agents) + }, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + # Define some stopping criterion. Note that this criterion is an avg episode return + # to be reached. The stop criterion does not consider the built-in crash we are + # triggering through our callback. + stop = { + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + + # Run tune for some iterations and generate checkpoints. + tuner = tune.Tuner( + trainable=config.algo_class, + param_space=config, + run_config=tune.RunConfig( + callbacks=tune_callbacks, + checkpoint_config=tune.CheckpointConfig( + checkpoint_frequency=args.checkpoint_freq, + checkpoint_at_end=args.checkpoint_at_end, + ), + stop=stop, + ), + ) + tuner_results = tuner.fit() + + # Perform a very quick test to make sure our algo (upon restoration) did not lose + # its ability to perform well in the env. + # - Extract the best checkpoint. + metric = f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}" + best_result = tuner_results.get_best_result(metric=metric, mode="max") + assert ( + best_result.metrics[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + >= args.stop_reward_crash + ) + # - Change our config, such that the restored algo will have an env on the local + # EnvRunner (to perform evaluation) and won't crash anymore (remove the crashing + # callback). + config.callbacks(None) + # Rebuild the algorithm (just for testing purposes). + test_algo = config.build() + # Load algo's state from best checkpoint. + test_algo.restore(best_result.checkpoint) + # Perform some checks on the restored state. + assert test_algo.training_iteration > 0 + # Evaluate on the restored algorithm. + test_eval_results = test_algo.evaluate() + assert ( + test_eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + >= args.stop_reward_crash + ), test_eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + # Train one iteration to make sure, the performance does not collapse (e.g. due + # to the optimizer weights not having been restored properly). + test_results = test_algo.train() + assert ( + test_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= args.stop_reward_crash + ), test_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + # Stop the test algorithm again. + test_algo.stop() + + # Create a new Tuner from the existing experiment path (which contains the tuner's + # own checkpoint file). Note that even the WandB logging will be continued without + # creating a new WandB run name. + restored_tuner = tune.Tuner.restore( + path=tuner_results.experiment_path, + trainable=config.algo_class, + param_space=config, + # Important to set this to True b/c the previous trial had failed (due to our + # `CrashAfterNIters` callback). + resume_errored=True, + ) + # Continue the experiment exactly where we left off. + tuner_results = restored_tuner.fit() + + # Not sure, whether this is really necessary, but we have observed the WandB + # logger sometimes not logging some of the last iterations. This sleep here might + # give it enough time to do so. + time.sleep(20) + + if args.as_test: + check_learning_achieved(tuner_results, args.stop_reward, metric=metric) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/onnx_torch_lstm.py b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/onnx_torch_lstm.py new file mode 100644 index 0000000000000000000000000000000000000000..f9f3a3dc2319c3c2d43c6b6c2a6eb1cc242b07a9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/onnx_torch_lstm.py @@ -0,0 +1,131 @@ +# @OldAPIStack + +import numpy as np +import onnxruntime + +import ray +import ray.rllib.algorithms.ppo as ppo +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import add_rllib_example_script_args, check +from ray.rllib.utils.torch_utils import convert_to_torch_tensor + +torch, _ = try_import_torch() + +parser = add_rllib_example_script_args() +parser.set_defaults( + num_env_runners=1, + # ONNX is not supported by RLModule API yet. + old_api_stack=True, +) + + +class ONNXCompatibleWrapper(torch.nn.Module): + def __init__(self, original_model): + super(ONNXCompatibleWrapper, self).__init__() + self.original_model = original_model + + def forward(self, a, b0, b1, c): + # Convert the separate tensor inputs back into the list format + # expected by the original model's forward method. + b = [b0, b1] + ret = self.original_model({"obs": a}, b, c) + # results, state_out_0, state_out_1 + return ret[0], ret[1][0], ret[1][1] + + +if __name__ == "__main__": + args = parser.parse_args() + + ray.init(local_mode=args.local_mode) + + # Configure our PPO Algorithm. + config = ( + ppo.PPOConfig() + .environment("CartPole-v1") + .env_runners(num_env_runners=args.num_env_runners) + .training(model={"use_lstm": True}) + ) + + B = 3 + T = 5 + LSTM_CELL = 256 + + # Input data for a python inference forward call. + test_data_python = { + "obs": np.random.uniform(0, 1.0, size=(B * T, 4)).astype(np.float32), + "state_ins": [ + np.random.uniform(0, 1.0, size=(B, LSTM_CELL)).astype(np.float32), + np.random.uniform(0, 1.0, size=(B, LSTM_CELL)).astype(np.float32), + ], + "seq_lens": np.array([T] * B, np.float32), + } + # Input data for the ONNX session. + test_data_onnx = { + "obs": test_data_python["obs"], + "state_in_0": test_data_python["state_ins"][0], + "state_in_1": test_data_python["state_ins"][1], + "seq_lens": test_data_python["seq_lens"], + } + + # Input data for compiling the ONNX model. + test_data_onnx_input = convert_to_torch_tensor(test_data_onnx) + + # Initialize a PPO Algorithm. + algo = config.build() + + # You could train the model here + # algo.train() + + # Let's run inference on the torch model + policy = algo.get_policy() + result_pytorch, _ = policy.model( + { + "obs": torch.tensor(test_data_python["obs"]), + }, + [ + torch.tensor(test_data_python["state_ins"][0]), + torch.tensor(test_data_python["state_ins"][1]), + ], + torch.tensor(test_data_python["seq_lens"]), + ) + + # Evaluate tensor to fetch numpy array + result_pytorch = result_pytorch.detach().numpy() + + # Wrap the actual ModelV2 with the torch wrapper above to make this all work with + # LSTMs (extra `state` in- and outputs and `seq_lens` inputs). + onnx_compatible = ONNXCompatibleWrapper(policy.model) + exported_model_file = "model.onnx" + input_names = [ + "obs", + "state_in_0", + "state_in_1", + "seq_lens", + ] + + # This line will export the model to ONNX. + torch.onnx.export( + onnx_compatible, + tuple(test_data_onnx_input[n] for n in input_names), + exported_model_file, + export_params=True, + opset_version=11, + do_constant_folding=True, + input_names=input_names, + output_names=[ + "output", + "state_out_0", + "state_out_1", + ], + dynamic_axes={k: {0: "batch_size"} for k in input_names}, + ) + # Start an inference session for the ONNX model. + session = onnxruntime.InferenceSession(exported_model_file, None) + result_onnx = session.run(["output"], test_data_onnx) + + # These results should be equal! + print("PYTORCH", result_pytorch) + print("ONNX", result_onnx[0]) + + check(result_pytorch, result_onnx[0]) + print("Model outputs are equal. PASSED") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/restore_1_of_n_agents_from_checkpoint.py b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/restore_1_of_n_agents_from_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..1c420ae6ef110f53f8e9d08e15c6dedd121f5870 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/checkpoints/restore_1_of_n_agents_from_checkpoint.py @@ -0,0 +1,170 @@ +"""Example demonstrating how to load module weights for 1 of n agents from a checkpoint. + +This example: + - Runs a multi-agent `Pendulum-v1` experiment with >= 2 policies, p0, p1, etc.. + - Saves a checkpoint of the `MultiRLModule` every `--checkpoint-freq` + iterations. + - Stops the experiments after the agents reach a combined return of -800. + - Picks the best checkpoint by combined return and restores p0 from it. + - Runs a second experiment with the restored `RLModule` for p0 and + a fresh `RLModule` for the other policies. + - Stops the second experiment after the agents reach a combined return of -800. + + +How to run this script +---------------------- +`python [script file name].py --num-agents=2 +--checkpoint-freq=20 --checkpoint-at-end` + +Control the number of agents and policies (RLModules) via --num-agents and +--num-policies. + +Control the number of checkpoints by setting `--checkpoint-freq` to a value > 0. +Note that the checkpoint frequency is per iteration and this example needs at +least a single checkpoint to load the RLModule weights for policy 0. +If `--checkpoint-at-end` is set, a checkpoint will be saved at the end of the +experiment. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should expect a reward of -400.0 eventually being achieved by a simple +single PPO policy. In the second run of the experiment, the MultiRLModule weights +for policy 0 are restored from the checkpoint of the first run. The reward for a +single agent should be -400.0 again, but the training time should be shorter +(around 30 iterations instead of 190) due to the fact that one policy is already +an expert from the get go. +""" + +from pathlib import Path + +from ray.rllib.algorithms.callbacks import DefaultCallbacks +from ray.rllib.core import ( + COMPONENT_LEARNER, + COMPONENT_LEARNER_GROUP, + COMPONENT_RL_MODULE, +) +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentPendulum +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + check, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env +from ray.tune.result import TRAINING_ITERATION + +parser = add_rllib_example_script_args( + # Pendulum-v1 sum of 2 agents (each agent reaches -250). + default_reward=-500.0, +) +parser.set_defaults( + checkpoint_freq=1, + num_agents=2, +) +# TODO (sven): This arg is currently ignored (hard-set to 2). +parser.add_argument("--num-policies", type=int, default=2) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + if args.num_agents > 1: + register_env( + "env", + lambda _: MultiAgentPendulum(config={"num_agents": args.num_agents}), + ) + else: + raise ValueError( + f"`num_agents` must be > 1, but is {args.num_agents}." + "Read the script docstring for more information." + ) + + assert args.checkpoint_freq > 0, ( + "This example requires at least one checkpoint to load the RLModule " + "weights for policy 0." + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env") + .training( + train_batch_size_per_learner=512, + minibatch_size=64, + lambda_=0.1, + gamma=0.95, + lr=0.0003, + vf_clip_param=10.0, + ) + .rl_module( + model_config=DefaultModelConfig(fcnet_activation="relu"), + ) + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + # Augment the base config with further settings and train the agents. + results = run_rllib_example_script_experiment(base_config, args, keep_ray_up=True) + + # Now swap in the RLModule weights for policy 0. + chkpt_path = results.get_best_result().checkpoint.path + p_0_module_state_path = ( + Path(chkpt_path) # <- algorithm's checkpoint dir + / COMPONENT_LEARNER_GROUP # <- learner group + / COMPONENT_LEARNER # <- learner + / COMPONENT_RL_MODULE # <- MultiRLModule + / "p0" # <- (single) RLModule + ) + + class LoadP0OnAlgoInitCallback(DefaultCallbacks): + def on_algorithm_init(self, *, algorithm, **kwargs): + module_p0 = algorithm.get_module("p0") + weight_before = convert_to_numpy(next(iter(module_p0.parameters()))) + algorithm.restore_from_path( + p_0_module_state_path, + component=( + COMPONENT_LEARNER_GROUP + + "/" + + COMPONENT_LEARNER + + "/" + + COMPONENT_RL_MODULE + + "/p0" + ), + ) + # Make sure weights were updated. + weight_after = convert_to_numpy(next(iter(module_p0.parameters()))) + check(weight_before, weight_after, false=True) + + base_config.callbacks(LoadP0OnAlgoInitCallback) + + # Define stopping criteria. + stop = { + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -800.0, + f"{ENV_RUNNER_RESULTS}/{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 100000, + TRAINING_ITERATION: 100, + } + + # Run the experiment again with the restored MultiRLModule. + run_rllib_example_script_experiment(base_config, args, stop=stop) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a009fca47a2b9875cced8b4c53095ca1170c764a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/count_based_curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/count_based_curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44e6cbe1945f7deafd138d16198f573846ad4893 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/count_based_curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c37abfad4d0110146ebaf508ca80ec9c92e640a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/flatten_observations_dict_space.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/flatten_observations_dict_space.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91a14390ad872c04635b74d205acefe8925a71b8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/flatten_observations_dict_space.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/frame_stacking.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/frame_stacking.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fca2c336ef8202292027963e15fa353a178b1da Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/frame_stacking.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/mean_std_filtering.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/mean_std_filtering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7ac7e23e314f572ab4ba3ee855ae332c8904e93 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/mean_std_filtering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/multi_agent_observation_preprocessor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/multi_agent_observation_preprocessor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23eb3ca2b5d09f01d26e5823ec6629ff25d4e5e7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/multi_agent_observation_preprocessor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/prev_actions_prev_rewards.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/prev_actions_prev_rewards.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cca83b4c884a7ce3ae9dcfe53b67e747978a2e2c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/prev_actions_prev_rewards.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/single_agent_observation_preprocessor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/single_agent_observation_preprocessor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..285c25086a5a964d16dcee0b0078f8460e3ab640 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/__pycache__/single_agent_observation_preprocessor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c9a8e5fdff87d94ee8743292d7b59f16dce0b7e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/add_other_agents_row_index_to_xy_pos.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/add_other_agents_row_index_to_xy_pos.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..504af680cf161fb0e781f686de7c1a4cc0f5dcdb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/add_other_agents_row_index_to_xy_pos.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/count_based_curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/count_based_curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a886b6664ce42e53e30c6e48251a0b88266673f7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/count_based_curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..893269cc1fe9a935a1b7f588cc7063f4699167c9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/protobuf_cartpole_observation_decoder.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/protobuf_cartpole_observation_decoder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..020ff0ed337fcbb5872d695d6d950c2941ba785d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/__pycache__/protobuf_cartpole_observation_decoder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/add_other_agents_row_index_to_xy_pos.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/add_other_agents_row_index_to_xy_pos.py new file mode 100644 index 0000000000000000000000000000000000000000..a1e3d49ce001b0c5373c92df7a452bb2099dec2a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/add_other_agents_row_index_to_xy_pos.py @@ -0,0 +1,114 @@ +from typing import Any + +import gymnasium as gym +import numpy as np + +from ray.rllib.connectors.env_to_module.observation_preprocessor import ( + MultiAgentObservationPreprocessor, +) +from ray.rllib.utils.annotations import override + + +class AddOtherAgentsRowIndexToXYPos(MultiAgentObservationPreprocessor): + """Adds other agent's row index to an x/y-observation for an agent. + + Run this connector with this env: + :py:class:`~ray.rllib.examples.env.classes.multi_agent.double_row_corridor_env.DoubleRowCorridorEnv` # noqa + + In this env, 2 agents walk around in a grid-world and must, each separately, reach + their individual goal position to receive a final reward. However, if they collide + while search for these goal positions, another larger reward is given to both + agents. Thus, optimal policies aim at seeking the other agent first, and only then + proceeding to their agent's goal position. + + Each agents' observation space is a 2-tuple encoding the x/y position + (x=row, y=column). + This connector converts these observations to: + A dict for `agent_0` of structure: + { + "agent": Discrete index encoding the position of the agent, + "other_agent_row": Discrete(2), indicating whether the other agent is in row 0 + or row 1, + } + And a 3-tuple for `agent_1`, encoding the x/y position of `agent_1` plus the row + index (0 or 1) of `agent_0`. + + Note that the row information for the respective other agent, which this connector + provides, is needed for learning an optimal policy for any of the agents, because + the env rewards the first collision between the two agents. Hence, an agent needs to + have information on which row the respective other agent is currently in, so it can + change to this row and try to collide with this other agent. + """ + + @override(MultiAgentObservationPreprocessor) + def recompute_output_observation_space( + self, + input_observation_space, + input_action_space, + ) -> gym.Space: + """Maps the original (input) observation space to the new one. + + Original observation space is `Dict({agent_n: Box(4,), ...})`. + Converts the space for `self.agent` into information specific to this agent, + plus the current row of the respective other agent. + Output observation space is then: + `Dict({`agent_n`: Dict(Discrete, Discrete), ...}), where the 1st Discrete + is the position index of the agent and the 2nd Discrete encodes the current row + of the other agent (0 or 1). If the other agent is already done with the episode + (has reached its goal state) a special value of 2 is used. + """ + agent_0_space = input_observation_space.spaces["agent_0"] + self._env_corridor_len = agent_0_space.high[1] + 1 # Box.high is inclusive. + # Env has always 2 rows (and `self._env_corridor_len` columns). + num_discrete = int(2 * self._env_corridor_len) + spaces = { + "agent_0": gym.spaces.Dict( + { + # Exact position of this agent (as an int index). + "agent": gym.spaces.Discrete(num_discrete), + # Row (0 or 1) of other agent. Or 2, if other agent is already done. + "other_agent_row": gym.spaces.Discrete(3), + } + ), + "agent_1": gym.spaces.Box( + 0, + agent_0_space.high[1], # 1=column + shape=(3,), + dtype=np.float32, + ), + } + return gym.spaces.Dict(spaces) + + @override(MultiAgentObservationPreprocessor) + def preprocess(self, observations, episode) -> Any: + # Observations: dict of keys "agent_0" and "agent_1", mapping to the respective + # x/y positions of these agents (x=row, y=col). + # For example: [1.0, 4.0] means the agent is in row 1 and column 4. + + new_obs = {} + # 2=agent is already done + row_agent_0 = observations.get("agent_0", [2])[0] + row_agent_1 = observations.get("agent_1", [2])[0] + + if "agent_0" in observations: + # Compute `agent_0` and `agent_1` enhanced observation. + index_obs_agent_0 = ( + observations["agent_0"][0] * self._env_corridor_len + + observations["agent_0"][1] + ) + new_obs["agent_0"] = { + "agent": index_obs_agent_0, + "other_agent_row": row_agent_1, + } + + if "agent_1" in observations: + new_obs["agent_1"] = np.array( + [ + observations["agent_1"][0], + observations["agent_1"][1], + row_agent_0, + ], + dtype=np.float32, + ) + + return new_obs diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/count_based_curiosity.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/count_based_curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..5985d3deee02eb0a990ec5256855d4951f471d35 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/count_based_curiosity.py @@ -0,0 +1,92 @@ +from collections import Counter +from typing import Any, List, Optional + +import gymnasium as gym + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.typing import EpisodeType + + +class CountBasedCuriosity(ConnectorV2): + """Learner ConnectorV2 piece to compute intrinsic rewards based on obs counts. + + Add this connector piece to your Learner pipeline, through your algo config: + ``` + config.training( + learner_connector=lambda obs_sp, act_sp: CountBasedCuriosity() + ) + ``` + + Intrinsic rewards are computed on the Learner side based on naive observation + counts, which is why this connector should only be used for simple environments + with a reasonable number of possible observations. The intrinsic reward for a given + timestep is: + r(i) = intrinsic_reward_coeff * (1 / C(obs(i))) + where C is the total (lifetime) count of the obs at timestep i. + + The intrinsic reward is added to the extrinsic reward and saved back into the + episode (under the main "rewards" key). + + Note that the computation and saving back to the episode all happens before the + actual train batch is generated from the episode data. Thus, the Learner and the + RLModule used do not take notice of the extra reward added. + + If you would like to use a more sophisticated mechanism for intrinsic reward + computations, take a look at the `EuclidianDistanceBasedCuriosity` connector piece + at `ray.rllib.examples.connectors.classes.euclidian_distance_based_curiosity` + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + intrinsic_reward_coeff: float = 1.0, + **kwargs, + ): + """Initializes a CountBasedCuriosity instance. + + Args: + intrinsic_reward_coeff: The weight with which to multiply the intrinsic + reward before adding (and saving) it back to the main (extrinsic) + reward of the episode at each timestep. + """ + super().__init__(input_observation_space, input_action_space) + + # Naive observation counter. + self._counts = Counter() + self.intrinsic_reward_coeff = intrinsic_reward_coeff + + def __call__( + self, + *, + rl_module: RLModule, + batch: Any, + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # Loop through all episodes and change the reward to + # [reward + intrinsic reward] + for sa_episode in self.single_agent_episode_iterator( + episodes=episodes, agents_that_stepped_only=False + ): + # Loop through all observations, except the last one. + observations = sa_episode.get_observations(slice(None, -1)) + # Get all respective extrinsic rewards. + rewards = sa_episode.get_rewards() + + for i, (obs, rew) in enumerate(zip(observations, rewards)): + # Add 1 to obs counter. + obs = tuple(obs) + self._counts[obs] += 1 + # Compute the count-based intrinsic reward and add it to the extrinsic + # reward. + rew += self.intrinsic_reward_coeff * (1 / self._counts[obs]) + # Store the new reward back to the episode (under the correct + # timestep/index). + sa_episode.set_rewards(new_data=rew, at_indices=i) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/euclidian_distance_based_curiosity.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/euclidian_distance_based_curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..c50a2caae5d744e40a95f1c78608ca0f99050398 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/euclidian_distance_based_curiosity.py @@ -0,0 +1,122 @@ +from collections import deque +from typing import Any, List, Optional + +import gymnasium as gym +import numpy as np + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.typing import EpisodeType + + +class EuclidianDistanceBasedCuriosity(ConnectorV2): + """Learner ConnectorV2 piece computing intrinsic rewards with euclidian distance. + + Add this connector piece to your Learner pipeline, through your algo config: + ``` + config.training( + learner_connector=lambda obs_sp, act_sp: EuclidianDistanceBasedCuriosity() + ) + ``` + + Intrinsic rewards are computed on the Learner side based on comparing the euclidian + distance of observations vs already seen ones. A configurable number of observations + will be stored in a FIFO buffer and all incoming observations have their distance + measured against those. + + The minimum distance measured is the intrinsic reward for the incoming obs + (multiplied by a fixed coeffieicnt and added to the "main" extrinsic reward): + r(i) = intrinsic_reward_coeff * min(ED(o, o(i)) for o in stored_obs)) + where `ED` is the euclidian distance and `stored_obs` is the buffer. + + The intrinsic reward is then added to the extrinsic reward and saved back into the + episode (under the main "rewards" key). + + Note that the computation and saving back to the episode all happens before the + actual train batch is generated from the episode data. Thus, the Learner and the + RLModule used do not take notice of the extra reward added. + + Only one observation per incoming episode will be stored as a new one in the buffer. + Thereby, we pick the observation with the largest `min(ED)` value over all already + stored observations to be stored per episode. + + If you would like to use a simpler, count-based mechanism for intrinsic reward + computations, take a look at the `CountBasedCuriosity` connector piece + at `ray.rllib.examples.connectors.classes.count_based_curiosity` + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + intrinsic_reward_coeff: float = 1.0, + max_buffer_size: int = 100, + **kwargs, + ): + """Initializes a CountBasedCuriosity instance. + + Args: + intrinsic_reward_coeff: The weight with which to multiply the intrinsic + reward before adding (and saving) it back to the main (extrinsic) + reward of the episode at each timestep. + """ + super().__init__(input_observation_space, input_action_space) + + # Create an observation buffer + self.obs_buffer = deque(maxlen=max_buffer_size) + self.intrinsic_reward_coeff = intrinsic_reward_coeff + + self._test = 0 + + def __call__( + self, + *, + rl_module: RLModule, + batch: Any, + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + if self._test > 10: + return batch + self._test += 1 + # Loop through all episodes and change the reward to + # [reward + intrinsic reward] + for sa_episode in self.single_agent_episode_iterator( + episodes=episodes, agents_that_stepped_only=False + ): + # Loop through all obs, except the last one. + observations = sa_episode.get_observations(slice(None, -1)) + # Get all respective (extrinsic) rewards. + rewards = sa_episode.get_rewards() + + max_dist_obs = None + max_dist = float("-inf") + for i, (obs, rew) in enumerate(zip(observations, rewards)): + # Compare obs to all stored observations and compute euclidian distance. + min_dist = 0.0 + if self.obs_buffer: + min_dist = min( + np.sqrt(np.sum((obs - stored_obs) ** 2)) + for stored_obs in self.obs_buffer + ) + if min_dist > max_dist: + max_dist = min_dist + max_dist_obs = obs + + # Compute our euclidian distance-based intrinsic reward and add it to + # the main (extrinsic) reward. + rew += self.intrinsic_reward_coeff * min_dist + # Store the new reward back to the episode (under the correct + # timestep/index). + sa_episode.set_rewards(new_data=rew, at_indices=i) + + # Add the one observation of this episode with the largest (min) euclidian + # dist to all already stored obs to the buffer (maybe throwing out the + # oldest obs in there). + if max_dist_obs is not None: + self.obs_buffer.append(max_dist_obs) + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/protobuf_cartpole_observation_decoder.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/protobuf_cartpole_observation_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..af7e81d481de2d4d63eb1569129834d01832b5e7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/classes/protobuf_cartpole_observation_decoder.py @@ -0,0 +1,82 @@ +from typing import Any, List, Optional + +import gymnasium as gym +import numpy as np + +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.examples.envs.classes.utils.cartpole_observations_proto import ( + CartPoleObservation, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType + + +class ProtobufCartPoleObservationDecoder(ConnectorV2): + """Env-to-module ConnectorV2 piece decoding protobuf obs into CartPole-v1 obs. + + Add this connector piece to your env-to-module pipeline, through your algo config: + ``` + config.env_runners( + env_to_module_connector=( + lambda env, spaces, device: ProtobufCartPoleObservationDecoder() + ) + ) + ``` + + The incoming observation space must be a 1D Box of dtype uint8 + (which is the same as a binary string). The outgoing observation space is the + normal CartPole-v1 1D space: Box(-inf, inf, (4,), float32). + """ + + @override(ConnectorV2) + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + # Make sure the incoming observation space is a protobuf (binary string). + assert ( + isinstance(input_observation_space, gym.spaces.Box) + and len(input_observation_space.shape) == 1 + and input_observation_space.dtype.name == "uint8" + ) + # Return CartPole-v1's natural observation space. + return gym.spaces.Box(float("-inf"), float("inf"), (4,), np.float32) + + def __call__( + self, + *, + rl_module: RLModule, + batch: Any, + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # Loop through all episodes and change the observation from a binary string + # to an actual 1D np.ndarray (normal CartPole-v1 obs). + for sa_episode in self.single_agent_episode_iterator(episodes=episodes): + # Get last obs (binary string). + obs = sa_episode.get_observations(-1) + obs_bytes = obs.tobytes() + obs_protobuf = CartPoleObservation() + obs_protobuf.ParseFromString(obs_bytes) + + # Set up the natural CartPole-v1 observation tensor from the protobuf + # values. + new_obs = np.array( + [ + obs_protobuf.x_pos, + obs_protobuf.x_veloc, + obs_protobuf.angle_pos, + obs_protobuf.angle_veloc, + ], + np.float32, + ) + + # Write the new observation (1D tensor) back into the Episode. + sa_episode.set_observations(new_data=new_obs, at_indices=-1) + + # Return `data` as-is. + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/count_based_curiosity.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/count_based_curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..ad09e4ceb6bf2665d0766ebe218d071ca2ada19f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/count_based_curiosity.py @@ -0,0 +1,14 @@ +"""Placeholder for training with count-based curiosity. + +The actual script can be found at a different location (see code below). +""" + +if __name__ == "__main__": + import subprocess + import sys + + # Forward to "python ../curiosity/[same script name].py [same options]" + command = [sys.executable, "../curiosity/", sys.argv[0]] + sys.argv[1:] + + # Run the script. + subprocess.run(command, capture_output=True) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/euclidian_distance_based_curiosity.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/euclidian_distance_based_curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..6e52de76791304545eebcf1aea76fe93cb2c6f39 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/euclidian_distance_based_curiosity.py @@ -0,0 +1,14 @@ +"""Placeholder for training with euclidian distance-based curiosity. + +The actual script can be found at a different location (see code below). +""" + +if __name__ == "__main__": + import subprocess + import sys + + # Forward to "python ../curiosity/[same script name].py [same options]" + command = [sys.executable, "../curiosity/", sys.argv[0]] + sys.argv[1:] + + # Run the script. + subprocess.run(command, capture_output=True) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/flatten_observations_dict_space.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/flatten_observations_dict_space.py new file mode 100644 index 0000000000000000000000000000000000000000..246f6131e75e5bad23057da00d44c4ef9348cd20 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/flatten_observations_dict_space.py @@ -0,0 +1,151 @@ +"""Example using a ConnectorV2 to flatten arbitrarily nested dict or tuple observations. + +An RLlib Algorithm has 3 distinct connector pipelines: +- An env-to-module pipeline in an EnvRunner accepting a list of episodes and producing +a batch for an RLModule to compute actions (`forward_inference()` or +`forward_exploration()`). +- A module-to-env pipeline in an EnvRunner taking the RLModule's output and converting +it into an action readable by the environment. +- A learner connector pipeline on a Learner taking a list of episodes and producing +a batch for an RLModule to perform the training forward pass (`forward_train()`). + +Each of these pipelines has a fixed set of default ConnectorV2 pieces that RLlib +adds/prepends to these pipelines in order to perform the most basic functionalities. +For example, RLlib adds the `AddObservationsFromEpisodesToBatch` ConnectorV2 into any +env-to-module pipeline to make sure the batch for computing actions contains - at the +minimum - the most recent observation. + +On top of these default ConnectorV2 pieces, users can define their own ConnectorV2 +pieces (or use the ones available already in RLlib) and add them to one of the 3 +different pipelines described above, as required. + +This example: + - shows how the `FlattenObservation` ConnectorV2 piece can be added to the + env-to-module pipeline. + - demonstrates that by using this connector, any arbitrarily nested dict or tuple + observations is properly flattened into a simple 1D tensor, for easier RLModule + processing. + - shows how - in a multi-agent setup - individual agents can be specified, whose + observations should be flattened (while other agents' observations will always + be left as-is). + - uses a variant of the CartPole-v1 environment, in which the 4 observation items + (x-pos, x-veloc, angle, and angle-veloc) are taken apart and put into a nested dict + with the structure: + { + "x-pos": [x-pos], + "angular-pos": { + "value": [angle], + "some_random_stuff": [random Discrete(3)], # <- should be ignored by algo + }, + "velocs": Tuple([x-veloc], [angle-veloc]), + } + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- + ++---------------------+------------+----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|---------------------+------------+----------------+--------+------------------+ +| PPO_env_a2fd6_00000 | TERMINATED | 127.0.0.1:7409 | 25 | 24.1426 | ++---------------------+------------+----------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| num_env_steps_sample | num_env_steps_traine | episode_return_mean | +| d_lifetime | d_lifetime | | ++------------------------+------------------------+------------------------| +| 100000 | 100000 | 421.42 | ++------------------------+------------------------+------------------------+ +""" +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.envs.classes.cartpole_with_dict_observation_space import ( + CartPoleWithDictObservationSpace, +) +from ray.rllib.examples.envs.classes.multi_agent import ( + MultiAgentCartPoleWithDictObservationSpace, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +# Read in common example script command line arguments. +parser = add_rllib_example_script_args(default_timesteps=200000, default_reward=400.0) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Define env-to-module-connector pipeline for the new stack. + def _env_to_module_pipeline(env, spaces, device): + return FlattenObservations(multi_agent=args.num_agents > 0) + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "env", + lambda _: MultiAgentCartPoleWithDictObservationSpace( + config={"num_agents": args.num_agents} + ), + ) + else: + register_env("env", lambda _: CartPoleWithDictObservationSpace()) + + # Define the AlgorithmConfig used. + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env") + .env_runners(env_to_module_connector=_env_to_module_pipeline) + .training( + gamma=0.99, + lr=0.0003, + ) + .rl_module( + model_config=DefaultModelConfig( + fcnet_hiddens=[32], + fcnet_activation="linear", + vf_share_layers=True, + ), + ) + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + # PPO-specific settings (for better learning behavior only). + if args.algo == "PPO": + base_config.training( + num_epochs=6, + vf_loss_coeff=0.01, + ) + # IMPALA-specific settings (for better learning behavior only). + elif args.algo == "IMPALA": + base_config.training( + lr=0.0005, + vf_loss_coeff=0.05, + entropy_coeff=0.0, + ) + + # Run everything as configured. + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/frame_stacking.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/frame_stacking.py new file mode 100644 index 0000000000000000000000000000000000000000..7eb83725b6087d94329489e2465d87cd3c7a2638 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/frame_stacking.py @@ -0,0 +1,227 @@ +"""Example using 2 ConnectorV2 for observation frame-stacking in Atari environments. + +An RLlib Algorithm has 3 distinct connector pipelines: +- An env-to-module pipeline in an EnvRunner accepting a list of episodes and producing +a batch for an RLModule to compute actions (`forward_inference()` or +`forward_exploration()`). +- A module-to-env pipeline in an EnvRunner taking the RLModule's output and converting +it into an action readable by the environment. +- A learner connector pipeline on a Learner taking a list of episodes and producing +a batch for an RLModule to perform the training forward pass (`forward_train()`). + +Each of these pipelines has a fixed set of default ConnectorV2 pieces that RLlib +adds/prepends to these pipelines in order to perform the most basic functionalities. +For example, RLlib adds the `AddObservationsFromEpisodesToBatch` ConnectorV2 into any +env-to-module pipeline to make sure the batch for computing actions contains - at the +minimum - the most recent observation. + +On top of these default ConnectorV2 pieces, users can define their own ConnectorV2 +pieces (or use the ones available already in RLlib) and add them to one of the 3 +different pipelines described above, as required. + +This example: + - shows how the `FrameStackingEnvToModule` ConnectorV2 piece can be added to the + env-to-module pipeline. + - shows how the `FrameStackingLearner` ConnectorV2 piece can be added to the + learner connector pipeline. + - demonstrates that using these two pieces (rather than performing framestacking + already inside the environment using a gymnasium wrapper) increases overall + performance by about 5%. + + +How to run this script +---------------------- +`python [script file name].py --num-frames=4 --env=ALE/Pong-v5` + +Use the `--num-frames` option to define the number of observations to framestack. +If you don't want to use Connectors to perform the framestacking, set the +`--use-gym-wrapper-framestacking` flag to perform framestacking already inside a +gymnasium observation wrapper. In this case though, be aware that the tensors being +sent through the network are `--num-frames` x larger than if you use the Connector +setup. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- + +With `--num-frames=4` and using the two extra ConnectorV2 pieces (in the env-to-module +and learner connector pipelines), you should see something like this using: +`--env ALE/Pong-v5 --num-learners=4 --num-gpus-per-learner=1 --num-env-runners=95` ++---------------------------+------------+--------+------------------+... +| Trial name | status | iter | total time (s) | +| | | | | +|---------------------------+------------+--------+------------------+... +| PPO_atari-env_2fc4a_00000 | TERMINATED | 200 | 335.837 | ++---------------------------+------------+--------+------------------+... + +Note that the time to run these 200 iterations is about ~5% faster than when +performing framestacking already inside the environment (using a +`gymnasium.wrappers.ObservationWrapper`), due to the additional network traffic +needed (sending back 4x[obs] batches instead of 1x[obs] to the learners). + +Thus, with the `--use-gym-wrapper-framestacking` option (all other options being equal), +the output looks like this: ++---------------------------+------------+--------+------------------+... +| Trial name | status | iter | total time (s) | +| | | | | +|---------------------------+------------+--------+------------------+... +| PPO_atari-env_2fc4a_00000 | TERMINATED | 200 | 351.505 | ++---------------------------+------------+--------+------------------+... +""" +import gymnasium as gym + +from ray.rllib.connectors.env_to_module.frame_stacking import FrameStackingEnvToModule +from ray.rllib.connectors.learner.frame_stacking import FrameStackingLearner +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack +from ray.rllib.examples.envs.classes.multi_agent import make_multi_agent +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +# Read in common example script command line arguments. +parser = add_rllib_example_script_args( + default_timesteps=5000000, default_reward=20.0, default_iters=200 +) +# Use Pong by default. +parser.set_defaults( + env="ale_py:ALE/Pong-v5", +) +parser.add_argument( + "--num-frames", + type=int, + default=4, + help="The number of observation frames to stack.", +) +parser.add_argument( + "--use-gym-wrapper-framestacking", + action="store_true", + help="Whether to use RLlib's Atari wrapper's framestacking capabilities (as " + "opposed to doing it via a specific ConenctorV2 pipeline).", +) + + +if __name__ == "__main__": + from ray import tune + + args = parser.parse_args() + + # Define our custom connector pipelines. + def _make_env_to_module_connector(env, spaces, device): + # Create the env-to-module connector. We return an individual connector piece + # here, which RLlib automatically integrates into a pipeline (and + # add its default connector piece to the end of that pipeline). + # The default pipeline automatically fixes the input- and output spaces of the + # individual connector pieces in it. + # Note that since the frame stacking connector does NOT write information + # back to the episode (in order to save memory and network traffic), we + # also need to perform the same procedure on the Learner end (see below + # where we set up the Learner pipeline). + return FrameStackingEnvToModule( + num_frames=args.num_frames, + multi_agent=args.num_agents > 0, + ) + + def _make_learner_connector(input_observation_space, input_action_space): + # Create the learner connector. + return FrameStackingLearner( + num_frames=args.num_frames, + multi_agent=args.num_agents > 0, + ) + + # Create a custom Atari setup (w/o the usual RLlib-hard-coded framestacking in it). + # We would like our frame stacking connector to do this job. + def _env_creator(cfg): + return wrap_atari_for_new_api_stack( + gym.make(args.env, **cfg, **{"render_mode": "rgb_array"}), + # Perform framestacking either through ConnectorV2 or right here through + # the observation wrapper. + framestack=( + args.num_frames if args.use_gym_wrapper_framestacking else None + ), + ) + + if args.num_agents > 0: + tune.register_env( + "atari-env", + lambda cfg: make_multi_agent(_env_creator)( + dict(cfg, **{"num_agents": args.num_agents}) + ), + ) + else: + tune.register_env("atari-env", _env_creator) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + "atari-env", + env_config={ + # Make analogous to old v4 + NoFrameskip. + "frameskip": 1, + "full_action_space": False, + "repeat_action_probability": 0.0, + }, + clip_rewards=True, + ) + .env_runners( + # ... new EnvRunner and our frame stacking env-to-module connector. + env_to_module_connector=( + None + if args.use_gym_wrapper_framestacking + else _make_env_to_module_connector + ), + num_envs_per_env_runner=2, + ) + .training( + # Use our frame stacking learner connector. + learner_connector=( + None if args.use_gym_wrapper_framestacking else _make_learner_connector + ), + entropy_coeff=0.01, + # Linearly adjust learning rate based on number of GPUs. + lr=0.00015 * (args.num_learners or 1), + grad_clip=100.0, + grad_clip_by="global_norm", + ) + .rl_module( + model_config=DefaultModelConfig( + vf_share_layers=True, + conv_filters=[(16, 4, 2), (32, 4, 2), (64, 4, 2), (128, 4, 2)], + conv_activation="relu", + head_fcnet_hiddens=[256], + ), + ) + ) + + # PPO specific settings. + if args.algo == "PPO": + base_config.training( + num_epochs=10, + minibatch_size=64, + lambda_=0.95, + kl_coeff=0.5, + clip_param=0.1, + vf_clip_param=10.0, + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + # Run everything as configured. + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/mean_std_filtering.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/mean_std_filtering.py new file mode 100644 index 0000000000000000000000000000000000000000..201ce369113897a7cff239d3f608b8137aa7bce0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/mean_std_filtering.py @@ -0,0 +1,193 @@ +"""Example using a ConnectorV2 for processing observations with a mean/std filter. + +An RLlib Algorithm has 3 distinct connector pipelines: +- An env-to-module pipeline in an EnvRunner accepting a list of episodes and producing +a batch for an RLModule to compute actions (`forward_inference()` or +`forward_exploration()`). +- A module-to-env pipeline in an EnvRunner taking the RLModule's output and converting +it into an action readable by the environment. +- A learner connector pipeline on a Learner taking a list of episodes and producing +a batch for an RLModule to perform the training forward pass (`forward_train()`). + +Each of these pipelines has a fixed set of default ConnectorV2 pieces that RLlib +adds/prepends to these pipelines in order to perform the most basic functionalities. +For example, RLlib adds the `AddObservationsFromEpisodesToBatch` ConnectorV2 into any +env-to-module pipeline to make sure the batch for computing actions contains - at the +minimum - the most recent observation. + +On top of these default ConnectorV2 pieces, users can define their own ConnectorV2 +pieces (or use the ones available already in RLlib) and add them to one of the 3 +different pipelines described above, as required. + +This example: + - shows how the `MeanStdFilter` ConnectorV2 piece can be added to the env-to-module + pipeline. + - demonstrates that using such a filter enhances learning behavior (or even makes + if possible to learn overall) in some environments, especially those with lopsided + observation spaces, for example `Box(-3000, -1000, ...)`. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +Running this example with the mean-std filter results in the normally expected Pendulum +learning behavior: ++-------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-------------------------------+------------+-----------------+--------+ +| PPO_lopsided-pend_f9c96_00000 | TERMINATED | 127.0.0.1:43612 | 77 | ++-------------------------------+------------+-----------------+--------+ ++------------------+------------------------+-----------------------+ +| total time (s) | num_env_steps_sample | episode_return_mean | +| | d_lifetime | | +|------------------+------------------------+-----------------------| +| 30.7466 | 40040 | -276.3 | ++------------------+------------------------+-----------------------+ + +If you try using the `--disable-mean-std-filter` (all other things being equal), you +will either see no learning progress at all (or a very slow one), but more likely some +numerical instability related error will be thrown: + +ValueError: Expected parameter loc (Tensor of shape (64, 1)) of distribution + Normal(loc: torch.Size([64, 1]), scale: torch.Size([64, 1])) to satisfy the + constraint Real(), but found invalid values: +tensor([[nan], + [nan], + [nan], + ... +""" +import gymnasium as gym +import numpy as np + +from ray.rllib.connectors.env_to_module.mean_std_filter import MeanStdFilter +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentPendulum +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +torch, _ = try_import_torch() + +parser = add_rllib_example_script_args( + default_iters=500, + default_timesteps=500000, + default_reward=-300.0, +) +parser.add_argument( + "--disable-mean-std-filter", + action="store_true", + help="Run w/o a mean/std env-to-module connector piece (filter).", +) + + +class LopsidedObs(gym.ObservationWrapper): + def __init__(self, env): + super().__init__(env) + self.observation_space = gym.spaces.Box(-4000.0, -1456.0, (3,), np.float32) + + def observation(self, observation): + # Lopside [-1.0, 1.0] Pendulum observations + return ((observation + 1.0) / 2.0) * (4000.0 - 1456.0) - 4000.0 + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "lopsided-pend", + lambda _: MultiAgentPendulum(config={"num_agents": args.num_agents}), + ) + else: + register_env("lopsided-pend", lambda _: LopsidedObs(gym.make("Pendulum-v1"))) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("lopsided-pend") + .env_runners( + num_envs_per_env_runner=1 if args.num_agents > 0 else 20, + # Define a single connector piece to be prepended to the env-to-module + # connector pipeline. + # Alternatively, return a list of n ConnectorV2 pieces (which will then be + # included in an automatically generated EnvToModulePipeline or return a + # EnvToModulePipeline directly. + env_to_module_connector=( + None + if args.disable_mean_std_filter + else lambda env, spaces, device: ( + MeanStdFilter(multi_agent=args.num_agents > 0) + ) + ), + ) + .training( + train_batch_size_per_learner=512, + gamma=0.95, + # Linearly adjust learning rate based on number of GPUs. + lr=0.0003 * (args.num_learners or 1), + vf_loss_coeff=0.01, + ) + .rl_module( + model_config=DefaultModelConfig( + fcnet_activation="relu", + fcnet_kernel_initializer=torch.nn.init.xavier_uniform_, + fcnet_bias_initializer=torch.nn.init.constant_, + fcnet_bias_initializer_kwargs={"val": 0.0}, + ), + ) + # In case you would like to run with a evaluation EnvRunners, make sure your + # `evaluation_config` key contains the `use_worker_filter_stats=False` setting + # (see below). This setting makes sure that the mean/std stats collected by the + # evaluation EnvRunners are NOT used for the training EnvRunners (unless you + # really want to mix these stats). It's normally a good idea to keep the stats + # collected during evaluation completely out of the training data (already for + # better reproducibility alone). + # .evaluation( + # evaluation_num_env_runners=1, + # evaluation_interval=1, + # evaluation_config={ + # "explore": False, + # # Do NOT use the eval EnvRunners' ConnectorV2 states. Instead, before + # # each round of evaluation, broadcast the latest training + # # EnvRunnerGroup's ConnectorV2 states (merged from all training remote + # # EnvRunners) to the eval EnvRunnerGroup (and discard the eval + # # EnvRunners' stats). + # "use_worker_filter_stats": False, + # }, + # ) + ) + + # PPO specific settings. + if args.algo == "PPO": + base_config.training( + minibatch_size=64, + lambda_=0.1, + vf_clip_param=10.0, + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/multi_agent_observation_preprocessor.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/multi_agent_observation_preprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..1fbf29de0267815dc0b18f0413f4b85010b2d9a8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/multi_agent_observation_preprocessor.py @@ -0,0 +1,138 @@ +"""Example of a ConnectorV2 mapping global observations to n per-module observations. + +An RLlib Algorithm has 3 distinct connector pipelines: +- An env-to-module pipeline in an EnvRunner accepting a list of episodes and producing +a batch for an RLModule to compute actions (`forward_inference()` or +`forward_exploration()`). +- A module-to-env pipeline in an EnvRunner taking the RLModule's output and converting +it into an action readable by the environment. +- A learner connector pipeline on a Learner taking a list of episodes and producing +a batch for an RLModule to perform the training forward pass (`forward_train()`). + +Each of these pipelines has a fixed set of default ConnectorV2 pieces that RLlib +adds/prepends to these pipelines in order to perform the most basic functionalities. +For example, RLlib adds the `AddObservationsFromEpisodesToBatch` ConnectorV2 into any +env-to-module pipeline to make sure the batch for computing actions contains - at the +minimum - the most recent observation. + +On top of these default ConnectorV2 pieces, users can define their own ConnectorV2 +pieces (or use the ones available already in RLlib) and add them to one of the 3 +different pipelines described above, as required. + +This example: + - shows how the custom `AddOtherAgentsRowIndexToXYPos` ConnectorV2 piece can be + added to the env-to-module pipeline. It serves as a multi-agent observation + preprocessor and makes sure than both agents' observations contain necessary + information about the respective other agent. Without this extra information, the + agents won't be able to learn to solve the problem optimally. + - demonstrates that using various such observation mapping connector pieces allows + users to map from global, multi-agent observations to individual modules' + observations. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see the algo reach an episode return of slightly above 20.0, which proves +that both agents learn how to utilize the other agents' row-index (0 or 1) in order +to collide with the other agent and receive an extra +5 reward. Without this collision +during the episode (if one agent reaches its goal, it's removed from the scene and no +collision can occur any longer), the maximum return per agent is under 10.0. + ++--------------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|--------------------------------------+------------+-----------------+--------+ +| PPO_DoubleRowCorridorEnv_ba678_00000 | TERMINATED | 127.0.0.1:73310 | 37 | ++--------------------------------------+------------+-----------------+--------+ ++------------------+-------+-------------------+-------------+-------------+ +| total time (s) | ts | combined return | return p1 | return p0 | +|------------------+-------+-------------------+-------------+-------------| +| 41.5389 | 19998 | 23.072 | 11.418 | 11.654 | ++------------------+-------+-------------------+-------------+-------------+ +""" +from ray.rllib.connectors.env_to_module.flatten_observations import ( + FlattenObservations, +) +from ray.rllib.examples.connectors.classes.add_other_agents_row_index_to_xy_pos import ( + AddOtherAgentsRowIndexToXYPos, +) +from ray.rllib.examples.envs.classes.multi_agent.double_row_corridor_env import ( + DoubleRowCorridorEnv, +) +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +torch, _ = try_import_torch() + +parser = add_rllib_example_script_args( + default_iters=200, + default_timesteps=200000, + default_reward=22.0, +) +parser.set_defaults( + num_agents=2, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment(DoubleRowCorridorEnv) + .env_runners( + num_envs_per_env_runner=20, + # Define a list of two connector piece to be prepended to the env-to-module + # connector pipeline: + # 1) The custom connector piece: A MultiAgentObservationPreprocessor, which + # enhances each agents' individual observations through adding the + # respective other agent's row index to the observation. + # 2) A FlattenObservations connector to flatten the integer observations + # for `agent_0`, which the AddOtherAgentsRowIndexToXYPos outputs. + env_to_module_connector=lambda env, spaces, device: [ + AddOtherAgentsRowIndexToXYPos(), + # Only flatten agent_0's observations (b/c these are ints that need to + # be one-hot'd). + FlattenObservations(multi_agent=True, agent_ids=["agent_0"]), + ], + ) + .training( + train_batch_size_per_learner=512, + gamma=0.95, + # Linearly adjust learning rate based on number of GPUs. + lr=0.0003 * (args.num_learners or 1), + vf_loss_coeff=0.01, + ) + .multi_agent( + policies={"p0", "p1"}, + policy_mapping_fn=lambda aid, eps, **kw: "p0" if aid == "agent_0" else "p1", + ) + ) + + # PPO specific settings. + if args.algo == "PPO": + base_config.training( + minibatch_size=64, + lambda_=0.1, + vf_clip_param=10.0, + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/prev_actions_prev_rewards.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/prev_actions_prev_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..439956c7eceb2d9c039d381ce7e6ecad5cf240e7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/prev_actions_prev_rewards.py @@ -0,0 +1,163 @@ +"""Example using a ConnectorV2 to add previous rewards/actions to an RLModule's input. + +An RLlib Algorithm has 3 distinct connector pipelines: +- An env-to-module pipeline in an EnvRunner accepting a list of episodes and producing +a batch for an RLModule to compute actions (`forward_inference()` or +`forward_exploration()`). +- A module-to-env pipeline in an EnvRunner taking the RLModule's output and converting +it into an action readable by the environment. +- A learner connector pipeline on a Learner taking a list of episodes and producing +a batch for an RLModule to perform the training forward pass (`forward_train()`). + +Each of these pipelines has a fixed set of default ConnectorV2 pieces that RLlib +adds/prepends to these pipelines in order to perform the most basic functionalities. +For example, RLlib adds the `AddObservationsFromEpisodesToBatch` ConnectorV2 into any +env-to-module pipeline to make sure the batch for computing actions contains - at the +minimum - the most recent observation. + +On top of these default ConnectorV2 pieces, users can define their own ConnectorV2 +pieces (or use the ones available already in RLlib) and add them to one of the 3 +different pipelines described above, as required. + +This example: + - shows how the `PrevActionsPrevRewards` ConnectorV2 piece can be added to the + env-to-module pipeline to extract previous rewards and/or actions from the ongoing + episodes. + - shows how this connector creates and wraps this new information (rewards and + actions) together with the original observations into the RLModule's input dict + under a new `gym.spaces.Dict` structure (for example, if your observation space + is `O=Box(shape=(3,))` and you add the most recent 1 reward, the new observation + space will be `Dict({"_original_obs": O, "prev_n_rewards": Box(shape=())})`. + - demonstrates how to use RLlib's `FlattenObservations` right after the + `PrevActionsPrevRewards` to flatten that new dict observation structure again into + a single 1D tensor. + - uses the StatelessCartPole environment, a CartPole-v1 derivative that's missing + both x-veloc and angle-veloc observation components and is therefore non-Markovian + (only partially observable). An LSTM default model is used for training. Adding + the additional context to the observations (for example, prev. actions) helps the + LSTM to more quickly learn in this environment. + + +How to run this script +---------------------- +`python [script file name].py --num-frames=4 --env=ALE/Pong-v5` + +Use the `--num-frames` option to define the number of observations to framestack. +If you don't want to use Connectors to perform the framestacking, set the +`--use-gym-wrapper-framestacking` flag to perform framestacking already inside a +gymnasium observation wrapper. In this case though, be aware that the tensors being +sent through the network are `--num-frames` x larger than if you use the Connector +setup. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- + +You should see something similar to this in your terminal output when running +ths script as described above: + ++---------------------+------------+-----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|---------------------+------------+-----------------+--------+------------------+ +| PPO_env_0edd2_00000 | TERMINATED | 127.0.0.1:12632 | 17 | 42.6898 | ++---------------------+------------+-----------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| num_env_steps_sample | num_env_steps_traine | episode_return_mean | +| d_lifetime | d_lifetime | | +|------------------------+------------------------+------------------------| +| 68000 | 68000 | 205.22 | ++------------------------+------------------------+------------------------+ +""" +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.connectors.env_to_module import ( + FlattenObservations, + PrevActionsPrevRewards, +) +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentStatelessCartPole +from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune import register_env + +torch, nn = try_import_torch() + + +parser = add_rllib_example_script_args( + default_reward=200.0, default_timesteps=1000000, default_iters=2000 +) +parser.add_argument("--n-prev-rewards", type=int, default=1) +parser.add_argument("--n-prev-actions", type=int, default=1) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Define our custom connector pipelines. + def _env_to_module(env, spaces, device): + # Create the env-to-module connector pipeline. + return [ + PrevActionsPrevRewards( + multi_agent=args.num_agents > 0, + n_prev_rewards=args.n_prev_rewards, + n_prev_actions=args.n_prev_actions, + ), + FlattenObservations(multi_agent=args.num_agents > 0), + ] + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "env", + lambda _: MultiAgentStatelessCartPole( + config={"num_agents": args.num_agents} + ), + ) + else: + register_env("env", lambda _: StatelessCartPole()) + + config = ( + PPOConfig() + .environment("env") + .env_runners(env_to_module_connector=_env_to_module) + .training( + num_epochs=6, + lr=0.0003, + train_batch_size=4000, + vf_loss_coeff=0.01, + ) + .rl_module( + model_config=DefaultModelConfig( + use_lstm=True, + max_seq_len=20, + fcnet_hiddens=[32], + fcnet_activation="linear", + fcnet_kernel_initializer=nn.init.xavier_uniform_, + fcnet_bias_initializer=nn.init.constant_, + fcnet_bias_initializer_kwargs={"val": 0.0}, + vf_share_layers=True, + ), + ) + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + config = config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + run_rllib_example_script_experiment(config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/connectors/single_agent_observation_preprocessor.py b/lib/python3.12/site-packages/ray/rllib/examples/connectors/single_agent_observation_preprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..f54bfce19066d952f911c5933197a65ae5d1c11b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/connectors/single_agent_observation_preprocessor.py @@ -0,0 +1,154 @@ +"""Example using a `SingleAgentObservationPreprocessor` to preprocess observations. + +The custom preprocessor here is part of the env-to-module connector pipeline and +alters the CartPole-v1 environment observations from the Markovian 4-tuple (x-pos, +angular-pos, x-velocity, angular-velocity) to a non-Markovian, simpler 2-tuple (only +x-pos and angular-pos). The resulting problem can only be solved through a +memory/stateful model, for example an LSTM. + +An RLlib Algorithm has 3 distinct connector pipelines: +- An env-to-module pipeline in an EnvRunner accepting a list of episodes and producing +a batch for an RLModule to compute actions (`forward_inference()` or +`forward_exploration()`). +- A module-to-env pipeline in an EnvRunner taking the RLModule's output and converting +it into an action readable by the environment. +- A learner connector pipeline on a Learner taking a list of episodes and producing +a batch for an RLModule to perform the training forward pass (`forward_train()`). + +Each of these pipelines has a fixed set of default ConnectorV2 pieces that RLlib +adds/prepends to these pipelines in order to perform the most basic functionalities. +For example, RLlib adds the `AddObservationsFromEpisodesToBatch` ConnectorV2 into any +env-to-module pipeline to make sure the batch for computing actions contains - at the +minimum - the most recent observation. + +On top of these default ConnectorV2 pieces, users can define their own ConnectorV2 +pieces (or use the ones available already in RLlib) and add them to one of the 3 +different pipelines described above, as required. + +This example: + - shows how to write a custom `SingleAgentObservationPreprocessor` ConnectorV2 + piece. + - shows how to add this custom class to the env-to-module pipeline through the + algorithm config. + - demonstrates that by using this connector, the normal CartPole observation + changes from a Markovian (fully observable) to a non-Markovian (partially + observable) observation. Only stateful, memory enhanced models can solve the + resulting RL problem. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- + +You should see something like this at the end in your console output. +Note that your setup wouldn't be able to solve the environment, preprocessed through +your custom `SingleAgentObservationPreprocessor`, without the help of the configured +LSTM since you convert the env from a Markovian one to a partially observable, +non-Markovian one. ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_0ecb5_00000 | TERMINATED | 127.0.0.1:57921 | 9 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | episode_return_mean | num_env_steps_sample | +| | | d_lifetime | +|------------------+------------------------+------------------------| +| 26.2305 | 224.38 | 36000 | ++------------------+------------------------+------------------------+ +""" +import gymnasium as gym +import numpy as np + +from ray.rllib.connectors.env_to_module.observation_preprocessor import ( + SingleAgentObservationPreprocessor, +) +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +# Read in common example script command line arguments. +parser = add_rllib_example_script_args(default_timesteps=200000, default_reward=200.0) + + +class ReduceCartPoleObservationsToNonMarkovian(SingleAgentObservationPreprocessor): + def recompute_output_observation_space( + self, + input_observation_space: gym.Space, + input_action_space: gym.Space, + ) -> gym.Space: + # The new observation space only has a shape of (2,), not (4,). + return gym.spaces.Box( + -5.0, + 5.0, + (input_observation_space.shape[0] - 2,), + np.float32, + ) + + def preprocess(self, observation, episode: SingleAgentEpisode): + # Extract only the positions (x-position and angular-position). + return np.array([observation[0], observation[2]], np.float32) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Define the AlgorithmConfig used. + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + # You use the normal CartPole-v1 env here and your env-to-module preprocessor + # converts this into a non-Markovian version of CartPole. + .environment("CartPole-v1") + .env_runners( + env_to_module_connector=( + lambda env, spaces, device: ReduceCartPoleObservationsToNonMarkovian() + ), + ) + .training( + gamma=0.99, + lr=0.0003, + ) + .rl_module( + model_config=DefaultModelConfig( + # Solve the non-Markovian env through using an LSTM-enhanced model. + use_lstm=True, + vf_share_layers=True, + ), + ) + ) + + # PPO-specific settings (for better learning behavior only). + if args.algo == "PPO": + base_config.training( + num_epochs=6, + vf_loss_coeff=0.01, + ) + # IMPALA-specific settings (for better learning behavior only). + elif args.algo == "IMPALA": + base_config.training( + lr=0.0005, + vf_loss_coeff=0.05, + entropy_coeff=0.0, + ) + + # Run everything as configured. + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f5eb10670992ce448f932571b142d7997a59b44 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/count_based_curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/count_based_curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..640f2d5937c585e3dbebcff20cb6c3879b80efa6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/count_based_curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1064ddb673488112633308a41524943b99e150f3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/euclidian_distance_based_curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/intrinsic_curiosity_model_based_curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/intrinsic_curiosity_model_based_curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..99a73ce9de07bff1c027d67c989a34ae275d63e0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__pycache__/intrinsic_curiosity_model_based_curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/count_based_curiosity.py b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/count_based_curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..7f2900b42a7dd7227d65e3122cc5748b326ad1b5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/count_based_curiosity.py @@ -0,0 +1,136 @@ +"""Example of using a count-based curiosity mechanism to learn in sparse-rewards envs. + +This example: + - demonstrates how to define your own count-based curiosity ConnectorV2 piece + that computes intrinsic rewards based on simple observation counts and adds these + intrinsic rewards to the "main" (extrinsic) rewards. + - shows how this connector piece overrides the main (extrinsic) rewards in the + episode and thus demonstrates how to do reward shaping in general with RLlib. + - shows how to plug this connector piece into your algorithm's config. + - uses Tune and RLlib to learn the env described above and compares 2 + algorithms, one that does use curiosity vs one that does not. + +We use a FrozenLake (sparse reward) environment with a map size of 8x8 and a time step +limit of 14 to make it almost impossible for a non-curiosity based policy to learn. + + +How to run this script +---------------------- +`python [script file name].py` + +Use the `--no-curiosity` flag to disable curiosity learning and force your policy +to be trained on the task w/o the use of intrinsic rewards. With this option, the +algorithm should NOT succeed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that only a PPO policy that uses curiosity can +actually learn. + +Policy using count-based curiosity: ++-------------------------------+------------+--------+------------------+ +| Trial name | status | iter | total time (s) | +| | | | | +|-------------------------------+------------+--------+------------------+ +| PPO_FrozenLake-v1_109de_00000 | TERMINATED | 48 | 44.46 | ++-------------------------------+------------+--------+------------------+ ++------------------------+-------------------------+------------------------+ +| episode_return_mean | num_episodes_lifetime | num_env_steps_traine | +| | | d_lifetime | +|------------------------+-------------------------+------------------------| +| 0.99 | 12960 | 194000 | ++------------------------+-------------------------+------------------------+ + +Policy NOT using curiosity: +[DOES NOT LEARN AT ALL] +""" +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.connectors.classes.count_based_curiosity import ( + CountBasedCuriosity, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +parser = add_rllib_example_script_args( + default_reward=0.99, default_iters=200, default_timesteps=1000000 +) +parser.add_argument( + "--intrinsic-reward-coeff", + type=float, + default=1.0, + help="The weight with which to multiply intrinsic rewards before adding them to " + "the extrinsic ones (default is 1.0).", +) +parser.add_argument( + "--no-curiosity", + action="store_true", + help="Whether to NOT use count-based curiosity.", +) + +ENV_OPTIONS = { + "is_slippery": False, + # Use this hard-to-solve 8x8 map with lots of holes (H) to fall into and only very + # few valid paths from the starting state (S) to the goal state (G). + "desc": [ + "SFFHFFFH", + "FFFHFFFF", + "FFFHHFFF", + "FFFFFFFH", + "HFFHFFFF", + "HHFHFFHF", + "FFFHFHHF", + "FHFFFFFG", + ], + # Limit the number of steps the agent is allowed to make in the env to + # make it almost impossible to learn without (count-based) curiosity. + "max_episode_steps": 14, +} + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + "FrozenLake-v1", + env_config=ENV_OPTIONS, + ) + .env_runners( + num_envs_per_env_runner=5, + # Flatten discrete observations (into one-hot vectors). + env_to_module_connector=lambda env, spaces, device: FlattenObservations(), + ) + .training( + # The main code in this example: We add the `CountBasedCuriosity` connector + # piece to our Learner connector pipeline. + # This pipeline is fed with collected episodes (either directly from the + # EnvRunners in on-policy fashion or from a replay buffer) and converts + # these episodes into the final train batch. The added piece computes + # intrinsic rewards based on simple observation counts and add them to + # the "main" (extrinsic) rewards. + learner_connector=( + None if args.no_curiosity else lambda *ags, **kw: CountBasedCuriosity() + ), + num_epochs=10, + vf_loss_coeff=0.01, + ) + .rl_module(model_config=DefaultModelConfig(vf_share_layers=True)) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/euclidian_distance_based_curiosity.py b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/euclidian_distance_based_curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..4e54c8ad7287e2e9f33d3deaae2413e41c12b488 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/euclidian_distance_based_curiosity.py @@ -0,0 +1,126 @@ +"""Example of a euclidian-distance curiosity mechanism to learn in sparse-rewards envs. + +This example: + - demonstrates how to define your own euclidian-distance-based curiosity ConnectorV2 + piece that computes intrinsic rewards based on the delta between incoming + observations and some set of already stored (prior) observations. Thereby, the + further away the incoming observation is from the already stored ones, the higher + its corresponding intrinsic reward. + - shows how this connector piece adds the intrinsic reward to the corresponding + "main" (extrinsic) reward and overrides the value in the "rewards" key in the + episode. It thus demonstrates how to do reward shaping in general with RLlib. + - shows how to plug this connector piece into your algorithm's config. + - uses Tune and RLlib to learn the env described above and compares 2 + algorithms, one that does use curiosity vs one that does not. + +We use the MountainCar-v0 environment, a sparse-reward env that is very hard to learn +for a regular PPO algorithm. + + +How to run this script +---------------------- +`python [script file name].py` + +Use the `--no-curiosity` flag to disable curiosity learning and force your policy +to be trained on the task w/o the use of intrinsic rewards. With this option, the +algorithm should NOT succeed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that only a PPO policy that uses curiosity can +actually learn. + +Policy using count-based curiosity: ++-------------------------------+------------+--------+------------------+ +| Trial name | status | iter | total time (s) | +| | | | | +|-------------------------------+------------+--------+------------------+ +| PPO_FrozenLake-v1_109de_00000 | TERMINATED | 48 | 44.46 | ++-------------------------------+------------+--------+------------------+ ++------------------------+-------------------------+------------------------+ +| episode_return_mean | num_episodes_lifetime | num_env_steps_traine | +| | | d_lifetime | +|------------------------+-------------------------+------------------------| +| 0.99 | 12960 | 194000 | ++------------------------+-------------------------+------------------------+ + +Policy NOT using curiosity: +[DOES NOT LEARN AT ALL] +""" +from ray.rllib.connectors.env_to_module import MeanStdFilter +from ray.rllib.examples.connectors.classes.euclidian_distance_based_curiosity import ( + EuclidianDistanceBasedCuriosity, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +# TODO (sven): SB3's PPO learns MountainCar-v0 until a reward of ~-110. +# We might have to play around some more with different initializations, etc.. +# to get to these results as well. +parser = add_rllib_example_script_args( + default_reward=-140.0, default_iters=2000, default_timesteps=1000000 +) +parser.set_defaults( + num_env_runners=4, +) +parser.add_argument( + "--intrinsic-reward-coeff", + type=float, + default=0.0001, + help="The weight with which to multiply intrinsic rewards before adding them to " + "the extrinsic ones (default is 0.0001).", +) +parser.add_argument( + "--no-curiosity", + action="store_true", + help="Whether to NOT use count-based curiosity.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("MountainCar-v0") + .env_runners( + env_to_module_connector=lambda env, spaces, device: MeanStdFilter(), + num_envs_per_env_runner=5, + ) + .training( + # The main code in this example: We add the + # `EuclidianDistanceBasedCuriosity` connector piece to our Learner connector + # pipeline. This pipeline is fed with collected episodes (either directly + # from the EnvRunners in on-policy fashion or from a replay buffer) and + # converts these episodes into the final train batch. The added piece + # computes intrinsic rewards based on simple observation counts and add them + # to the "main" (extrinsic) rewards. + learner_connector=( + None + if args.no_curiosity + else lambda *ags, **kw: EuclidianDistanceBasedCuriosity() + ), + # train_batch_size_per_learner=512, + grad_clip=20.0, + entropy_coeff=0.003, + gamma=0.99, + lr=0.0002, + lambda_=0.98, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68c75d9d2b1adc3d67b4112726f215122e2c3821 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/curriculum_learning.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/curriculum_learning.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..308a0c8c5cded436d16b4c87d3319cb79242a59e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/curriculum_learning.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/pong_curriculum_learning.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/pong_curriculum_learning.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..096c030faa15c273a1633a4cf3e7561c71d25d87 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/__pycache__/pong_curriculum_learning.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curriculum/curriculum_learning.py b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/curriculum_learning.py new file mode 100644 index 0000000000000000000000000000000000000000..db8fd311ae31936d48fe8b553b8b3320bf924adf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/curriculum_learning.py @@ -0,0 +1,249 @@ +"""Example of using an env-task curriculum by implementing a custom callback. + +This example: + - demonstrates how to define your own curriculum-capable environments using + gymnasium's FrozenLake env. + - defines a custom callback that gets called once per iteration and - if necessary - + changes the maps used by FrozenLake on all EnvRunners to a new task (by moving the + goal position further and further away from the starting position). + - also demonstrates an alternative approach via reloading/recreating an entirely new + env inside all EnvRunners. + - uses Tune and RLlib to curriculum-learn the env described above and compares 2 + algorithms, one that does use curriculum learning vs one that does not. + +We use a FrozenLake (sparse reward) environment with a map size of 8x8 and a time step +limit of 16 to make it almost impossible for a non-curriculum policy to learn. + + +How to run this script +---------------------- +`python [script file name].py` + +Use the `--no-curriculum` flag to disable curriculum learning and force your policy +to be trained on the hardest task right away. With this option, the algorithm should NOT +succeed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that only PPO policy that uses a curriculum can +actually learn, whereas the one that is thrown into the toughest task right from the +start never learns anything. + +Policy using the curriculum: ++-------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|-------------------------------+------------+-----------------+--------+ +| PPO_FrozenLake-v1_93ca4_00000 | TERMINATED | 127.0.0.1:73318 | 41 | ++-------------------------------+------------+-----------------+--------+ ++------------------+--------+----------+--------------------+ +| total time (s) | ts | reward | episode_len_mean | +|------------------+--------+----------+--------------------| +| 97.652 | 164000 | 1 | 14.0348 | ++------------------+--------+----------+--------------------+ + +Policy NOT using the curriculum (trying to solve the hardest task right away): +[DOES NOT LEARN AT ALL] +""" +from functools import partial + +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls +from ray.tune.result import TRAINING_ITERATION + +parser = add_rllib_example_script_args(default_iters=100, default_timesteps=600000) +parser.add_argument( + "--upgrade-task-threshold", + type=float, + default=0.99, + help="The mean episode return, upon reaching of which we increase the task by one.", +) +parser.add_argument( + "--no-curriculum", + action="store_true", + help="Whether to NOT use curriculum learning (and instead trying to solve the " + "hardest task right away).", +) + +# __curriculum_learning_example_env_options__ +ENV_OPTIONS = { + "is_slippery": False, + # Limit the number of steps the agent is allowed to make in the env to + # make it almost impossible to learn without the curriculum. + "max_episode_steps": 16, +} + +# Our 3 tasks: 0=easiest, 1=medium, 2=hard +ENV_MAPS = [ + # 0 + [ + "SFFHFFFH", + "FFFHFFFF", + "FFGFFFFF", + "FFFFFFFF", + "HFFFFFFF", + "HHFFFFHF", + "FFFFFHHF", + "FHFFFFFF", + ], + # 1 + [ + "SFFHFFFH", + "FFFHFFFF", + "FFFFFFFF", + "FFFFFFFF", + "HFFFFFFF", + "HHFFGFHF", + "FFFFFHHF", + "FHFFFFFF", + ], + # 2 + [ + "SFFHFFFH", + "FFFHFFFF", + "FFFFFFFF", + "FFFFFFFF", + "HFFFFFFF", + "HHFFFFHF", + "FFFFFHHF", + "FHFFFFFG", + ], +] +# __END_curriculum_learning_example_env_options__ + + +# Simple function sent to an EnvRunner to change the map of all its gym. Envs from +# the current one to a new (tougher) one, in which the goal position is further away +# from the starting position. Note that a map is a list of strings, each one +# representing one row in the map. Each character in the strings represent a single +# field (S=starting position, H=hole (bad), F=frozen/free field (ok), G=goal (great!)). +def _remote_fn(env_runner, new_task: int): + # We recreate the entire env object by changing the env_config on the worker, + # then calling its `make_env()` method. + env_runner.config.environment(env_config={"desc": ENV_MAPS[new_task]}) + env_runner.make_env() + + +class EnvTaskCallback(RLlibCallback): + """Custom callback implementing `on_train_result()` for changing the envs' maps.""" + + def on_algorithm_init( + self, + *, + algorithm: "Algorithm", + **kwargs, + ) -> None: + # Set the initial task to 0. + algorithm._counters["current_env_task"] = 0 + + def on_train_result( + self, + *, + algorithm: Algorithm, + metrics_logger=None, + result: dict, + **kwargs, + ) -> None: + # Hack: Store the current task inside a counter in our Algorithm. + # W/o a curriculum, the task is always 2 (hardest). + if args.no_curriculum: + algorithm._counters["current_env_task"] = 2 + current_task = algorithm._counters["current_env_task"] + + # If episode return is consistently `args.upgrade_task_threshold`, we switch + # to a more difficult task (if possible). If we already mastered the most + # difficult task, we publish our victory in the result dict. + result["task_solved"] = 0.0 + current_return = result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + if current_return > args.upgrade_task_threshold: + if current_task < 2: + new_task = current_task + 1 + print( + f"Switching task/map on all EnvRunners to #{new_task} (0=easiest, " + f"2=hardest), b/c R={current_return} on current task." + ) + algorithm.env_runner_group.foreach_env_runner( + func=partial(_remote_fn, new_task=new_task) + ) + algorithm._counters["current_env_task"] = new_task + + # Hardest task was solved (1.0) -> report this in the results dict. + elif current_return == 1.0: + result["task_solved"] = 1.0 + # Emergency brake: If return is 0.0 AND we are already at a harder task (1 or + # 2), we go back to task=0. + elif current_return == 0.0 and current_task > 0: + print( + "Emergency brake: Our policy seemed to have collapsed -> Setting task " + "back to 0." + ) + algorithm.env_runner_group.foreach_env_runner( + func=partial(_remote_fn, new_task=0) + ) + algorithm._counters["current_env_task"] = 0 + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + # Plug in our curriculum callbacks that controls when we should upgrade the env + # task based on the received return for the current task. + .callbacks(EnvTaskCallback) + .environment( + "FrozenLake-v1", + env_config={ + # w/ curriculum: start with task=0 (easiest) + # w/o curriculum: start directly with hardest task 2. + "desc": ENV_MAPS[2 if args.no_curriculum else 0], + **ENV_OPTIONS, + }, + ) + .env_runners( + num_envs_per_env_runner=5, + env_to_module_connector=lambda env, spaces, device: FlattenObservations(), + ) + .training( + num_epochs=6, + vf_loss_coeff=0.01, + lr=0.0002, + ) + .rl_module(model_config=DefaultModelConfig(vf_share_layers=True)) + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + # Reward directly does not matter to us as we would like to continue + # after the policy reaches a return of ~1.0 on the 0-task (easiest). + # But we DO want to stop, once the entire task is learned (policy achieves + # return of 1.0 on the most difficult task=2). + "task_solved": 1.0, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + } + + run_rllib_example_script_experiment( + base_config, args, stop=stop, success_metric={"task_solved": 1.0} + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curriculum/pong_curriculum_learning.py b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/pong_curriculum_learning.py new file mode 100644 index 0000000000000000000000000000000000000000..e03e478717a59f04758cefdd77fb2d40fd665951 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/curriculum/pong_curriculum_learning.py @@ -0,0 +1,281 @@ +"""Example of using curriculum learning for Atari Pong by implementing a custom callback. + +This example: + - demonstrates how to define a curriculum for an agent playing gymnasium's Atari + Pong. + - defines a custom callback that gets called once per iteration and - if the agent + performs well enough - increases the task difficulty, i.e. the `frameskip` for all + environments on all EnvRunners (the agent must act now faster). + - also demonstrates how to provide the callback with varying curriculum parameters + (like threshold maps, returns at which the curriculum ends, etc.). + - uses Ray Tune and RLlib to curriculum-learn Atari Pong with a high frameskip. + +We use Atari Pong with a framestack of 4 images (i.e. observation dimensions of 64x64x4) +and start with a frameskip of 1. At a return of 15.0 we increase the frameskip to 2, at +a return of 17.0 to 3, at 19.0 to 4, and the task is solved at a frameskip of 21.0. + +How to run this script +---------------------- +`python [script file name].py` + +Use the `--solved-return` flag to define the threshold at which curriculum learning ends. +Note that a PPO agent on Atari Pong will need a long time to learn. + +To ensure the agent has not collapsed, but rather made had a bad seed, we only decrease +the frameskip when the agent performed worse than the next lower threshold. The margin by +which the agent has to be worse is defined by the `--demotion-margin` argument and defaults +to 2.0. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +""" + +import functools +from typing import Callable + +import gymnasium as gym + +from ray import tune +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.connectors.env_to_module.frame_stacking import FrameStackingEnvToModule +from ray.rllib.connectors.learner.frame_stacking import FrameStackingLearner +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack +from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS, EPISODE_RETURN_MEAN +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +parser = add_rllib_example_script_args( + default_reward=float("inf"), + default_timesteps=3000000, + default_iters=100000000000, +) +parser.set_defaults( + env="ale_py:ALE/Pong-v5", +) +parser.add_argument( + "--solved-return", + type=float, + default=21.0, + help=("The mean episode return at which we consider the task to be fully solved."), +) +parser.add_argument( + "--demotion-margin", + type=float, + default=2.0, + help=( + "The margin below the next lower task threshold, beneath which the agent " + " is considered to have collapsed, prompting a downgrade of the task." + ), +) +# Use `parser` to add your own custom command line options to this script +# and (if needed) use their values to set up `config` below. +args = parser.parse_args() + +NUM_LEARNERS = args.num_learners or 1 +ENV = args.env + + +class PongEnvTaskCallback(RLlibCallback): + """Custom callback changing the frameskip in Atari Pong dependent on return.""" + + def __init__( + self, + task_threshold_map: dict, + remote_fn: Callable, + demotion_margin: float = 0.0, + solved_return: float = float("inf"), + ): + self.task_threshold_map = task_threshold_map + self.remote_fn = remote_fn + self.demotion_margin = demotion_margin + self.solved_return = solved_return + + def on_algorithm_init( + self, + *, + algorithm: "Algorithm", + **kwargs, + ) -> None: + # Set the initial task to 1, which corresponds to a frameskip of 1. + algorithm.metrics.log_value("current_env_task", 1, reduce="sum") + + def on_train_result( + self, + *, + algorithm: Algorithm, + metrics_logger=None, + result: dict, + **kwargs, + ) -> None: + # Store the current task inside the metrics logger in our Algorithm. + current_task = metrics_logger.peek("current_env_task") + + # If episode return is consistently above `task_threshold_map[current_task]`, + # we switch to a more difficult task (i.e. higher `frameskip`` if possible). + # If we already mastered the most difficult task, we publish our victory in + # the result dict. + result["task_solved"] = 0.0 + + # Note, in the first callback executions there may be no completed episode + # (and therefore no episode return) reported. In this case we will skip the + # the logic to manage task difficulty. + if EPISODE_RETURN_MEAN in result[ENV_RUNNER_RESULTS]: + current_return = result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + else: + return + + # Get the threshold of the current task from the threshold map. + threshold = self.task_threshold_map.get(current_task, float("inf")) + + # Check, if curriculum is solved. + final_task = max(self.task_threshold_map.keys()) + if current_task == final_task and current_return >= self.solved_return: + # Hardest task was solved -> report this in the results dict. + result["task_solved"] = 1.0 + + # Check promotion (increasing task). Note, we could use here also a promotion_patience + # that ensures that the return is collected in a stable manner instead of a lucky shot. + if ( + current_return >= threshold + ): # & result[ENV_RUNNER_RESULTS][NUM_EPISODES] > promotion_patience. + next_task = current_task + 1 + if next_task in self.task_threshold_map: + print( + f"Switching task on all EnvRunners up to #{next_task} (1=easiest, " + f"4=hardest), b/c R={current_return} on current task." + ) + # Increase task. + algorithm.env_runner_group.foreach_env_runner( + func=functools.partial(self.remote_fn, new_task=next_task) + ) + metrics_logger.log_value("current_env_task", next_task, window=1) + + # Check demotion (decreasing task). The demotion is used to avoid decreasing the task + # in case of an unlucky episode run. Only if the return is singificantly lower we + # decrease the task. + previous_task = current_task - 1 + if previous_task in self.task_threshold_map: + previous_threshold = self.task_threshold_map[previous_task] + if current_return < previous_threshold - self.demotion_margin: + print( + f"Switching task on all EnvRunners back to #{previous_task} (1=easiest, " + f"4=hardest), b/c R={current_return} on current task." + ) + # Decrease to previous level. + algorithm.env_runner_group.foreach_env_runner( + func=functools.partial(self.remote_fn, new_task=previous_task) + ) + metrics_logger.log_value("current_env_task", previous_task, window=1) + + +# These tags allow extracting portions of this script on Anyscale. +# ws-template-code-start +def _make_env_to_module_connector(env, spaces, device): + return FrameStackingEnvToModule(num_frames=4) + + +def _make_learner_connector(input_observation_space, input_action_space): + return FrameStackingLearner(num_frames=4) + + +# Create a custom Atari setup (w/o the usual RLlib-hard-coded framestacking in it). +# We would like our frame stacking connector to do this job. +def _env_creator(cfg): + return wrap_atari_for_new_api_stack( + gym.make(ENV, **cfg, render_mode="rgb_array"), + # Perform frame-stacking through ConnectorV2 API. + framestack=None, + ) + + +# Simple function sent to an EnvRunner to change the map of all its gym. Envs from +# the current one to a new (tougher) one, in which the frameskip is higher +# and the agent must therefore act faster. +def _remote_fn(env_runner, new_task: int): + # Override the env_config with the new setting. + env_runner.config.env_config.update( + { + "frameskip": new_task, + } + ) + # We recreate the entire env object by changing the env_config on the worker, + # then calling its `make_env()` method. + env_runner.make_env() + + +# Task threshold map keeps track of thresholds for each task. If the threshold has +# been surpassed the task difficulty is increased. +task_threshold_map = { + # Frameskip: Return. + 1: 15.0, + 2: 17.0, + 3: 19.0, + 4: float("inf"), +} + +tune.register_env("env", _env_creator) + +config = ( + PPOConfig() + .environment( + "env", + env_config={ + # Make analogous to old v4 + NoFrameskip. + "frameskip": 1, + "full_action_space": False, + "repeat_action_probability": 0.0, + }, + clip_rewards=True, + ) + .env_runners( + env_to_module_connector=_make_env_to_module_connector, + ) + .training( + learner_connector=_make_learner_connector, + train_batch_size_per_learner=4000, + minibatch_size=128, + lambda_=0.95, + kl_coeff=0.5, + clip_param=0.1, + vf_clip_param=10.0, + entropy_coeff=0.01, + num_epochs=10, + lr=0.00015 * NUM_LEARNERS, + grad_clip=100.0, + grad_clip_by="global_norm", + ) + .rl_module( + model_config=DefaultModelConfig( + conv_filters=[[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]], + conv_activation="relu", + head_fcnet_hiddens=[256], + vf_share_layers=True, + ), + ) + .callbacks( + functools.partial( + PongEnvTaskCallback, + task_threshold_map=task_threshold_map, + remote_fn=_remote_fn, + # Avoids downgrading the task to early when the agent had an unlucky run. + demotion_margin=args.demotion_margin, + # The return at which the task is learned. + solved_return=args.solved_return, + ) + ) +) + +if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + run_rllib_example_script_experiment(config, args=args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/debugging/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/debugging/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/debugging/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/debugging/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a7fba8f422c9ec356722703a19c8cfea9f0d508 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/debugging/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/debugging/__pycache__/deterministic_sampling_and_training.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/debugging/__pycache__/deterministic_sampling_and_training.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24e7ee41303c98e1149134c1044fdf2b92e354aa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/debugging/__pycache__/deterministic_sampling_and_training.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/debugging/deterministic_sampling_and_training.py b/lib/python3.12/site-packages/ray/rllib/examples/debugging/deterministic_sampling_and_training.py new file mode 100644 index 0000000000000000000000000000000000000000..b56cff72bb4f74502c04fc9f702773f7185f4016 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/debugging/deterministic_sampling_and_training.py @@ -0,0 +1,155 @@ +"""Example of how to seed your experiment with the `config.debugging(seed=...)` option. + +This example shows: + - how to seed an experiment, both on the Learner and on the EnvRunner side. + - that different experiments run with the exact same seed always yield the exact + same results (use the `--as-test` option to enforce assertions on the results). + Results checked range from EnvRunner stats, such as episode return, to Learner + stats, such as losses and gradient averages. + +Note that some algorithms, such as APPO which rely on asynchronous sampling in +combination with Ray network communication always behave stochastically, no matter +whether you set a seed or not. Therefore, make sure your `--algo` option is set to +a non-asynchronous algorithm, like "PPO" or "DQN". + + +How to run this script +---------------------- +`python [script file name].py --seed 1234` + +Use the `--num-learners=2` option to run with multiple Learner workers and, if GPUs +are available, place these workers on multiple GPUs. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0 --num-learners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should expect to see 2 experiments running and finishing in your console. +After the second experiment, you should see the confirmation that both experiments +yielded the exact same metrics. + ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_fb6d2_00000 | TERMINATED | 127.0.0.1:86298 | 3 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | episode_return_mean | num_env_steps_sample | +| | | d_lifetime | +|------------------+------------------------+------------------------| +| 6.2416 | 67.52 | 12004 | ++------------------+------------------------+------------------------+ + +... + +Determinism works! ok +""" +import ray +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + LEARNER_RESULTS, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + check, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args(default_iters=3) +parser.set_defaults( + # Test by default with more than one Env per EnvRunner. + num_envs_per_env_runner=2, +) +parser.add_argument("--seed", type=int, default=42) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "env", + lambda _: MultiAgentCartPole(config={"num_agents": args.num_agents}), + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env" if args.num_agents > 0 else "CartPole-v1") + # Make sure every environment gets a fixed seed. + .debugging(seed=args.seed) + # Log gradients and check them in the test. + .reporting(log_gradients=True) + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + results1 = run_rllib_example_script_experiment( + base_config, + args, + keep_ray_up=True, + success_metric={ENV_RUNNER_RESULTS + "/" + EPISODE_RETURN_MEAN: 10.0}, + ) + results2 = run_rllib_example_script_experiment( + base_config, + args, + keep_ray_up=True, + success_metric={ENV_RUNNER_RESULTS + "/" + EPISODE_RETURN_MEAN: 10.0}, + ) + + if args.as_test: + results1 = results1.get_best_result().metrics + results2 = results2.get_best_result().metrics + + # Test EnvRunner behaviors. + check( + results1[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN], + results2[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN], + ) + # As well as training behavior (minibatch sequence during SGD + # iterations). + for key in [ + # Losses and coefficients. + "curr_kl_coeff", + "vf_loss", + "policy_loss", + "entropy", + "total_loss", + "module_train_batch_size_mean", + # Optimizer stuff. + "gradients_default_optimizer_global_norm", + ]: + if args.num_agents > 0: + for aid in range(args.num_agents): + check( + results1[LEARNER_RESULTS][f"p{aid}"][key], + results2[LEARNER_RESULTS][f"p{aid}"][key], + ) + else: + check( + results1[LEARNER_RESULTS][DEFAULT_MODULE_ID][key], + results2[LEARNER_RESULTS][DEFAULT_MODULE_ID][key], + ) + + print("Determinism works! ok") + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7464d2bf41f22613f30afecd41600938f685f435 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/custom_evaluation.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/custom_evaluation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..577bc4329b4c426d1fdfea78a9f15b10c5826a72 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/custom_evaluation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/evaluation_parallel_to_training.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/evaluation_parallel_to_training.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..755293fd5400950911df817e0191367f91b3ef2b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/__pycache__/evaluation_parallel_to_training.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/evaluation/custom_evaluation.py b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/custom_evaluation.py new file mode 100644 index 0000000000000000000000000000000000000000..be0e4a0d71e32c58a1b4027898985f04ed734ee4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/custom_evaluation.py @@ -0,0 +1,230 @@ +"""Example of customizing the evaluation procedure for an RLlib Algorithm. + +Note, that you should only choose to provide a custom eval function, in case the already +built-in eval options are not sufficient. Normally, though, RLlib's eval utilities +that come with each Algorithm are enough to properly evaluate the learning progress +of your Algorithm. + +This script uses the SimpleCorridor environment, a simple 1D gridworld, in which +the agent can only walk left (action=0) or right (action=1). The goal state is located +at the end of the (1D) corridor. The env exposes an API to change the length of the +corridor on-the-fly. We use this API here to extend the size of the corridor for the +evaluation runs. + +For demonstration purposes only, we define a simple custom evaluation method that does +the following: +- It changes the corridor length of all environments used on the evaluation EnvRunners. +- It runs a defined number of episodes for evaluation purposes. +- It collects the metrics from those runs, summarizes these metrics and returns them. + + +How to run this script +---------------------- +`python [script file name].py + +You can switch off custom evaluation (and use RLlib's default evaluation procedure) +with the `--no-custom-eval` flag. + +You can switch on parallel evaluation to training using the +`--evaluation-parallel-to-training` flag. See this example script here: +https://github.com/ray-project/ray/blob/master/rllib/examples/evaluation/evaluation_parallel_to_training.py # noqa +for more details on running evaluation parallel to training. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see the following (or very similar) console output when running this script. +Note that for each iteration, due to the definition of our custom evaluation function, +we run 3 evaluation rounds per single training round. + +... +Training iteration 1 -> evaluation round 0 +Training iteration 1 -> evaluation round 1 +Training iteration 1 -> evaluation round 2 +... +... ++--------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|--------------------------------+------------+-----------------+--------+ +| PPO_SimpleCorridor_06582_00000 | TERMINATED | 127.0.0.1:69905 | 4 | ++--------------------------------+------------+-----------------+--------+ ++------------------+-------+----------+--------------------+ +| total time (s) | ts | reward | episode_len_mean | +|------------------+-------+----------+--------------------| +| 26.1973 | 16000 | 0.872034 | 13.7966 | ++------------------+-------+----------+--------------------+ +""" +from typing import Tuple + +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.env.env_runner_group import EnvRunnerGroup +from ray.rllib.examples.envs.classes.simple_corridor import SimpleCorridor +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.rllib.utils.typing import ResultDict +from ray.tune.registry import get_trainable_cls +from ray.tune.result import TRAINING_ITERATION + +parser = add_rllib_example_script_args( + default_iters=50, + default_reward=0.7, + default_timesteps=50000, +) +parser.add_argument("--no-custom-eval", action="store_true") +parser.add_argument("--corridor-length-training", type=int, default=10) +parser.add_argument("--corridor-length-eval-worker-1", type=int, default=20) +parser.add_argument("--corridor-length-eval-worker-2", type=int, default=30) + + +def custom_eval_function( + algorithm: Algorithm, + eval_workers: EnvRunnerGroup, +) -> Tuple[ResultDict, int, int]: + """Example of a custom evaluation function. + + Args: + algorithm: Algorithm class to evaluate. + eval_workers: Evaluation EnvRunnerGroup. + + Returns: + metrics: Evaluation metrics dict. + """ + # Set different env settings for each (eval) EnvRunner. Here we use the EnvRunner's + # `worker_index` property to figure out the actual length. + # Loop through all workers and all sub-envs (gym.Env) on each worker and call the + # `set_corridor_length` method on these. + eval_workers.foreach_env_runner( + func=lambda worker: ( + env.unwrapped.set_corridor_length( + args.corridor_length_eval_worker_1 + if worker.worker_index == 1 + else args.corridor_length_eval_worker_2 + ) + for env in worker.env.unwrapped.envs + ) + ) + + # Collect metrics results collected by eval workers in this list for later + # processing. + env_runner_metrics = [] + sampled_episodes = [] + # For demonstration purposes, run through some number of evaluation + # rounds within this one call. Note that this function is called once per + # training iteration (`Algorithm.train()` call) OR once per `Algorithm.evaluate()` + # (which can be called manually by the user). + for i in range(3): + print(f"Training iteration {algorithm.iteration} -> evaluation round {i}") + # Sample episodes from the EnvRunners AND have them return only the thus + # collected metrics. + episodes_and_metrics_all_env_runners = eval_workers.foreach_env_runner( + # Return only the metrics, NOT the sampled episodes (we don't need them + # anymore). + func=lambda worker: (worker.sample(), worker.get_metrics()), + local_env_runner=False, + ) + sampled_episodes.extend( + eps + for eps_and_mtrcs in episodes_and_metrics_all_env_runners + for eps in eps_and_mtrcs[0] + ) + env_runner_metrics.extend( + eps_and_mtrcs[1] for eps_and_mtrcs in episodes_and_metrics_all_env_runners + ) + + # You can compute metrics from the episodes manually, or use the Algorithm's + # convenient MetricsLogger to store all evaluation metrics inside the main + # algo. + algorithm.metrics.aggregate( + env_runner_metrics, key=(EVALUATION_RESULTS, ENV_RUNNER_RESULTS) + ) + eval_results = algorithm.metrics.peek((EVALUATION_RESULTS, ENV_RUNNER_RESULTS)) + # Alternatively, you could manually reduce over the n returned `env_runner_metrics` + # dicts, but this would be much harder as you might not know, which metrics + # to sum up, which ones to average over, etc.. + + # Compute env and agent steps from sampled episodes. + env_steps = sum(eps.env_steps() for eps in sampled_episodes) + agent_steps = sum(eps.agent_steps() for eps in sampled_episodes) + + return eval_results, env_steps, agent_steps + + +if __name__ == "__main__": + args = parser.parse_args() + args.local_mode = True + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + # For training, we use a corridor length of n. For evaluation, we use different + # values, depending on the eval worker index (1 or 2). + .environment( + SimpleCorridor, + env_config={"corridor_length": args.corridor_length_training}, + ) + .env_runners(create_env_on_local_worker=True) + .evaluation( + # Do we use the custom eval function defined above? + custom_evaluation_function=( + None if args.no_custom_eval else custom_eval_function + ), + # Number of eval EnvRunners to use. + evaluation_num_env_runners=2, + # Enable evaluation, once per training iteration. + evaluation_interval=1, + # Run 10 episodes each time evaluation runs (OR "auto" if parallel to + # training). + evaluation_duration="auto" if args.evaluation_parallel_to_training else 10, + # Evaluate parallelly to training? + evaluation_parallel_to_training=args.evaluation_parallel_to_training, + # Override the env settings for the eval workers. + # Note, though, that this setting here is only used in case --no-custom-eval + # is set, b/c in case the custom eval function IS used, we override the + # length of the eval environments in that custom function, so this setting + # here is simply ignored. + evaluation_config=AlgorithmConfig.overrides( + env_config={"corridor_length": args.corridor_length_training * 2}, + # TODO (sven): Add support for window=float(inf) and reduce=mean for + # evaluation episode_return_mean reductions (identical to old stack + # behavior, which does NOT use a window (100 by default) to reduce + # eval episode returns. + metrics_num_episodes_for_smoothing=5, + ), + ) + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": ( + args.stop_reward + ), + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + } + + run_rllib_example_script_experiment( + base_config, + args, + stop=stop, + success_metric={ + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": ( + args.stop_reward + ), + }, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/evaluation/evaluation_parallel_to_training.py b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/evaluation_parallel_to_training.py new file mode 100644 index 0000000000000000000000000000000000000000..7a9c35062b8222ea3d4636e50327a3faf54b8119 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/evaluation/evaluation_parallel_to_training.py @@ -0,0 +1,256 @@ +"""Example showing how one can set up evaluation running in parallel to training. + +Such a setup saves a considerable amount of time during RL Algorithm training, b/c +the next training step does NOT have to wait for the previous evaluation procedure to +finish, but can already start running (in parallel). + +See RLlib's documentation for more details on the effect of the different supported +evaluation configuration options: +https://docs.ray.io/en/latest/rllib/rllib-advanced-api.html#customized-evaluation-during-training # noqa + +For an example of how to write a fully customized evaluation function (which normally +is not necessary as the config options are sufficient and offer maximum flexibility), +see this example script here: + +https://github.com/ray-project/ray/blob/master/rllib/examples/evaluation/custom_evaluation.py # noqa + + +How to run this script +---------------------- +`python [script file name].py` + +Use the `--evaluation-num-workers` option to scale up the evaluation workers. Note +that the requested evaluation duration (`--evaluation-duration` measured in +`--evaluation-duration-unit`, which is either "timesteps" (default) or "episodes") is +shared between all configured evaluation workers. For example, if the evaluation +duration is 10 and the unit is "episodes" and you configured 5 workers, then each of the +evaluation workers will run exactly 2 episodes. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see the following output (at the end of the experiment) in your console when +running with a fixed number of 100k training timesteps +(`--evaluation-duration=auto --stop-timesteps=100000 +--stop-reward=100000`): ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_1377a_00000 | TERMINATED | 127.0.0.1:73330 | 25 | ++-----------------------------+------------+-----------------+--------+ ++------------------+--------+----------+--------------------+ +| total time (s) | ts | reward | episode_len_mean | +|------------------+--------+----------+--------------------| +| 71.7485 | 100000 | 476.51 | 476.51 | ++------------------+--------+----------+--------------------+ + +When running without parallel evaluation (no `--evaluation-parallel-to-training` flag), +the experiment takes considerably longer (~70sec vs ~80sec): ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_f1788_00000 | TERMINATED | 127.0.0.1:75135 | 25 | ++-----------------------------+------------+-----------------+--------+ ++------------------+--------+----------+--------------------+ +| total time (s) | ts | reward | episode_len_mean | +|------------------+--------+----------+--------------------| +| 81.7371 | 100000 | 494.68 | 494.68 | ++------------------+--------+----------+--------------------+ +""" +from typing import Optional + +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_EPISODES, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.rllib.utils.typing import ResultDict +from ray.tune.registry import get_trainable_cls, register_env +from ray.tune.result import TRAINING_ITERATION + +parser = add_rllib_example_script_args( + default_timesteps=200000, + default_reward=500.0, +) +parser.set_defaults( + evaluation_num_env_runners=2, + evaluation_interval=1, +) + + +class AssertEvalCallback(RLlibCallback): + def on_train_result( + self, + *, + algorithm: Algorithm, + metrics_logger: Optional[MetricsLogger] = None, + result: ResultDict, + **kwargs, + ): + # The eval results can be found inside the main `result` dict + # (old API stack: "evaluation"). + eval_results = result.get(EVALUATION_RESULTS, {}) + # In there, there is a sub-key: ENV_RUNNER_RESULTS. + eval_env_runner_results = eval_results.get(ENV_RUNNER_RESULTS) + # Make sure we always run exactly the given evaluation duration, + # no matter what the other settings are (such as + # `evaluation_num_env_runners` or `evaluation_parallel_to_training`). + if eval_env_runner_results and NUM_EPISODES in eval_env_runner_results: + num_episodes_done = eval_env_runner_results[NUM_EPISODES] + if algorithm.config.enable_env_runner_and_connector_v2: + num_timesteps_reported = eval_env_runner_results[NUM_ENV_STEPS_SAMPLED] + else: + num_timesteps_reported = eval_results["timesteps_this_iter"] + + # We run for automatic duration (as long as training takes). + if algorithm.config.evaluation_duration == "auto": + # If duration=auto: Expect at least as many timesteps as workers + # (each worker's `sample()` is at least called once). + # UNLESS: All eval workers were completely busy during the auto-time + # with older (async) requests and did NOT return anything from the async + # fetch. + assert ( + num_timesteps_reported == 0 + or num_timesteps_reported + >= algorithm.config.evaluation_num_env_runners + ) + # We count in episodes. + elif algorithm.config.evaluation_duration_unit == "episodes": + # Compare number of entries in episode_lengths (this is the + # number of episodes actually run) with desired number of + # episodes from the config. + assert ( + algorithm.iteration + 1 % algorithm.config.evaluation_interval != 0 + or num_episodes_done == algorithm.config.evaluation_duration + ), (num_episodes_done, algorithm.config.evaluation_duration) + print( + "Number of run evaluation episodes: " f"{num_episodes_done} (ok)!" + ) + # We count in timesteps. + else: + # TODO (sven): This assertion works perfectly fine locally, but breaks + # the CI for no reason. The observed collected timesteps is +500 more + # than desired (~2500 instead of 2011 and ~1250 vs 1011). + # num_timesteps_wanted = algorithm.config.evaluation_duration + # delta = num_timesteps_wanted - num_timesteps_reported + # Expect roughly the same (desired // num-eval-workers). + # assert abs(delta) < 20, ( + # delta, + # num_timesteps_wanted, + # num_timesteps_reported, + # ) + print( + "Number of run evaluation timesteps: " + f"{num_timesteps_reported} (ok?)!" + ) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "env", + lambda _: MultiAgentCartPole(config={"num_agents": args.num_agents}), + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env" if args.num_agents > 0 else "CartPole-v1") + .env_runners(create_env_on_local_worker=True) + # Use a custom callback that asserts that we are running the + # configured exact number of episodes per evaluation OR - in auto + # mode - run at least as many episodes as we have eval workers. + .callbacks(AssertEvalCallback) + .evaluation( + # Parallel evaluation+training config. + # Switch on evaluation in parallel with training. + evaluation_parallel_to_training=args.evaluation_parallel_to_training, + # Use two evaluation workers. Must be >0, otherwise, + # evaluation will run on a local worker and block (no parallelism). + evaluation_num_env_runners=args.evaluation_num_env_runners, + # Evaluate every other training iteration (together + # with every other call to Algorithm.train()). + evaluation_interval=args.evaluation_interval, + # Run for n episodes/timesteps (properly distribute load amongst + # all eval workers). The longer it takes to evaluate, the more sense + # it makes to use `evaluation_parallel_to_training=True`. + # Use "auto" to run evaluation for roughly as long as the training + # step takes. + evaluation_duration=args.evaluation_duration, + # "episodes" or "timesteps". + evaluation_duration_unit=args.evaluation_duration_unit, + # Switch off exploratory behavior for better (greedy) results. + evaluation_config={ + "explore": False, + # TODO (sven): Add support for window=float(inf) and reduce=mean for + # evaluation episode_return_mean reductions (identical to old stack + # behavior, which does NOT use a window (100 by default) to reduce + # eval episode returns. + "metrics_num_episodes_for_smoothing": 5, + }, + ) + ) + + # Set the minimum time for an iteration to 10sec, even for algorithms like PPO + # that naturally limit their iteration times to exactly one `training_step` + # call. This provides enough time for the eval EnvRunners in the + # "evaluation_duration=auto" setting to sample at least one complete episode. + if args.evaluation_duration == "auto": + base_config.reporting(min_time_s_per_iteration=10) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + # Set some PPO-specific tuning settings to learn better in the env (assumed to be + # CartPole-v1). + if args.algo == "PPO": + base_config.training( + lr=0.0003, + num_epochs=6, + vf_loss_coeff=0.01, + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": ( + args.stop_reward + ), + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + } + + run_rllib_example_script_experiment( + base_config, + args, + stop=stop, + success_metric={ + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": ( + args.stop_reward + ), + }, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d0fbce8b7d378b6a76f065b8d2b29eb0141ab36 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__pycache__/crashing_and_stalling_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__pycache__/crashing_and_stalling_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4aef6712698d1af69d2198c9425be03fac5c27a7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/__pycache__/crashing_and_stalling_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/crashing_and_stalling_env.py b/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/crashing_and_stalling_env.py new file mode 100644 index 0000000000000000000000000000000000000000..39b700edde700dd91ed52f4c1470348ff49c4cc5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/fault_tolerance/crashing_and_stalling_env.py @@ -0,0 +1,175 @@ +"""Example demonstrating that RLlib can learn (at scale) in unstable environments. + +This script uses the `CartPoleCrashing` environment, an adapted cartpole env whose +instability is configurable through setting the probability of a crash and/or stall +(sleep for a configurable amount of time) during `reset()` and/or `step()`. + +RLlib has two major flags for EnvRunner fault tolerance, which can be independently +set to True: +1) `config.fault_tolerance(restart_failed_sub_environments=True)` causes only the +(gymnasium) environment object on an EnvRunner to be closed (try calling `close()` on +the faulty object), garbage collected, and finally recreated from scratch. Note that +during this process, the containing EnvRunner remaing up and running and sampling +simply continues after the env recycling. This is the lightest and fastest form of +fault tolerance and should be attempted first. +2) `config.fault_tolerance(restart_failed_env_runners=True)` causes the entire +EnvRunner (a Ray remote actor) to be restarted. This restart logically includes the +gymnasium environment, the RLModule, and all connector pipelines on the EnvRunner. +Use this option only if you face problems with the first option +(restart_failed_sub_environments=True), such as incomplete cleanups and memory leaks. + + +How to run this script +---------------------- +`python [script file name].py + +You can switch on the fault tolerant behavior (1) (restart_failed_sub_environments) +through the `--restart-failed-envs` flag. If this flag is not set, the script will +recreate the entire (faulty) EnvRunner. + +You can switch on stalling (besides crashing) through the `--stall` command line flag. +If set, besides crashing on `reset()` and/or `step()`, there is also a chance of +stalling for a few seconds on each of these events. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see the following (or very similar) console output when running this script +with: +`--algo=PPO --stall --restart-failed-envs --stop-reward=450.0` ++---------------------+------------+----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|---------------------+------------+----------------+--------+------------------+ +| PPO_env_ba39b_00000 | TERMINATED | 127.0.0.1:1401 | 22 | 133.497 | ++---------------------+------------+----------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| episode_return_mean | num_episodes_lifetim | num_env_steps_traine | +| | e | d_lifetime | +|------------------------+------------------------+------------------------| +| 450.24 | 542 | 88628 | ++------------------------+------------------------+------------------------+ + +For APPO and testing restarting the entire EnvRunners, you could run the script with: +`--algo=APPO --stall --stop-reward=450.0` ++----------------------+------------+----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|----------------------+------------+----------------+--------+------------------+ +| APPO_env_ba39b_00000 | TERMINATED | 127.0.0.1:4653 | 10 | 101.531 | ++----------------------+------------+----------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| episode_return_mean | num_episodes_lifetim | num_env_steps_traine | +| | e | d_lifetime | +|------------------------+------------------------+------------------------| +| 478.85 | 2546 | 321500 | ++------------------------+------------------------+------------------------+ +""" +from gymnasium.wrappers import TimeLimit + +from ray import tune +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.envs.classes.cartpole_crashing import ( + CartPoleCrashing, + MultiAgentCartPoleCrashing, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_reward=450.0, + default_timesteps=2000000, +) +parser.set_defaults( + num_env_runners=4, +) +# Use `parser` to add your own custom command line options to this script +# and (if needed) use their values to set up `config` below. +parser.add_argument( + "--stall", + action="store_true", + help="Whether to also stall the env from time to time", +) +parser.add_argument( + "--restart-failed-envs", + action="store_true", + help="Whether to restart a failed environment (vs restarting the entire " + "EnvRunner).", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + if args.num_agents > 0: + tune.register_env("env", lambda cfg: MultiAgentCartPoleCrashing(cfg)) + else: + tune.register_env( + "env", + lambda cfg: TimeLimit(CartPoleCrashing(cfg), max_episode_steps=500), + ) + + base_config = ( + tune.registry.get_trainable_cls(args.algo) + .get_default_config() + .environment( + "env", + env_config={ + "num_agents": args.num_agents, + # Probability to crash during step(). + "p_crash": 0.0001, + # Probability to crash during reset(). + "p_crash_reset": 0.001, + "crash_on_worker_indices": [1, 2], + "init_time_s": 2.0, + # Probability to stall during step(). + "p_stall": 0.0005, + # Probability to stall during reset(). + "p_stall_reset": 0.001, + # Stall from 2 to 5sec (or 0.0 if --stall not set). + "stall_time_sec": (2, 5) if args.stall else 0.0, + # EnvRunner indices to stall on. + "stall_on_worker_indices": [2, 3], + }, + ) + # Switch on resiliency. + .fault_tolerance( + # Recreate any failed EnvRunners. + restart_failed_env_runners=True, + # Restart any failed environment (w/o recreating the EnvRunner). Note that + # this is the much faster option. + restart_failed_sub_environments=args.restart_failed_envs, + ) + ) + + # Use more stabilizing hyperparams for APPO. + if args.algo == "APPO": + base_config.training( + grad_clip=40.0, + entropy_coeff=0.0, + vf_loss_coeff=0.05, + ) + base_config.rl_module( + model_config=DefaultModelConfig(vf_share_layers=True), + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + run_rllib_example_script_experiment(base_config, args=args) diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__init__.py b/lib/python3.12/site-packages/ray/rllib/execution/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4538bf7e92177cf73eb10fca92de90ac35f56a88 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/__init__.py @@ -0,0 +1,23 @@ +from ray.rllib.execution.learner_thread import LearnerThread +from ray.rllib.execution.minibatch_buffer import MinibatchBuffer +from ray.rllib.execution.multi_gpu_learner_thread import MultiGPULearnerThread +from ray.rllib.execution.replay_ops import SimpleReplayBuffer +from ray.rllib.execution.rollout_ops import ( + standardize_fields, + synchronous_parallel_sample, +) +from ray.rllib.execution.train_ops import ( + multi_gpu_train_one_step, + train_one_step, +) + +__all__ = [ + "multi_gpu_train_one_step", + "standardize_fields", + "synchronous_parallel_sample", + "train_one_step", + "LearnerThread", + "MultiGPULearnerThread", + "SimpleReplayBuffer", + "MinibatchBuffer", +] diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eee01f56a41ca8c60a069a78a67a70cf41454702 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/learner_thread.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/learner_thread.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bad46a7bc7024d2afed44bf0e53f2d37f72747b2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/learner_thread.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/minibatch_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/minibatch_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80af87a3731673fa8af832467222bfd16f1c9a68 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/minibatch_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/multi_gpu_learner_thread.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/multi_gpu_learner_thread.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a5879386fecc2025ecddaf9d507a9690e5ec150 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/multi_gpu_learner_thread.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/replay_ops.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/replay_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b4d4b9e0a6577a1271ebd2a4530a045228c9580 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/replay_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/rollout_ops.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/rollout_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06250e92685060d06332911704bb85f20880ac41 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/rollout_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/segment_tree.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/segment_tree.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f178e32521140150571e27dc971ce13a32cf064d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/segment_tree.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/train_ops.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/train_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8d68a39b928502565d3c614f5c914b94f12edcb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/__pycache__/train_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/buffers/__init__.py b/lib/python3.12/site-packages/ray/rllib/execution/buffers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/execution/buffers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/buffers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0dcb3ddb4876f136b6745c3c0616c6dfac2a38d5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/buffers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/buffers/__pycache__/mixin_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/execution/buffers/__pycache__/mixin_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a4952e65d1f00fe80e9f7d6a6f7882c782be6eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/execution/buffers/__pycache__/mixin_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/execution/buffers/mixin_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/execution/buffers/mixin_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..57195ddbd181d64dd1a66fe740dfc70d692d42d0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/buffers/mixin_replay_buffer.py @@ -0,0 +1,173 @@ +import collections +import platform +import random +from typing import Optional + +from ray.rllib.execution.replay_ops import SimpleReplayBuffer +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, concat_samples +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.replay_buffers.multi_agent_replay_buffer import ReplayMode +from ray.rllib.utils.replay_buffers.replay_buffer import _ALL_POLICIES +from ray.rllib.utils.typing import PolicyID, SampleBatchType +from ray.util.timer import _Timer + + +@OldAPIStack +class MixInMultiAgentReplayBuffer: + """This buffer adds replayed samples to a stream of new experiences. + + - Any newly added batch (`add()`) is immediately returned upon + the next `replay` call (close to on-policy) as well as being moved + into the buffer. + - Additionally, a certain number of old samples is mixed into the + returned sample according to a given "replay ratio". + - If >1 calls to `add()` are made without any `replay()` calls + in between, all newly added batches are returned (plus some older samples + according to the "replay ratio"). + + .. testcode:: + + from ray.rllib.execution.buffers.mixin_replay_buffer import ( + MixInMultiAgentReplayBuffer) + from ray.rllib.policy.sample_batch import SampleBatch + # replay ratio 0.66 (2/3 replayed, 1/3 new samples): + buffer = MixInMultiAgentReplayBuffer(capacity=100, + replay_ratio=0.66) + A, B, C = (SampleBatch({"obs": [1]}), SampleBatch({"obs": [2]}), + SampleBatch({"obs": [3]})) + buffer.add(A) + buffer.add(B) + buffer.add(B) + print(buffer.replay()["obs"]) + + .. testoutput:: + :hide: + + ... + """ + + def __init__( + self, + capacity: int, + replay_ratio: float, + replay_mode: ReplayMode = ReplayMode.INDEPENDENT, + ): + """Initializes MixInReplay instance. + + Args: + capacity: Number of batches to store in total. + replay_ratio: Ratio of replayed samples in the returned + batches. E.g. a ratio of 0.0 means only return new samples + (no replay), a ratio of 0.5 means always return newest sample + plus one old one (1:1), a ratio of 0.66 means always return + the newest sample plus 2 old (replayed) ones (1:2), etc... + """ + self.capacity = capacity + self.replay_ratio = replay_ratio + self.replay_proportion = None + if self.replay_ratio != 1.0: + self.replay_proportion = self.replay_ratio / (1.0 - self.replay_ratio) + + if replay_mode in ["lockstep", ReplayMode.LOCKSTEP]: + self.replay_mode = ReplayMode.LOCKSTEP + elif replay_mode in ["independent", ReplayMode.INDEPENDENT]: + self.replay_mode = ReplayMode.INDEPENDENT + else: + raise ValueError("Unsupported replay mode: {}".format(replay_mode)) + + def new_buffer(): + return SimpleReplayBuffer(num_slots=capacity) + + self.replay_buffers = collections.defaultdict(new_buffer) + + # Metrics. + self.add_batch_timer = _Timer() + self.replay_timer = _Timer() + self.update_priorities_timer = _Timer() + + # Added timesteps over lifetime. + self.num_added = 0 + + # Last added batch(es). + self.last_added_batches = collections.defaultdict(list) + + def add(self, batch: SampleBatchType) -> None: + """Adds a batch to the appropriate policy's replay buffer. + + Turns the batch into a MultiAgentBatch of the DEFAULT_POLICY_ID if + it is not a MultiAgentBatch. Subsequently adds the individual policy + batches to the storage. + + Args: + batch: The batch to be added. + """ + # Make a copy so the replay buffer doesn't pin plasma memory. + batch = batch.copy() + batch = batch.as_multi_agent() + + with self.add_batch_timer: + if self.replay_mode == ReplayMode.LOCKSTEP: + # Lockstep mode: Store under _ALL_POLICIES key (we will always + # only sample from all policies at the same time). + # This means storing a MultiAgentBatch to the underlying buffer + self.replay_buffers[_ALL_POLICIES].add_batch(batch) + self.last_added_batches[_ALL_POLICIES].append(batch) + else: + # Store independent SampleBatches + for policy_id, sample_batch in batch.policy_batches.items(): + self.replay_buffers[policy_id].add_batch(sample_batch) + self.last_added_batches[policy_id].append(sample_batch) + + self.num_added += batch.count + + def replay( + self, policy_id: PolicyID = DEFAULT_POLICY_ID + ) -> Optional[SampleBatchType]: + if self.replay_mode == ReplayMode.LOCKSTEP and policy_id != _ALL_POLICIES: + raise ValueError( + "Trying to sample from single policy's buffer in lockstep " + "mode. In lockstep mode, all policies' experiences are " + "sampled from a single replay buffer which is accessed " + "with the policy id `{}`".format(_ALL_POLICIES) + ) + + buffer = self.replay_buffers[policy_id] + # Return None, if: + # - Buffer empty or + # - `replay_ratio` < 1.0 (new samples required in returned batch) + # and no new samples to mix with replayed ones. + if len(buffer) == 0 or ( + len(self.last_added_batches[policy_id]) == 0 and self.replay_ratio < 1.0 + ): + return None + + # Mix buffer's last added batches with older replayed batches. + with self.replay_timer: + output_batches = self.last_added_batches[policy_id] + self.last_added_batches[policy_id] = [] + + # No replay desired -> Return here. + if self.replay_ratio == 0.0: + return concat_samples(output_batches) + # Only replay desired -> Return a (replayed) sample from the + # buffer. + elif self.replay_ratio == 1.0: + return buffer.replay() + + # Replay ratio = old / [old + new] + # Replay proportion: old / new + num_new = len(output_batches) + replay_proportion = self.replay_proportion + while random.random() < num_new * replay_proportion: + replay_proportion -= 1 + output_batches.append(buffer.replay()) + return concat_samples(output_batches) + + def get_host(self) -> str: + """Returns the computer's network name. + + Returns: + The computer's networks name or an empty string, if the network + name could not be determined. + """ + return platform.node() diff --git a/lib/python3.12/site-packages/ray/rllib/execution/learner_thread.py b/lib/python3.12/site-packages/ray/rllib/execution/learner_thread.py new file mode 100644 index 0000000000000000000000000000000000000000..2c98829ba34a61d167a474bf1427e6b6ef94705a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/learner_thread.py @@ -0,0 +1,137 @@ +import copy +import queue +import threading +from typing import Dict, Optional + +from ray.rllib.evaluation.rollout_worker import RolloutWorker +from ray.rllib.execution.minibatch_buffer import MinibatchBuffer +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics.learner_info import LEARNER_INFO, LearnerInfoBuilder +from ray.rllib.utils.metrics.window_stat import WindowStat +from ray.util.iter import _NextValueNotReady +from ray.util.timer import _Timer + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class LearnerThread(threading.Thread): + """Background thread that updates the local model from sample trajectories. + + The learner thread communicates with the main thread through Queues. This + is needed since Ray operations can only be run on the main thread. In + addition, moving heavyweight gradient ops session runs off the main thread + improves overall throughput. + """ + + def __init__( + self, + local_worker: RolloutWorker, + minibatch_buffer_size: int, + num_sgd_iter: int, + learner_queue_size: int, + learner_queue_timeout: int, + ): + """Initialize the learner thread. + + Args: + local_worker: process local rollout worker holding + policies this thread will call learn_on_batch() on + minibatch_buffer_size: max number of train batches to store + in the minibatching buffer + num_sgd_iter: number of passes to learn on per train batch + learner_queue_size: max size of queue of inbound + train batches to this thread + learner_queue_timeout: raise an exception if the queue has + been empty for this long in seconds + """ + threading.Thread.__init__(self) + self.learner_queue_size = WindowStat("size", 50) + self.local_worker = local_worker + self.inqueue = queue.Queue(maxsize=learner_queue_size) + self.outqueue = queue.Queue() + self.minibatch_buffer = MinibatchBuffer( + inqueue=self.inqueue, + size=minibatch_buffer_size, + timeout=learner_queue_timeout, + num_passes=num_sgd_iter, + init_num_passes=num_sgd_iter, + ) + self.queue_timer = _Timer() + self.grad_timer = _Timer() + self.load_timer = _Timer() + self.load_wait_timer = _Timer() + self.daemon = True + self.policy_ids_updated = [] + self.learner_info = {} + self.stopped = False + self.num_steps = 0 + + def run(self) -> None: + # Switch on eager mode if configured. + if self.local_worker.config.framework_str == "tf2": + tf1.enable_eager_execution() + while not self.stopped: + self.step() + + def step(self) -> Optional[_NextValueNotReady]: + with self.queue_timer: + try: + batch, _ = self.minibatch_buffer.get() + except queue.Empty: + return _NextValueNotReady() + with self.grad_timer: + # Use LearnerInfoBuilder as a unified way to build the final + # results dict from `learn_on_loaded_batch` call(s). + # This makes sure results dicts always have the same structure + # no matter the setup (multi-GPU, multi-agent, minibatch SGD, + # tf vs torch). + learner_info_builder = LearnerInfoBuilder(num_devices=1) + if self.local_worker.config.policy_states_are_swappable: + self.local_worker.lock() + multi_agent_results = self.local_worker.learn_on_batch(batch) + if self.local_worker.config.policy_states_are_swappable: + self.local_worker.unlock() + self.policy_ids_updated.extend(list(multi_agent_results.keys())) + for pid, results in multi_agent_results.items(): + learner_info_builder.add_learn_on_batch_results(results, pid) + self.learner_info = learner_info_builder.finalize() + + self.num_steps += 1 + # Put tuple: env-steps, agent-steps, and learner info into the queue. + self.outqueue.put((batch.count, batch.agent_steps(), self.learner_info)) + self.learner_queue_size.push(self.inqueue.qsize()) + + def add_learner_metrics(self, result: Dict, overwrite_learner_info=True) -> Dict: + """Add internal metrics to a result dict.""" + + def timer_to_ms(timer): + return round(1000 * timer.mean, 3) + + if overwrite_learner_info: + result["info"].update( + { + "learner_queue": self.learner_queue_size.stats(), + LEARNER_INFO: copy.deepcopy(self.learner_info), + "timing_breakdown": { + "learner_grad_time_ms": timer_to_ms(self.grad_timer), + "learner_load_time_ms": timer_to_ms(self.load_timer), + "learner_load_wait_time_ms": timer_to_ms(self.load_wait_timer), + "learner_dequeue_time_ms": timer_to_ms(self.queue_timer), + }, + } + ) + else: + result["info"].update( + { + "learner_queue": self.learner_queue_size.stats(), + "timing_breakdown": { + "learner_grad_time_ms": timer_to_ms(self.grad_timer), + "learner_load_time_ms": timer_to_ms(self.load_timer), + "learner_load_wait_time_ms": timer_to_ms(self.load_wait_timer), + "learner_dequeue_time_ms": timer_to_ms(self.queue_timer), + }, + } + ) + return result diff --git a/lib/python3.12/site-packages/ray/rllib/execution/minibatch_buffer.py b/lib/python3.12/site-packages/ray/rllib/execution/minibatch_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..873cf3c5f9f403c1a471e27ce202c1cb0dd38349 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/minibatch_buffer.py @@ -0,0 +1,61 @@ +import queue +from typing import Any, Tuple + +from ray.rllib.utils.annotations import OldAPIStack + + +@OldAPIStack +class MinibatchBuffer: + """Ring buffer of recent data batches for minibatch SGD. + + This is for use with AsyncSamplesOptimizer. + """ + + def __init__( + self, + inqueue: queue.Queue, + size: int, + timeout: float, + num_passes: int, + init_num_passes: int = 1, + ): + """Initialize a minibatch buffer. + + Args: + inqueue (queue.Queue): Queue to populate the internal ring buffer + from. + size: Max number of data items to buffer. + timeout: Queue timeout + num_passes: Max num times each data item should be emitted. + init_num_passes: Initial passes for each data item. + Maxiumum number of passes per item are increased to num_passes over + time. + """ + self.inqueue = inqueue + self.size = size + self.timeout = timeout + self.max_initial_ttl = num_passes + self.cur_initial_ttl = init_num_passes + self.buffers = [None] * size + self.ttl = [0] * size + self.idx = 0 + + def get(self) -> Tuple[Any, bool]: + """Get a new batch from the internal ring buffer. + + Returns: + buf: Data item saved from inqueue. + released: True if the item is now removed from the ring buffer. + """ + if self.ttl[self.idx] <= 0: + self.buffers[self.idx] = self.inqueue.get(timeout=self.timeout) + self.ttl[self.idx] = self.cur_initial_ttl + if self.cur_initial_ttl < self.max_initial_ttl: + self.cur_initial_ttl += 1 + buf = self.buffers[self.idx] + self.ttl[self.idx] -= 1 + released = self.ttl[self.idx] <= 0 + if released: + self.buffers[self.idx] = None + self.idx = (self.idx + 1) % len(self.buffers) + return buf, released diff --git a/lib/python3.12/site-packages/ray/rllib/execution/multi_gpu_learner_thread.py b/lib/python3.12/site-packages/ray/rllib/execution/multi_gpu_learner_thread.py new file mode 100644 index 0000000000000000000000000000000000000000..65bd97f084bf488c8d021ace601ec73e578e345a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/multi_gpu_learner_thread.py @@ -0,0 +1,245 @@ +import logging +import queue +import threading + +from ray._common.deprecation import deprecation_warning +from ray.rllib.evaluation.rollout_worker import RolloutWorker +from ray.rllib.execution.learner_thread import LearnerThread +from ray.rllib.execution.minibatch_buffer import MinibatchBuffer +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics.learner_info import LearnerInfoBuilder +from ray.util.timer import _Timer + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class MultiGPULearnerThread(LearnerThread): + """Learner that can use multiple GPUs and parallel loading. + + This class is used for async sampling algorithms. + + Example workflow: 2 GPUs and 3 multi-GPU tower stacks. + -> On each GPU, there are 3 slots for batches, indexed 0, 1, and 2. + + Workers collect data from env and push it into inqueue: + Workers -> (data) -> self.inqueue + + We also have two queues, indicating, which stacks are loaded and which + are not. + - idle_tower_stacks = [0, 1, 2] <- all 3 stacks are free at first. + - ready_tower_stacks = [] <- None of the 3 stacks is loaded with data. + + `ready_tower_stacks` is managed by `ready_tower_stacks_buffer` for + possible minibatch-SGD iterations per loaded batch (this avoids a reload + from CPU to GPU for each SGD iter). + + n _MultiGPULoaderThreads: self.inqueue -get()-> + policy.load_batch_into_buffer() -> ready_stacks = [0 ...] + + This thread: self.ready_tower_stacks_buffer -get()-> + policy.learn_on_loaded_batch() -> if SGD-iters done, + put stack index back in idle_tower_stacks queue. + """ + + def __init__( + self, + local_worker: RolloutWorker, + num_gpus: int = 1, + lr=None, # deprecated. + train_batch_size: int = 500, + num_multi_gpu_tower_stacks: int = 1, + num_sgd_iter: int = 1, + learner_queue_size: int = 16, + learner_queue_timeout: int = 300, + num_data_load_threads: int = 16, + _fake_gpus: bool = False, + # Deprecated arg, use + minibatch_buffer_size=None, + ): + """Initializes a MultiGPULearnerThread instance. + + Args: + local_worker: Local RolloutWorker holding + policies this thread will call `load_batch_into_buffer` and + `learn_on_loaded_batch` on. + num_gpus: Number of GPUs to use for data-parallel SGD. + train_batch_size: Size of batches (minibatches if + `num_sgd_iter` > 1) to learn on. + num_multi_gpu_tower_stacks: Number of buffers to parallelly + load data into on one device. Each buffer is of size of + `train_batch_size` and hence increases GPU memory usage + accordingly. + num_sgd_iter: Number of passes to learn on per train batch + (minibatch if `num_sgd_iter` > 1). + learner_queue_size: Max size of queue of inbound + train batches to this thread. + num_data_load_threads: Number of threads to use to load + data into GPU memory in parallel. + """ + # Deprecated: No need to specify as we don't need the actual + # minibatch-buffer anyways. + if minibatch_buffer_size: + deprecation_warning( + old="MultiGPULearnerThread.minibatch_buffer_size", + error=True, + ) + super().__init__( + local_worker=local_worker, + minibatch_buffer_size=0, + num_sgd_iter=num_sgd_iter, + learner_queue_size=learner_queue_size, + learner_queue_timeout=learner_queue_timeout, + ) + # Delete reference to parent's minibatch_buffer, which is not needed. + # Instead, in multi-GPU mode, we pull tower stack indices from the + # `self.ready_tower_stacks_buffer` buffer, whose size is exactly + # `num_multi_gpu_tower_stacks`. + self.minibatch_buffer = None + + self.train_batch_size = train_batch_size + + self.policy_map = self.local_worker.policy_map + self.devices = next(iter(self.policy_map.values())).devices + + logger.info("MultiGPULearnerThread devices {}".format(self.devices)) + assert self.train_batch_size % len(self.devices) == 0 + assert self.train_batch_size >= len(self.devices), "batch too small" + + self.tower_stack_indices = list(range(num_multi_gpu_tower_stacks)) + + # Two queues for tower stacks: + # a) Those that are loaded with data ("ready") + # b) Those that are ready to be loaded with new data ("idle"). + self.idle_tower_stacks = queue.Queue() + self.ready_tower_stacks = queue.Queue() + # In the beginning, all stacks are idle (no loading has taken place + # yet). + for idx in self.tower_stack_indices: + self.idle_tower_stacks.put(idx) + # Start n threads that are responsible for loading data into the + # different (idle) stacks. + for i in range(num_data_load_threads): + self.loader_thread = _MultiGPULoaderThread(self, share_stats=(i == 0)) + self.loader_thread.start() + + # Create a buffer that holds stack indices that are "ready" + # (loaded with data). Those are stacks that we can call + # "learn_on_loaded_batch" on. + self.ready_tower_stacks_buffer = MinibatchBuffer( + self.ready_tower_stacks, + num_multi_gpu_tower_stacks, + learner_queue_timeout, + num_sgd_iter, + ) + + @override(LearnerThread) + def step(self) -> None: + if not self.loader_thread.is_alive(): + raise RuntimeError( + "The `_MultiGPULoaderThread` has died! Will therefore also terminate " + "the `MultiGPULearnerThread`." + ) + + with self.load_wait_timer: + buffer_idx, released = self.ready_tower_stacks_buffer.get() + + get_num_samples_loaded_into_buffer = 0 + with self.grad_timer: + # Use LearnerInfoBuilder as a unified way to build the final + # results dict from `learn_on_loaded_batch` call(s). + # This makes sure results dicts always have the same structure + # no matter the setup (multi-GPU, multi-agent, minibatch SGD, + # tf vs torch). + learner_info_builder = LearnerInfoBuilder(num_devices=len(self.devices)) + + for pid in self.policy_map.keys(): + # Not a policy-to-train. + if ( + self.local_worker.is_policy_to_train is not None + and not self.local_worker.is_policy_to_train(pid) + ): + continue + policy = self.policy_map[pid] + default_policy_results = policy.learn_on_loaded_batch( + offset=0, buffer_index=buffer_idx + ) + learner_info_builder.add_learn_on_batch_results( + default_policy_results, policy_id=pid + ) + self.policy_ids_updated.append(pid) + get_num_samples_loaded_into_buffer += ( + policy.get_num_samples_loaded_into_buffer(buffer_idx) + ) + + self.learner_info = learner_info_builder.finalize() + + if released: + self.idle_tower_stacks.put(buffer_idx) + + # Put tuple: env-steps, agent-steps, and learner info into the queue. + self.outqueue.put( + ( + get_num_samples_loaded_into_buffer, + get_num_samples_loaded_into_buffer, + self.learner_info, + ) + ) + self.learner_queue_size.push(self.inqueue.qsize()) + + +class _MultiGPULoaderThread(threading.Thread): + def __init__( + self, multi_gpu_learner_thread: MultiGPULearnerThread, share_stats: bool + ): + threading.Thread.__init__(self) + self.multi_gpu_learner_thread = multi_gpu_learner_thread + self.daemon = True + if share_stats: + self.queue_timer = multi_gpu_learner_thread.queue_timer + self.load_timer = multi_gpu_learner_thread.load_timer + else: + self.queue_timer = _Timer() + self.load_timer = _Timer() + + def run(self) -> None: + while True: + self._step() + + def _step(self) -> None: + s = self.multi_gpu_learner_thread + policy_map = s.policy_map + + # Get a new batch from the data (inqueue). + with self.queue_timer: + batch = s.inqueue.get() + + # Get next idle stack for loading. + buffer_idx = s.idle_tower_stacks.get() + + # Load the batch into the idle stack. + with self.load_timer: + for pid in policy_map.keys(): + if ( + s.local_worker.is_policy_to_train is not None + and not s.local_worker.is_policy_to_train(pid, batch) + ): + continue + policy = policy_map[pid] + if isinstance(batch, SampleBatch): + policy.load_batch_into_buffer( + batch=batch, + buffer_index=buffer_idx, + ) + elif pid in batch.policy_batches: + policy.load_batch_into_buffer( + batch=batch.policy_batches[pid], + buffer_index=buffer_idx, + ) + + # Tag just-loaded stack as "ready". + s.ready_tower_stacks.put(buffer_idx) diff --git a/lib/python3.12/site-packages/ray/rllib/execution/replay_ops.py b/lib/python3.12/site-packages/ray/rllib/execution/replay_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..13282d68b4a5d309777e72459628c970f5e27b78 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/replay_ops.py @@ -0,0 +1,37 @@ +import random +from typing import Optional + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.replay_buffers.replay_buffer import warn_replay_capacity +from ray.rllib.utils.typing import SampleBatchType + + +@OldAPIStack +class SimpleReplayBuffer: + """Simple replay buffer that operates over batches.""" + + def __init__(self, num_slots: int, replay_proportion: Optional[float] = None): + """Initialize SimpleReplayBuffer. + + Args: + num_slots: Number of batches to store in total. + """ + self.num_slots = num_slots + self.replay_batches = [] + self.replay_index = 0 + + def add_batch(self, sample_batch: SampleBatchType) -> None: + warn_replay_capacity(item=sample_batch, num_items=self.num_slots) + if self.num_slots > 0: + if len(self.replay_batches) < self.num_slots: + self.replay_batches.append(sample_batch) + else: + self.replay_batches[self.replay_index] = sample_batch + self.replay_index += 1 + self.replay_index %= self.num_slots + + def replay(self) -> SampleBatchType: + return random.choice(self.replay_batches) + + def __len__(self): + return len(self.replay_batches) diff --git a/lib/python3.12/site-packages/ray/rllib/execution/rollout_ops.py b/lib/python3.12/site-packages/ray/rllib/execution/rollout_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..9509a814880ef36a755d201c31fd9aceb68c9181 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/rollout_ops.py @@ -0,0 +1,208 @@ +import logging +from typing import List, Optional, Union + +import tree + +from ray.rllib.env.env_runner_group import EnvRunnerGroup +from ray.rllib.policy.sample_batch import ( + DEFAULT_POLICY_ID, + SampleBatch, + concat_samples, +) +from ray.rllib.utils.annotations import ExperimentalAPI, OldAPIStack +from ray.rllib.utils.metrics import NUM_AGENT_STEPS_SAMPLED, NUM_ENV_STEPS_SAMPLED +from ray.rllib.utils.sgd import standardized +from ray.rllib.utils.typing import EpisodeType, SampleBatchType + +logger = logging.getLogger(__name__) + + +@ExperimentalAPI +def synchronous_parallel_sample( + *, + worker_set: EnvRunnerGroup, + max_agent_steps: Optional[int] = None, + max_env_steps: Optional[int] = None, + concat: bool = True, + sample_timeout_s: Optional[float] = None, + random_actions: bool = False, + _uses_new_env_runners: bool = False, + _return_metrics: bool = False, +) -> Union[List[SampleBatchType], SampleBatchType, List[EpisodeType], EpisodeType]: + """Runs parallel and synchronous rollouts on all remote workers. + + Waits for all workers to return from the remote calls. + + If no remote workers exist (num_workers == 0), use the local worker + for sampling. + + Alternatively to calling `worker.sample.remote()`, the user can provide a + `remote_fn()`, which will be applied to the worker(s) instead. + + Args: + worker_set: The EnvRunnerGroup to use for sampling. + remote_fn: If provided, use `worker.apply.remote(remote_fn)` instead + of `worker.sample.remote()` to generate the requests. + max_agent_steps: Optional number of agent steps to be included in the + final batch or list of episodes. + max_env_steps: Optional number of environment steps to be included in the + final batch or list of episodes. + concat: Whether to aggregate all resulting batches or episodes. in case of + batches the list of batches is concatinated at the end. in case of + episodes all episode lists from workers are flattened into a single list. + sample_timeout_s: The timeout in sec to use on the `foreach_env_runner` call. + After this time, the call will return with a result (or not if all + EnvRunners are stalling). If None, will block indefinitely and not timeout. + _uses_new_env_runners: Whether the new `EnvRunner API` is used. In this case + episodes instead of `SampleBatch` objects are returned. + + Returns: + The list of collected sample batch types or episode types (one for each parallel + rollout worker in the given `worker_set`). + + .. testcode:: + + # Define an RLlib Algorithm. + from ray.rllib.algorithms.ppo import PPO, PPOConfig + config = ( + PPOConfig() + .environment("CartPole-v1") + ) + algorithm = config.build() + # 2 remote EnvRunners (num_env_runners=2): + episodes = synchronous_parallel_sample( + worker_set=algorithm.env_runner_group, + _uses_new_env_runners=True, + concat=False, + ) + print(len(episodes)) + + .. testoutput:: + + 2 + """ + # Only allow one of `max_agent_steps` or `max_env_steps` to be defined. + assert not (max_agent_steps is not None and max_env_steps is not None) + + agent_or_env_steps = 0 + max_agent_or_env_steps = max_agent_steps or max_env_steps or None + sample_batches_or_episodes = [] + all_stats_dicts = [] + + random_action_kwargs = {} if not random_actions else {"random_actions": True} + + # Stop collecting batches as soon as one criterium is met. + while (max_agent_or_env_steps is None and agent_or_env_steps == 0) or ( + max_agent_or_env_steps is not None + and agent_or_env_steps < max_agent_or_env_steps + ): + # No remote workers in the set -> Use local worker for collecting + # samples. + if worker_set.num_remote_workers() <= 0: + sampled_data = [worker_set.local_env_runner.sample(**random_action_kwargs)] + if _return_metrics: + stats_dicts = [worker_set.local_env_runner.get_metrics()] + # Loop over remote workers' `sample()` method in parallel. + else: + sampled_data = worker_set.foreach_env_runner( + ( + (lambda w: w.sample(**random_action_kwargs)) + if not _return_metrics + else (lambda w: (w.sample(**random_action_kwargs), w.get_metrics())) + ), + local_env_runner=False, + timeout_seconds=sample_timeout_s, + ) + # Nothing was returned (maybe all workers are stalling) or no healthy + # remote workers left: Break. + # There is no point staying in this loop, since we will not be able to + # get any new samples if we don't have any healthy remote workers left. + if not sampled_data or worker_set.num_healthy_remote_workers() <= 0: + if not sampled_data: + logger.warning( + "No samples returned from remote workers. If you have a " + "slow environment or model, consider increasing the " + "`sample_timeout_s` or decreasing the " + "`rollout_fragment_length` in `AlgorithmConfig.env_runners()." + ) + elif worker_set.num_healthy_remote_workers() <= 0: + logger.warning( + "No healthy remote workers left. Trying to restore workers ..." + ) + break + + if _return_metrics: + stats_dicts = [s[1] for s in sampled_data] + sampled_data = [s[0] for s in sampled_data] + + # Update our counters for the stopping criterion of the while loop. + if _return_metrics: + if max_agent_steps: + agent_or_env_steps += sum( + int(agent_stat) + for stat_dict in stats_dicts + for agent_stat in stat_dict[NUM_AGENT_STEPS_SAMPLED].values() + ) + else: + agent_or_env_steps += sum( + int(stat_dict[NUM_ENV_STEPS_SAMPLED]) for stat_dict in stats_dicts + ) + sample_batches_or_episodes.extend(sampled_data) + all_stats_dicts.extend(stats_dicts) + else: + for batch_or_episode in sampled_data: + if max_agent_steps: + agent_or_env_steps += ( + sum(e.agent_steps() for e in batch_or_episode) + if _uses_new_env_runners + else batch_or_episode.agent_steps() + ) + else: + agent_or_env_steps += ( + sum(e.env_steps() for e in batch_or_episode) + if _uses_new_env_runners + else batch_or_episode.env_steps() + ) + sample_batches_or_episodes.append(batch_or_episode) + # Break out (and ignore the remaining samples) if max timesteps (batch + # size) reached. We want to avoid collecting batches that are too large + # only because of a failed/restarted worker causing a second iteration + # of the main loop. + if ( + max_agent_or_env_steps is not None + and agent_or_env_steps >= max_agent_or_env_steps + ): + break + + if concat is True: + # If we have episodes flatten the episode list. + if _uses_new_env_runners: + sample_batches_or_episodes = tree.flatten(sample_batches_or_episodes) + # Otherwise we concatenate the `SampleBatch` objects + else: + sample_batches_or_episodes = concat_samples(sample_batches_or_episodes) + + if _return_metrics: + return sample_batches_or_episodes, all_stats_dicts + return sample_batches_or_episodes + + +@OldAPIStack +def standardize_fields(samples: SampleBatchType, fields: List[str]) -> SampleBatchType: + """Standardize fields of the given SampleBatch""" + wrapped = False + + if isinstance(samples, SampleBatch): + samples = samples.as_multi_agent() + wrapped = True + + for policy_id in samples.policy_batches: + batch = samples.policy_batches[policy_id] + for field in fields: + if field in batch: + batch[field] = standardized(batch[field]) + + if wrapped: + samples = samples.policy_batches[DEFAULT_POLICY_ID] + + return samples diff --git a/lib/python3.12/site-packages/ray/rllib/execution/segment_tree.py b/lib/python3.12/site-packages/ray/rllib/execution/segment_tree.py new file mode 100644 index 0000000000000000000000000000000000000000..5316fe34eecfcbff3dd98d59d39dc9ca33847571 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/segment_tree.py @@ -0,0 +1,217 @@ +import operator +from typing import Any, Optional + + +class SegmentTree: + """A Segment Tree data structure. + + https://en.wikipedia.org/wiki/Segment_tree + + Can be used as regular array, but with two important differences: + + a) Setting an item's value is slightly slower. It is O(lg capacity), + instead of O(1). + b) Offers efficient `reduce` operation which reduces the tree's values + over some specified contiguous subsequence of items in the array. + Operation could be e.g. min/max/sum. + + The data is stored in a list, where the length is 2 * capacity. + The second half of the list stores the actual values for each index, so if + capacity=8, values are stored at indices 8 to 15. The first half of the + array contains the reduced-values of the different (binary divided) + segments, e.g. (capacity=4): + 0=not used + 1=reduced-value over all elements (array indices 4 to 7). + 2=reduced-value over array indices (4 and 5). + 3=reduced-value over array indices (6 and 7). + 4-7: values of the tree. + NOTE that the values of the tree are accessed by indices starting at 0, so + `tree[0]` accesses `internal_array[4]` in the above example. + """ + + def __init__( + self, capacity: int, operation: Any, neutral_element: Optional[Any] = None + ): + """Initializes a Segment Tree object. + + Args: + capacity: Total size of the array - must be a power of two. + operation: Lambda obj, obj -> obj + The operation for combining elements (eg. sum, max). + Must be a mathematical group together with the set of + possible values for array elements. + neutral_element (Optional[obj]): The neutral element for + `operation`. Use None for automatically finding a value: + max: float("-inf"), min: float("inf"), sum: 0.0. + """ + + assert ( + capacity > 0 and capacity & (capacity - 1) == 0 + ), "Capacity must be positive and a power of 2!" + self.capacity = capacity + if neutral_element is None: + neutral_element = ( + 0.0 + if operation is operator.add + else float("-inf") + if operation is max + else float("inf") + ) + self.neutral_element = neutral_element + self.value = [self.neutral_element for _ in range(2 * capacity)] + self.operation = operation + + def reduce(self, start: int = 0, end: Optional[int] = None) -> Any: + """Applies `self.operation` to subsequence of our values. + + Subsequence is contiguous, includes `start` and excludes `end`. + + self.operation( + arr[start], operation(arr[start+1], operation(... arr[end]))) + + Args: + start: Start index to apply reduction to. + end (Optional[int]): End index to apply reduction to (excluded). + + Returns: + any: The result of reducing self.operation over the specified + range of `self._value` elements. + """ + if end is None: + end = self.capacity + elif end < 0: + end += self.capacity + + # Init result with neutral element. + result = self.neutral_element + # Map start/end to our actual index space (second half of array). + start += self.capacity + end += self.capacity + + # Example: + # internal-array (first half=sums, second half=actual values): + # 0 1 2 3 | 4 5 6 7 + # - 6 1 5 | 1 0 2 3 + + # tree.sum(0, 3) = 3 + # internally: start=4, end=7 -> sum values 1 0 2 = 3. + + # Iterate over tree starting in the actual-values (second half) + # section. + # 1) start=4 is even -> do nothing. + # 2) end=7 is odd -> end-- -> end=6 -> add value to result: result=2 + # 3) int-divide start and end by 2: start=2, end=3 + # 4) start still smaller end -> iterate once more. + # 5) start=2 is even -> do nothing. + # 6) end=3 is odd -> end-- -> end=2 -> add value to result: result=1 + # NOTE: This adds the sum of indices 4 and 5 to the result. + + # Iterate as long as start != end. + while start < end: + + # If start is odd: Add its value to result and move start to + # next even value. + if start & 1: + result = self.operation(result, self.value[start]) + start += 1 + + # If end is odd: Move end to previous even value, then add its + # value to result. NOTE: This takes care of excluding `end` in any + # situation. + if end & 1: + end -= 1 + result = self.operation(result, self.value[end]) + + # Divide both start and end by 2 to make them "jump" into the + # next upper level reduce-index space. + start //= 2 + end //= 2 + + # Then repeat till start == end. + + return result + + def __setitem__(self, idx: int, val: float) -> None: + """ + Inserts/overwrites a value in/into the tree. + + Args: + idx: The index to insert to. Must be in [0, `self.capacity`) + val: The value to insert. + """ + assert 0 <= idx < self.capacity, f"idx={idx} capacity={self.capacity}" + + # Index of the leaf to insert into (always insert in "second half" + # of the tree, the first half is reserved for already calculated + # reduction-values). + idx += self.capacity + self.value[idx] = val + + # Recalculate all affected reduction values (in "first half" of tree). + idx = idx >> 1 # Divide by 2 (faster than division). + while idx >= 1: + update_idx = 2 * idx # calculate only once + # Update the reduction value at the correct "first half" idx. + self.value[idx] = self.operation( + self.value[update_idx], self.value[update_idx + 1] + ) + idx = idx >> 1 # Divide by 2 (faster than division). + + def __getitem__(self, idx: int) -> Any: + assert 0 <= idx < self.capacity + return self.value[idx + self.capacity] + + def get_state(self): + return self.value + + def set_state(self, state): + assert len(state) == self.capacity * 2 + self.value = state + + +class SumSegmentTree(SegmentTree): + """A SegmentTree with the reduction `operation`=operator.add.""" + + def __init__(self, capacity: int): + super(SumSegmentTree, self).__init__(capacity=capacity, operation=operator.add) + + def sum(self, start: int = 0, end: Optional[Any] = None) -> Any: + """Returns the sum over a sub-segment of the tree.""" + return self.reduce(start, end) + + def find_prefixsum_idx(self, prefixsum: float) -> int: + """Finds highest i, for which: sum(arr[0]+..+arr[i - i]) <= prefixsum. + + Args: + prefixsum: `prefixsum` upper bound in above constraint. + + Returns: + int: Largest possible index (i) satisfying above constraint. + """ + assert 0 <= prefixsum <= self.sum() + 1e-5 + # Global sum node. + idx = 1 + + # Edge case when prefixsum can clip into the invalid regions + # https://github.com/ray-project/ray/issues/54284 + if prefixsum >= self.value[idx]: + prefixsum = self.value[idx] - 1e-5 + + # While non-leaf (first half of tree). + while idx < self.capacity: + update_idx = 2 * idx + if self.value[update_idx] > prefixsum: + idx = update_idx + else: + prefixsum -= self.value[update_idx] + idx = update_idx + 1 + return idx - self.capacity + + +class MinSegmentTree(SegmentTree): + def __init__(self, capacity: int): + super(MinSegmentTree, self).__init__(capacity=capacity, operation=min) + + def min(self, start: int = 0, end: Optional[Any] = None) -> Any: + """Returns min(arr[start], ..., arr[end])""" + return self.reduce(start, end) diff --git a/lib/python3.12/site-packages/ray/rllib/execution/train_ops.py b/lib/python3.12/site-packages/ray/rllib/execution/train_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..4de15d6ccff3b6fd0f779b4281f11471d897740a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/execution/train_ops.py @@ -0,0 +1,205 @@ +import logging +import math +from typing import Dict + +import numpy as np + +from ray._common.deprecation import deprecation_warning +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import ( + LEARN_ON_BATCH_TIMER, + LOAD_BATCH_TIMER, + NUM_AGENT_STEPS_TRAINED, + NUM_ENV_STEPS_TRAINED, +) +from ray.rllib.utils.metrics.learner_info import LearnerInfoBuilder +from ray.rllib.utils.sgd import do_minibatch_sgd +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + +logger = logging.getLogger(__name__) + + +@OldAPIStack +def train_one_step(algorithm, train_batch, policies_to_train=None) -> Dict: + """Function that improves the all policies in `train_batch` on the local worker. + + .. testcode:: + :skipif: True + + from ray.rllib.execution.rollout_ops import synchronous_parallel_sample + algo = [...] + train_batch = synchronous_parallel_sample(algo.env_runner_group) + # This trains the policy on one batch. + print(train_one_step(algo, train_batch))) + + .. testoutput:: + + {"default_policy": ...} + + Updates the NUM_ENV_STEPS_TRAINED and NUM_AGENT_STEPS_TRAINED counters as well as + the LEARN_ON_BATCH_TIMER timer of the `algorithm` object. + """ + config = algorithm.config + workers = algorithm.env_runner_group + local_worker = workers.local_env_runner + num_sgd_iter = config.get("num_epochs", config.get("num_sgd_iter", 1)) + minibatch_size = config.get("minibatch_size") + if minibatch_size is None: + minibatch_size = config.get("sgd_minibatch_size", 0) + + learn_timer = algorithm._timers[LEARN_ON_BATCH_TIMER] + with learn_timer: + # Subsample minibatches (size=`minibatch_size`) from the + # train batch and loop through train batch `num_sgd_iter` times. + if num_sgd_iter > 1 or minibatch_size > 0: + info = do_minibatch_sgd( + train_batch, + { + pid: local_worker.get_policy(pid) + for pid in policies_to_train + or local_worker.get_policies_to_train(train_batch) + }, + local_worker, + num_sgd_iter, + minibatch_size, + [], + ) + # Single update step using train batch. + else: + info = local_worker.learn_on_batch(train_batch) + + learn_timer.push_units_processed(train_batch.count) + algorithm._counters[NUM_ENV_STEPS_TRAINED] += train_batch.count + algorithm._counters[NUM_AGENT_STEPS_TRAINED] += train_batch.agent_steps() + + if algorithm.reward_estimators: + info[DEFAULT_POLICY_ID]["off_policy_estimation"] = {} + for name, estimator in algorithm.reward_estimators.items(): + info[DEFAULT_POLICY_ID]["off_policy_estimation"][name] = estimator.train( + train_batch + ) + return info + + +@OldAPIStack +def multi_gpu_train_one_step(algorithm, train_batch) -> Dict: + """Multi-GPU version of train_one_step. + + Uses the policies' `load_batch_into_buffer` and `learn_on_loaded_batch` methods + to be more efficient wrt CPU/GPU data transfers. For example, when doing multiple + passes through a train batch (e.g. for PPO) using `config.num_sgd_iter`, the + actual train batch is only split once and loaded once into the GPU(s). + + .. testcode:: + :skipif: True + + from ray.rllib.execution.rollout_ops import synchronous_parallel_sample + algo = [...] + train_batch = synchronous_parallel_sample(algo.env_runner_group) + # This trains the policy on one batch. + print(multi_gpu_train_one_step(algo, train_batch))) + + .. testoutput:: + + {"default_policy": ...} + + Updates the NUM_ENV_STEPS_TRAINED and NUM_AGENT_STEPS_TRAINED counters as well as + the LOAD_BATCH_TIMER and LEARN_ON_BATCH_TIMER timers of the Algorithm instance. + """ + if log_once("mulit_gpu_train_one_step_deprecation_warning"): + deprecation_warning( + old=("ray.rllib.execution.train_ops.multi_gpu_train_one_step") + ) + config = algorithm.config + workers = algorithm.env_runner_group + local_worker = workers.local_env_runner + num_sgd_iter = config.get("num_epochs", config.get("num_sgd_iter", 1)) + minibatch_size = config.get("minibatch_size") + if minibatch_size is None: + minibatch_size = config["train_batch_size"] + + # Determine the number of devices (GPUs or 1 CPU) we use. + num_devices = int(math.ceil(config["num_gpus"] or 1)) + + # Make sure total batch size is dividable by the number of devices. + # Batch size per tower. + per_device_batch_size = minibatch_size // num_devices + # Total batch size. + batch_size = per_device_batch_size * num_devices + assert batch_size % num_devices == 0 + assert batch_size >= num_devices, "Batch size too small!" + + # Handle everything as if multi-agent. + train_batch = train_batch.as_multi_agent() + + # Load data into GPUs. + load_timer = algorithm._timers[LOAD_BATCH_TIMER] + with load_timer: + num_loaded_samples = {} + for policy_id, batch in train_batch.policy_batches.items(): + # Not a policy-to-train. + if ( + local_worker.is_policy_to_train is not None + and not local_worker.is_policy_to_train(policy_id, train_batch) + ): + continue + + # Decompress SampleBatch, in case some columns are compressed. + batch.decompress_if_needed() + + # Load the entire train batch into the Policy's only buffer + # (idx=0). Policies only have >1 buffers, if we are training + # asynchronously. + num_loaded_samples[policy_id] = local_worker.policy_map[ + policy_id + ].load_batch_into_buffer(batch, buffer_index=0) + + # Execute minibatch SGD on loaded data. + learn_timer = algorithm._timers[LEARN_ON_BATCH_TIMER] + with learn_timer: + # Use LearnerInfoBuilder as a unified way to build the final + # results dict from `learn_on_loaded_batch` call(s). + # This makes sure results dicts always have the same structure + # no matter the setup (multi-GPU, multi-agent, minibatch SGD, + # tf vs torch). + learner_info_builder = LearnerInfoBuilder(num_devices=num_devices) + + for policy_id, samples_per_device in num_loaded_samples.items(): + policy = local_worker.policy_map[policy_id] + num_batches = max(1, int(samples_per_device) // int(per_device_batch_size)) + logger.debug("== sgd epochs for {} ==".format(policy_id)) + for _ in range(num_sgd_iter): + permutation = np.random.permutation(num_batches) + for batch_index in range(num_batches): + # Learn on the pre-loaded data in the buffer. + # Note: For minibatch SGD, the data is an offset into + # the pre-loaded entire train batch. + results = policy.learn_on_loaded_batch( + permutation[batch_index] * per_device_batch_size, buffer_index=0 + ) + + learner_info_builder.add_learn_on_batch_results(results, policy_id) + + # Tower reduce and finalize results. + learner_info = learner_info_builder.finalize() + + load_timer.push_units_processed(train_batch.count) + learn_timer.push_units_processed(train_batch.count) + + # TODO: Move this into Algorithm's `training_step` method for + # better transparency. + algorithm._counters[NUM_ENV_STEPS_TRAINED] += train_batch.count + algorithm._counters[NUM_AGENT_STEPS_TRAINED] += train_batch.agent_steps() + + if algorithm.reward_estimators: + learner_info[DEFAULT_POLICY_ID]["off_policy_estimation"] = {} + for name, estimator in algorithm.reward_estimators.items(): + learner_info[DEFAULT_POLICY_ID]["off_policy_estimation"][ + name + ] = estimator.train(train_batch) + + return learner_info