chore: import upstream snapshot with attribution

rllib/examples/_old_api_stack/catalogs/mobilenet_v2_encoder.py

@@ -0,0 +1,78 @@
+# @OldAPIStack
+
+"""
+This example shows two modifications:
+- How to write a custom Encoder (using MobileNet v2)
+- How to enhance Catalogs with this custom Encoder
+
+With the pattern shown in this example, we can enhance Catalogs such that they extend
+to new observation- or action spaces while retaining their original functionality.
+"""
+# __sphinx_doc_begin__
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.algorithms.ppo.ppo import PPOConfig
+from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
+from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+from ray.rllib.examples._old_api_stack.models.mobilenet_v2_encoder import (
+    MOBILENET_INPUT_SHAPE,
+    MobileNetV2EncoderConfig,
+)
+from ray.rllib.examples.envs.classes.random_env import RandomEnv
+
+
+# Define a PPO Catalog that we can use to inject our MobileNetV2 Encoder into RLlib's
+# decision tree of what model to choose
+class MobileNetEnhancedPPOCatalog(PPOCatalog):
+    @classmethod
+    def _get_encoder_config(
+        cls,
+        observation_space: gym.Space,
+        **kwargs,
+    ):
+        if (
+            isinstance(observation_space, gym.spaces.Box)
+            and observation_space.shape == MOBILENET_INPUT_SHAPE
+        ):
+            # Inject our custom encoder here, only if the observation space fits it
+            return MobileNetV2EncoderConfig()
+        else:
+            return super()._get_encoder_config(observation_space, **kwargs)
+
+
+# Create a generic config with our enhanced Catalog
+ppo_config = (
+    PPOConfig()
+    .rl_module(rl_module_spec=RLModuleSpec(catalog_class=MobileNetEnhancedPPOCatalog))
+    .env_runners(num_env_runners=0)
+    # The following training settings make it so that a training iteration is very
+    # quick. This is just for the sake of this example. PPO will not learn properly
+    # with these settings!
+    .training(train_batch_size_per_learner=32, minibatch_size=16, num_epochs=1)
+)
+
+# CartPole's observation space is not compatible with our MobileNetV2 Encoder, so
+# this will use the default behaviour of Catalogs
+ppo_config.environment("CartPole-v1")
+results = ppo_config.build().train()
+print(results)
+
+# For this training, we use a RandomEnv with observations of shape
+# MOBILENET_INPUT_SHAPE. This will use our custom Encoder.
+ppo_config.environment(
+    RandomEnv,
+    env_config={
+        "action_space": gym.spaces.Discrete(2),
+        # Test a simple Image observation space.
+        "observation_space": gym.spaces.Box(
+            0.0,
+            1.0,
+            shape=MOBILENET_INPUT_SHAPE,
+            dtype=np.float32,
+        ),
+    },
+)
+results = ppo_config.build().train()
+print(results)
+# __sphinx_doc_end__

rllib/examples/_old_api_stack/catalogs/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.examples.utils import add_rllib_example_script_args, check
+from ray.rllib.utils.framework import try_import_torch
+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()
+
+    # 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")