R2D2

.gitignore

@@ -1,3 +1,6 @@
 bazel-*
 __pycache__
 uv.lock
+# mlflow
+mlruns
+logs

.vscode/settings.json

@@ -24,5 +24,6 @@
     "bazel.buildifierExecutable": "@buildifier_prebuilt//:buildifier",
     "python.testing.pytestArgs": [
         "earl"
-    ]
+    ],
+    "git.ignoreLimitWarning": true
 }

BUILD.bazel

@@ -8,8 +8,10 @@ pip_compile(
         "--emit-index-url",
         "--no-strip-extras",
         "--extra=test",
+        "--extra=agent-r2d2",
+        "--index=https://download.pytorch.org/whl/cpu",
     ],
-    python_platform = "x86_64-unknown-linux-gnu",
+    python_platform = "x86_64-manylinux_2_28",  # envpool needs at least 2_24
     requirements_in = "//:pyproject.toml",
     requirements_txt = "requirements_linux_x86_64.txt",
 )

MODULE.bazel

@@ -6,15 +6,15 @@ module(
 )
 
 # BEGIN python toolchain
-_PYTHON_VERSION = "3.12"
+_PYTHON_VERSION = "3.11"  # latest that supports envpool
 
 bazel_dep(name = "rules_python", version = "1.1.0", dev_dependency = True)
 
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(python_version = _PYTHON_VERSION)
 # END python toolchain
 
-# BEGIN python dependencies
+# BEGIN python dependenciesP
 bazel_dep(name = "rules_uv", version = "0.53.0", dev_dependency = True)
 
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")

earl/agents/r2d2/BUILD.bazel

@@ -0,0 +1,67 @@
+load("@aspect_rules_py//py:defs.bzl", "py_library")
+load("//tools/py_test:py_test.bzl", "py_test")
+
+py_library(
+    name = "r2d2",
+    srcs = [
+        "networks.py",
+        "r2d2.py",
+        "utils.py",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//earl:core",
+        "@pypi//chex",
+        "@pypi//distrax",
+        "@pypi//equinox",
+        "@pypi//jax",
+        "@pypi//jaxtyping",
+        "@pypi//optax",
+    ],
+)
+
+py_test(
+    name = "test_r2d2",
+    timeout = "long",
+    srcs = ["test_r2d2.py"],
+    filterwarnings = [
+        "ignore:jax.interpreters.xla.pytype_aval_mappings is deprecated.:DeprecationWarning",
+        "ignore:Shape is deprecated; use StableHLO instead.:DeprecationWarning",
+    ],
+    shard_count = 2,
+    deps = [
+        ":r2d2",
+        "//earl:core",
+        "//earl/environment_loop:gymnasium_loop",
+        "//earl/environment_loop:gymnax_loop",
+        "@pypi//envpool",
+        "@pypi//gymnasium",
+        "@pypi//gymnax",
+        "@pypi//jax",
+        "@pypi//jax_loop_utils",
+        "@pypi//numpy",
+        "@pypi//optax",
+        "@pypi//pytest",
+    ],
+)
+
+py_test(
+    name = "test_r2d2_learns",
+    timeout = "long",
+    srcs = ["test_r2d2_learns.py"],
+    filterwarnings = [
+        "ignore:jax.interpreters.xla.pytype_aval_mappings is deprecated.:DeprecationWarning",
+        "ignore:Shape is deprecated; use StableHLO instead.:DeprecationWarning",
+    ],
+    tags = ["manual"],
+    deps = [
+        ":r2d2",
+        "//earl:core",
+        "//earl/environment_loop:gymnax_loop",
+        "@pypi//gymnax",
+        "@pypi//jax",
+        "@pypi//jax_loop_utils",
+        "@pypi//numpy",
+        "@pypi//pytest",
+    ],
+)