[RLlib-contrib] Decision Transformer

.buildkite/pipeline.ml.yml

@@ -564,3 +564,12 @@
     - (cd rllib_contrib/maml && pip install -r requirements.txt && pip install -e .)
     - ./ci/env/env_info.sh
     - pytest rllib_contrib/maml/tests/test_maml.py
+
+- label: ":exploding_death_star: RLlib Contrib: Decision Transformer Tests"
+  conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_RLLIB_CONTRIB_AFFECTED"]
+  commands:
+    - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
+    - (cd rllib_contrib/dt && pip install -r requirements.txt && pip install -e .)
+    - ./ci/env/env_info.sh
+    - pytest rllib_contrib/dt/tests/
+    - python rllib_contrib/dt/examples/dt_cartpole_v1.py --run-as-test

rllib_contrib/TOC.md

@@ -3,6 +3,7 @@
 
 * [A3C](./a3c)
 * [MAML](./maml)
+* [Decision Transformer](./dt)
 
 
 

rllib_contrib/dt/README.md

@@ -0,0 +1,18 @@
+# Decision Transformer
+
+[Decision Transformer](https://arxiv.org/abs/2106.01345) is an offline-rl algorithm that trains a transformer to generate
+optimal actions based on desired returns, past states, and actions.
+
+
+## Installation
+
+```
+conda create -n rllib-dt python=3.10
+conda activate rllib-dt
+pip install -r requirements.txt
+pip install -e '.[development]'
+```
+
+## Usage
+
+[DT Example]()

rllib_contrib/dt/examples/dt_cartpole_v1.py

@@ -0,0 +1,84 @@
+import argparse
+
+from rllib_dt.dt import DT, DTConfig
+
+import ray
+from ray import air, tune
+from ray.rllib.utils.test_utils import check_learning_achieved
+
+
+def get_cli_args():
+    """Create CLI parser and return parsed arguments"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--run-as-test", action="store_true", default=False)
+    args = parser.parse_args()
+    print(f"Running with following CLI args: {args}")
+    return args
+
+
+if __name__ == "__main__":
+    args = get_cli_args()
+
+    ray.init()
+    config = (
+        DTConfig()
+        .environment(env="CartPole-v1", clip_actions=True)
+        .framework("torch")
+        .offline_data(
+            input_="dataset",
+            input_config={
+                "format": "json",
+                "paths": ["s3://anonymous@air-example-data/rllib/cartpole/large.json"],
+            },
+            actions_in_input_normalized=True,
+        )
+        .training(
+            train_batch_size=512,
+            lr=0.01,
+            optimizer={
+                "weight_decay": 0.1,
+                "betas": [0.9, 0.999],
+            },
+            replay_buffer_config={"capacity": 20},
+            # model
+            model={"max_seq_len": 3},
+            num_layers=1,
+            num_heads=1,
+            embed_dim=64,
+            horizon=500,
+        )
+        .evaluation(
+            evaluation_interval=1,
+            evaluation_num_workers=1,
+            evaluation_duration=10,
+            target_return=200,
+            evaluation_duration_unit="episodes",
+            evaluation_parallel_to_training=True,
+            evaluation_config=DTConfig.overrides(input_="sampler", explore=False),
+        )
+        # Episode horizon: Must match environment's time limit, if any.
+        .rollouts(num_rollout_workers=3)
+        .reporting(min_train_timesteps_per_iteration=5000)
+    )
+
+    stop_reward = 200
+
+    tuner = tune.Tuner(
+        DT,
+        param_space=config.to_dict(),
+        run_config=air.RunConfig(
+            stop={
+                "evaluation/sampler_results/episode_reward_mean": stop_reward,
+                "training_iteration": 100,
+            },
+            failure_config=air.FailureConfig(fail_fast="raise"),
+        ),
+    )
+    results = tuner.fit()
+
+    if args.run_as_test:
+        check_learning_achieved(
+            results,
+            stop_reward,
+            metric="evaluation/sampler_results/episode_reward_mean",
+        )

rllib_contrib/dt/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[project]
+name = "rllib-dt"
+authors = [{name = "Anyscale Inc."}]
+version = "0.1.0"
+description = ""
+readme = "README.md"
+requires-python = ">=3.7, <3.11"
+dependencies = ["gymnasium", "ray[rllib]==2.5.0"]
+
+[project.optional-dependencies]
+development = ["pytest>=7.2.2", "pre-commit==2.21.0", "torch==1.12.0"]

rllib_contrib/dt/requirements.txt

@@ -0,0 +1 @@
+torch==1.12.0

rllib_contrib/dt/src/rllib_dt/dt/__init__.py

@@ -0,0 +1,9 @@
+from rllib_dt.dt.dt import DT, DTConfig
+from rllib_dt.dt.dt_torch_model import DTTorchModel
+from rllib_dt.dt.dt_torch_policy import DTTorchPolicy
+
+from ray.tune.registry import register_trainable
+
+__all__ = ["DT", "DTConfig", "DTTorchModel", "DTTorchPolicy"]
+
+register_trainable("rllib-contrib-dt", DT)