chore: import upstream snapshot with attribution

rllib/core/testing/bc_algorithm.py

@@ -0,0 +1,40 @@
+"""Contains example implementation of a custom algorithm.
+
+Note: It doesn't include any real use-case functionality; it only serves as an example
+to test the algorithm construction and customization.
+"""
+
+from ray.rllib.algorithms import Algorithm, AlgorithmConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.core.testing.torch.bc_learner import BCTorchLearner
+from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule
+from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import ResultDict
+
+
+class BCConfigTest(AlgorithmConfig):
+    def __init__(self, algo_class=None):
+        super().__init__(algo_class=algo_class or BCAlgorithmTest)
+
+    def get_default_rl_module_spec(self):
+        if self.framework_str == "torch":
+            return RLModuleSpec(module_class=DiscreteBCTorchModule)
+
+    def get_default_learner_class(self):
+        if self.framework_str == "torch":
+            return BCTorchLearner
+
+
+class BCAlgorithmTest(Algorithm):
+    @classmethod
+    def get_default_policy_class(cls, config: AlgorithmConfig):
+        if config.framework_str == "torch":
+            return TorchPolicyV2
+        else:
+            raise ValueError("Unknown framework: {}".format(config.framework_str))
+
+    @override(Algorithm)
+    def training_step(self) -> ResultDict:
+        # do nothing.
+        return {}

rllib/core/testing/testing_learner.py

@@ -0,0 +1,67 @@
+from typing import Type
+
+import numpy as np
+
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.learner.learner import Learner
+from ray.rllib.core.rl_module.multi_rl_module import (
+    MultiRLModule,
+    MultiRLModuleSpec,
+)
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.typing import RLModuleSpecType
+
+
+class BaseTestingAlgorithmConfig(AlgorithmConfig):
+    # A test setting to activate metrics on mean weights.
+    report_mean_weights: bool = True
+
+    @override(AlgorithmConfig)
+    def get_default_learner_class(self) -> Type["Learner"]:
+        if self.framework_str == "torch":
+            from ray.rllib.core.testing.torch.bc_learner import BCTorchLearner
+
+            return BCTorchLearner
+        else:
+            raise ValueError(f"Unsupported framework: {self.framework_str}")
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self) -> "RLModuleSpecType":
+        if self.framework_str == "torch":
+            from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule
+
+            cls = DiscreteBCTorchModule
+        else:
+            raise ValueError(f"Unsupported framework: {self.framework_str}")
+
+        spec = RLModuleSpec(
+            module_class=cls,
+            model_config={"fcnet_hiddens": [32]},
+        )
+
+        if self.is_multi_agent:
+            # TODO (Kourosh): Make this more multi-agent for example with policy ids
+            #  "1" and "2".
+            return MultiRLModuleSpec(
+                multi_rl_module_class=MultiRLModule,
+                rl_module_specs={DEFAULT_MODULE_ID: spec},
+            )
+        else:
+            return spec
+
+
+class BaseTestingLearner(Learner):
+    @override(Learner)
+    def after_gradient_based_update(self, *, timesteps):
+        # This is to check if in the multi-gpu case, the weights across workers are
+        # the same. It is really only needed during testing.
+        if self.config.report_mean_weights:
+            for module_id in self.module.keys():
+                parameters = convert_to_numpy(
+                    self.get_parameters(self.module[module_id])
+                )
+                mean_ws = np.mean([w.mean() for w in parameters])
+                self.metrics.log_value((module_id, "mean_weight"), mean_ws, window=1)

rllib/core/testing/torch/bc_learner.py

@@ -0,0 +1,35 @@
+from typing import TYPE_CHECKING, Any, Dict
+
+import torch
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.learner.torch.torch_learner import TorchLearner
+from ray.rllib.core.testing.testing_learner import BaseTestingLearner
+from ray.rllib.utils.typing import ModuleID, TensorType
+
+if TYPE_CHECKING:
+    from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+
+
+class BCTorchLearner(TorchLearner, BaseTestingLearner):
+    def compute_loss_for_module(
+        self,
+        *,
+        module_id: ModuleID,
+        config: "AlgorithmConfig",
+        batch: Dict[str, Any],
+        fwd_out: Dict[str, TensorType],
+    ) -> TensorType:
+        BaseTestingLearner.compute_loss_for_module(
+            self,
+            module_id=module_id,
+            config=config,
+            batch=batch,
+            fwd_out=fwd_out,
+        )
+        action_dist_inputs = fwd_out[Columns.ACTION_DIST_INPUTS]
+        action_dist_class = self._module[module_id].get_train_action_dist_cls()
+        action_dist = action_dist_class.from_logits(action_dist_inputs)
+        loss = -torch.mean(action_dist.logp(batch[Columns.ACTIONS]))
+
+        return loss

rllib/core/testing/torch/bc_module.py

@@ -0,0 +1,149 @@
+from typing import Any, Dict
+
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.distribution.torch.torch_distribution import TorchCategorical
+from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule
+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
+
+torch, nn = try_import_torch()
+
+
+class DiscreteBCTorchModule(TorchRLModule):
+    def setup(self):
+        input_dim = self.observation_space.shape[0]
+        hidden_dim = self.model_config["fcnet_hiddens"][0]
+        output_dim = self.action_space.n
+
+        self.policy = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, output_dim),
+        )
+
+        self.input_dim = input_dim
+
+    def get_train_action_dist_cls(self):
+        return TorchCategorical
+
+    def get_exploration_action_dist_cls(self):
+        return TorchCategorical
+
+    def get_inference_action_dist_cls(self):
+        return TorchCategorical
+
+    @override(RLModule)
+    def _forward_inference(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        with torch.no_grad():
+            return self._forward_train(batch)
+
+    @override(RLModule)
+    def _forward_exploration(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        with torch.no_grad():
+            return self._forward_train(batch)
+
+    @override(RLModule)
+    def _forward_train(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        action_logits = self.policy(batch["obs"])
+        return {Columns.ACTION_DIST_INPUTS: action_logits}
+
+
+class BCTorchRLModuleWithSharedGlobalEncoder(TorchRLModule):
+    """An example of an RLModule that uses an encoder shared with other things.
+
+    For example, we could consider a multi-agent case where for inference each agent
+    needs to know the global state of the environment, as well as the local state of
+    itself. For better representation learning we would like to share the encoder
+    across all the modules. So this module simply accepts the encoder object as its
+    input argument and uses it to encode the global state. The local state is passed
+    through as is. The policy head is then a simple MLP that takes the concatenation of
+    the global and local state as input and outputs the action logits.
+
+    """
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        local_dim: int,
+        hidden_dim: int,
+        action_dim: int,
+        config=None,
+    ) -> None:
+        super().__init__(config=config)
+
+        self.encoder = encoder
+        self.policy_head = nn.Sequential(
+            nn.Linear(hidden_dim + local_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, action_dim),
+        )
+
+    def get_train_action_dist_cls(self):
+        return TorchCategorical
+
+    def get_exploration_action_dist_cls(self):
+        return TorchCategorical
+
+    def get_inference_action_dist_cls(self):
+        return TorchCategorical
+
+    @override(RLModule)
+    def _default_input_specs(self):
+        return [("obs", "global"), ("obs", "local")]
+
+    @override(RLModule)
+    def _forward_inference(self, batch):
+        with torch.no_grad():
+            return self._common_forward(batch)
+
+    @override(RLModule)
+    def _forward_exploration(self, batch):
+        with torch.no_grad():
+            return self._common_forward(batch)
+
+    @override(RLModule)
+    def _forward_train(self, batch):
+        return self._common_forward(batch)
+
+    def _common_forward(self, batch):
+        obs = batch["obs"]
+        global_enc = self.encoder(obs["global"])
+        policy_in = torch.cat([global_enc, obs["local"]], dim=-1)
+        action_logits = self.policy_head(policy_in)
+
+        return {Columns.ACTION_DIST_INPUTS: action_logits}
+
+
+class BCTorchMultiAgentModuleWithSharedEncoder(MultiRLModule):
+    def setup(self):
+        module_specs = self.config.modules
+        module_spec = next(iter(module_specs.values()))
+        global_dim = module_spec.observation_space["global"].shape[0]
+        hidden_dim = module_spec.model_config_dict["fcnet_hiddens"][0]
+        shared_encoder = nn.Sequential(
+            nn.Linear(global_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+
+        rl_modules = {}
+        for module_id, module_spec in module_specs.items():
+            rl_modules[module_id] = module_spec.module_class(
+                config=self.config.modules[module_id].get_rl_module_config(),
+                encoder=shared_encoder,
+                local_dim=module_spec.observation_space["local"].shape[0],
+                hidden_dim=hidden_dim,
+                action_dim=module_spec.action_space.n,
+            )
+
+        self._rl_modules = rl_modules
+
+    def serialize(self):
+        # TODO (Kourosh): Implement when needed.
+        raise NotImplementedError
+
+    def deserialize(self, data):
+        # TODO (Kourosh): Implement when needed.
+        raise NotImplementedError