chore: import upstream snapshot with attribution

rllib/algorithms/bc/__init__.py

@@ -0,0 +1,6 @@
+from ray.rllib.algorithms.bc.bc import BC, BCConfig
+
+__all__ = [
+    "BC",
+    "BCConfig",
+]

rllib/algorithms/bc/bc.py

@@ -0,0 +1,120 @@
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.algorithms.marwil.marwil import MARWIL, MARWILConfig
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import RLModuleSpecType
+
+
+class BCConfig(MARWILConfig):
+    """Defines a configuration class from which a new BC Algorithm can be built
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.algorithms.bc import BCConfig
+        # Run this from the ray directory root.
+        config = BCConfig().training(lr=0.00001, gamma=0.99)
+        config = config.offline_data(
+            input_="./rllib/offline/tests/data/cartpole/large.json")
+
+        # Build an Algorithm object from the config and run 1 training iteration.
+        algo = config.build()
+        algo.train()
+
+    .. testcode::
+        :skipif: True
+
+        from ray.rllib.algorithms.bc import BCConfig
+        from ray import tune
+        config = BCConfig()
+        # Print out some default values.
+        print(config.beta)
+        # Update the config object.
+        config.training(
+            lr=tune.grid_search([0.001, 0.0001]), beta=0.75
+        )
+        # Set the config object's data path.
+        # Run this from the ray directory root.
+        config.offline_data(
+            input_="./rllib/offline/tests/data/cartpole/large.json"
+        )
+        # Set the config object's env, used for evaluation.
+        config.environment(env="CartPole-v1")
+        # Use to_dict() to get the old-style python config dict
+        # when running with tune.
+        tune.Tuner(
+            "BC",
+            param_space=config.to_dict(),
+        ).fit()
+    """
+
+    def __init__(self, algo_class=None):
+        super().__init__(algo_class=algo_class or BC)
+
+        # fmt: off
+        # __sphinx_doc_begin__
+        # No need to calculate advantages (or do anything else with the rewards).
+        self.beta = 0.0
+        # Advantages (calculated during postprocessing)
+        # not important for behavioral cloning.
+        self.postprocess_inputs = False
+
+        # Materialize only the mapped data. This is optimal as long
+        # as no connector in the connector pipeline holds a state.
+        self.materialize_data = False
+        self.materialize_mapped_data = True
+        # __sphinx_doc_end__
+        # fmt: on
+
+    @override(AlgorithmConfig)
+    def get_default_rl_module_spec(self) -> RLModuleSpecType:
+        if self.framework_str == "torch":
+            from ray.rllib.algorithms.bc.torch.default_bc_torch_rl_module import (
+                DefaultBCTorchRLModule,
+            )
+
+            return RLModuleSpec(module_class=DefaultBCTorchRLModule)
+        else:
+            raise ValueError(
+                f"The framework {self.framework_str} is not supported. "
+                "Use `torch` instead."
+            )
+
+    @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,
+        )
+
+        # Remove unneeded connectors from the MARWIL connector pipeline.
+        pipeline.remove("AddOneTsToEpisodesAndTruncate")
+        pipeline.remove("GeneralAdvantageEstimation")
+
+        return pipeline
+
+    @override(MARWILConfig)
+    def validate(self) -> None:
+        # Call super's validation method.
+        super().validate()
+
+        if self.beta != 0.0:
+            self._value_error("For behavioral cloning, `beta` parameter must be 0.0!")
+
+
+class BC(MARWIL):
+    """Behavioral Cloning (derived from MARWIL).
+
+    Uses MARWIL with beta force-set to 0.0.
+    """
+
+    @classmethod
+    @override(MARWIL)
+    def get_default_config(cls) -> BCConfig:
+        return BCConfig()

rllib/algorithms/bc/bc_catalog.py

@@ -0,0 +1,112 @@
+# __sphinx_doc_begin__
+import gymnasium as gym
+
+from ray.rllib.algorithms.ppo.ppo_catalog import _check_if_diag_gaussian
+from ray.rllib.core.models.base import Model
+from ray.rllib.core.models.catalog import Catalog
+from ray.rllib.core.models.configs import FreeLogStdMLPHeadConfig, MLPHeadConfig
+from ray.rllib.utils.annotations import OverrideToImplementCustomLogic
+
+
+class BCCatalog(Catalog):
+    """The Catalog class used to build models for BC.
+
+    BCCatalog provides the following models:
+        - Encoder: The encoder used to encode the observations.
+        - Pi Head: The head used for the policy logits.
+
+    The default encoder is chosen by RLlib dependent on the observation space.
+    See `ray.rllib.core.models.encoders::Encoder` for details. To define the
+    network architecture use the `model_config_dict[fcnet_hiddens]` and
+    `model_config_dict[fcnet_activation]`.
+
+    To implement custom logic, override `BCCatalog.build_encoder()` or modify the
+    `EncoderConfig` at `BCCatalog.encoder_config`.
+
+    Any custom head can be built by overriding the `build_pi_head()` method.
+    Alternatively, the `PiHeadConfig` can be overridden to build a custom
+    policy head during runtime. To change solely the network architecture,
+    `model_config_dict["head_fcnet_hiddens"]` and
+    `model_config_dict["head_fcnet_activation"]` can be used.
+    """
+
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        model_config_dict: dict,
+    ):
+        """Initializes the BCCatalog.
+
+        Args:
+            observation_space: The observation space if the Encoder.
+            action_space: The action space for the Pi Head.
+            model_cnfig_dict: The model config to use..
+        """
+        super().__init__(
+            observation_space=observation_space,
+            action_space=action_space,
+            model_config_dict=model_config_dict,
+        )
+
+        self.pi_head_hiddens = self._model_config_dict["head_fcnet_hiddens"]
+        self.pi_head_activation = self._model_config_dict["head_fcnet_activation"]
+
+        # At this time we do not have the precise (framework-specific) action
+        # distribution class, i.e. we do  not know the output dimension of the
+        # policy head. The config for the policy head is therefore build in the
+        # `self.build_pi_head()` method.
+        self.pi_head_config = None
+
+    @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 BC specific RLModule implementation.
+
+        Args:
+            framework: The framework to use. Either "torch" or "tf2".
+
+        Returns:
+            The policy head.
+        """
+
+        # Define the output dimension via the action distribution.
+        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
+        )
+        # With the action distribution class and the number of outputs defined,
+        # we can build the config for the policy head.
+        pi_head_config_cls = (
+            FreeLogStdMLPHeadConfig
+            if self._model_config_dict["free_log_std"]
+            else MLPHeadConfig
+        )
+        self.pi_head_config = pi_head_config_cls(
+            input_dims=self._latent_dims,
+            hidden_layer_dims=self.pi_head_hiddens,
+            hidden_layer_activation=self.pi_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)
+
+
+# __sphinx_doc_end__

rllib/algorithms/bc/tests/test_bc.py

@@ -0,0 +1,146 @@
+import unittest
+from pathlib import Path
+
+import ray
+from ray.rllib.algorithms.bc import BCConfig
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.metrics import (
+    ENV_RUNNER_RESULTS,
+    EPISODE_RETURN_MEAN,
+    EVALUATION_RESULTS,
+    LEARNER_RESULTS,
+    NUM_ENV_STEPS_SAMPLED_LIFETIME,
+)
+
+
+class TestBC(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        ray.shutdown()
+
+    def test_bc_compilation_and_learning_from_offline_file(self):
+        # Define the data paths.
+        data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+        base_path = Path(__file__).parents[3]
+        print(f"base_path={base_path}")
+        data_path = "local://" / base_path / data_path
+
+        print(f"data_path={data_path}")
+
+        # Define the BC config.
+        config = (
+            BCConfig()
+            .environment(env="CartPole-v1")
+            .learners(
+                num_learners=0,
+            )
+            .evaluation(
+                evaluation_interval=3,
+                evaluation_num_env_runners=1,
+                evaluation_duration=5,
+                evaluation_parallel_to_training=True,
+            )
+            # Note, the `input_` argument is the major argument for the
+            # new offline API.
+            .offline_data(
+                input_=[data_path.as_posix()],
+                dataset_num_iters_per_learner=1,
+            )
+            .training(
+                lr=0.0008,
+                train_batch_size_per_learner=2000,
+            )
+        )
+
+        num_iterations = 350
+        min_return_to_reach = 120.0
+
+        # TODO (simon): Add support for recurrent modules.
+        algo = config.build()
+        learnt = False
+        for i in range(num_iterations):
+            results = algo.train()
+            print(results)
+
+            eval_results = results.get(EVALUATION_RESULTS, {})
+            if eval_results:
+                episode_return_mean = eval_results[ENV_RUNNER_RESULTS][
+                    EPISODE_RETURN_MEAN
+                ]
+                print(f"iter={i}, R={episode_return_mean}")
+                if episode_return_mean > min_return_to_reach:
+                    print("BC has learnt the task!")
+                    learnt = True
+                    break
+
+        if not learnt:
+            raise ValueError(
+                f"`BC` did not reach {min_return_to_reach} reward from "
+                "expert offline data!"
+            )
+
+        algo.stop()
+
+    def test_bc_lr_schedule(self):
+        # Define the data paths.
+        data_path = "offline/tests/data/cartpole/cartpole-v1_large"
+        base_path = Path(__file__).parents[3]
+        data_path = "local://" / base_path / data_path
+
+        config = (
+            BCConfig()
+            .environment(env="CartPole-v1")
+            .learners(
+                num_learners=0,
+            )
+            .evaluation(
+                evaluation_interval=3,
+                evaluation_num_env_runners=1,
+                evaluation_duration=5,
+                evaluation_parallel_to_training=True,
+            )
+            # Note, the `input_` argument is the major argument for the
+            # new offline API.
+            .offline_data(
+                input_=[data_path.as_posix()],
+                dataset_num_iters_per_learner=1,
+            )
+            .training(
+                lr=[
+                    [0, 0.001],
+                    [3000, 0.01],
+                ],
+                train_batch_size_per_learner=2000,
+            )
+        )
+        algo = config.build()
+
+        done = False
+        while not done:
+            results = algo.train()
+            ts = results[NUM_ENV_STEPS_SAMPLED_LIFETIME]
+            assert ts > 0
+            lr = results[LEARNER_RESULTS][DEFAULT_POLICY_ID][
+                "default_optimizer_learning_rate"
+            ]
+            if ts < 3000:
+                # The learning rate should be linearly interpolated.
+                expected_lr = 0.001 + (ts / 3000) * (0.01 - 0.001)
+                self.assertAlmostEqual(lr, expected_lr, places=6)
+            else:
+                self.assertEqual(lr, 0.01)
+                done = True
+
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/bc/torch/default_bc_torch_rl_module.py

@@ -0,0 +1,56 @@
+import abc
+from typing import Any, Dict
+
+from ray.rllib.algorithms.bc.bc_catalog import BCCatalog
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.models.base import ENCODER_OUT
+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.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class DefaultBCTorchRLModule(TorchRLModule, abc.ABC):
+    """The default TorchRLModule used, if no custom RLModule is provided.
+
+    Builds an encoder net based on the observation space.
+    Builds a pi head based on the action space.
+
+    Passes observations from the input batch through the encoder, then the pi head to
+    compute action logits.
+    """
+
+    def __init__(self, *args, **kwargs):
+        catalog_class = kwargs.pop("catalog_class", None)
+        if catalog_class is None:
+            catalog_class = BCCatalog
+        super().__init__(*args, **kwargs, catalog_class=catalog_class)
+
+    @override(RLModule)
+    def setup(self):
+        # Build model components (encoder and pi head) from catalog.
+        super().setup()
+        self._encoder = self.catalog.build_encoder(framework=self.framework)
+        self._pi_head = self.catalog.build_pi_head(framework=self.framework)
+
+    @override(TorchRLModule)
+    def _forward(self, batch: Dict, **kwargs) -> Dict[str, Any]:
+        """Generic BC forward pass (for all phases of training/evaluation)."""
+        # Encoder embeddings.
+        encoder_outs = self._encoder(batch)
+        # Action dist inputs.
+        outputs = {Columns.ACTION_DIST_INPUTS: self._pi_head(encoder_outs[ENCODER_OUT])}
+
+        # Add the state if the encoder is stateful.
+        if Columns.STATE_OUT in encoder_outs:
+            outputs[Columns.STATE_OUT] = encoder_outs[Columns.STATE_OUT]
+        # Return the outputs.
+        return outputs
+
+    @override(RLModule)
+    def get_initial_state(self) -> dict:
+        if hasattr(self._encoder, "get_initial_state"):
+            return self._encoder.get_initial_state()
+        else:
+            return {}