chore: import upstream snapshot with attribution

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

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"])

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

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"])

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