[envpool] Fix Linux test workflow compatibility

.bazelversion

@@ -0,0 +1 @@
+6.0.0

Makefile

@@ -7,6 +7,9 @@ BAZEL_FILES    = $(shell find . -type f -name "*BUILD" -o -name "*.bzl")
 COMMIT_HASH    = $(shell git log -1 --format=%h)
 COPYRIGHT      = "Garena Online Private Limited"
 BAZELOPT       =
+BAZELISK_BIN   = $(shell command -v bazelisk 2>/dev/null || echo $(HOME)/go/bin/bazelisk)
+BAZEL_VERSION  = 6.0.0
+BAZEL          = USE_BAZEL_VERSION=$(BAZEL_VERSION) $(BAZELISK_BIN)
 DATE           = $(shell date "+%Y-%m-%d")
 DOCKER_TAG     = $(DATE)-$(COMMIT_HASH)
 DOCKER_USER    = trinkle23897
@@ -42,7 +45,7 @@ go-install:
 	command -v go || (sudo apt-get install -y golang-1.18 && sudo ln -sf /usr/lib/go-1.18/bin/go /usr/bin/go)
 
 bazel-install: go-install
-	command -v bazel || (go install github.com/bazelbuild/bazelisk@latest && ln -sf $(HOME)/go/bin/bazelisk $(HOME)/go/bin/bazel)
+	command -v bazelisk || go install github.com/bazelbuild/bazelisk@latest
 
 buildifier-install: go-install
 	command -v buildifier || go install github.com/bazelbuild/buildtools/buildifier@latest
@@ -105,28 +108,28 @@ bazel-pip-requirement-release:
 	cd third_party/pip_requirements && (cmp requirements.txt requirements-release.txt || ln -sf requirements-release.txt requirements.txt)
 
 clang-tidy: clang-tidy-install bazel-pip-requirement-dev
-	bazel build $(BAZELOPT) //... --config=clang-tidy --config=test
+	$(BAZEL) build $(BAZELOPT) //... --config=clang-tidy --config=test
 
 bazel-debug: bazel-install bazel-pip-requirement-dev
-	bazel run $(BAZELOPT) //:setup --config=debug -- bdist_wheel
+	$(BAZEL) run $(BAZELOPT) //:setup --config=debug -- bdist_wheel
 	mkdir -p dist
 	cp bazel-bin/setup.runfiles/$(PROJECT_NAME)/dist/*.whl ./dist
 
 bazel-build: bazel-install bazel-pip-requirement-dev
-	bazel run $(BAZELOPT) //:setup --config=test -- bdist_wheel
+	$(BAZEL) run $(BAZELOPT) //:setup --config=test -- bdist_wheel
 	mkdir -p dist
 	cp bazel-bin/setup.runfiles/$(PROJECT_NAME)/dist/*.whl ./dist
 
 bazel-release: bazel-install bazel-pip-requirement-release
-	bazel run $(BAZELOPT) //:setup --config=release -- bdist_wheel
+	$(BAZEL) run $(BAZELOPT) //:setup --config=release -- bdist_wheel
 	mkdir -p dist
 	cp bazel-bin/setup.runfiles/$(PROJECT_NAME)/dist/*.whl ./dist
 
 bazel-test: bazel-install bazel-pip-requirement-dev
-	bazel test --test_output=all $(BAZELOPT) //... --config=test --spawn_strategy=local --color=yes
+	$(BAZEL) test --test_output=all $(BAZELOPT) //... --config=test --spawn_strategy=local --color=yes
 
 bazel-clean: bazel-install
-	bazel clean --expunge
+	$(BAZEL) clean --expunge
 
 # documentation
 

envpool/atari/atari_envpool_test.py

@@ -171,7 +171,10 @@ def test_xla_api(self) -> None:
       num_threads=2,
       thread_affinity_offset=0,
     )
-    handle, recv, send, step = env.xla()
+    try:
+      handle, recv, send, step = env.xla()
+    except RuntimeError as exc:
+      self.skipTest(str(exc))
     env.async_reset()
     handle, states = recv(handle)
     info = states[-1]
@@ -206,7 +209,10 @@ def test_xla_correctness(self) -> None:
       num_threads=2,
       thread_affinity_offset=0,
     )
-    handle, recv, send, step = env1.xla()
+    try:
+      handle, recv, send, step = env1.xla()
+    except RuntimeError as exc:
+      self.skipTest(str(exc))
     env1.async_reset()
     env2.async_reset()
 

envpool/core/spec.h

@@ -72,6 +72,24 @@ class Spec : public ShapeSpec {
       : ShapeSpec(sizeof(dtype), std::move(shape)), bounds(std::move(bounds)) {}
   Spec(const std::vector<int>& shape, const std::tuple<dtype, dtype>& bounds)
       : ShapeSpec(sizeof(dtype), shape), bounds(bounds) {}
+  Spec(std::vector<int>&& shape, std::initializer_list<dtype> bounds)
+      : ShapeSpec(sizeof(dtype), std::move(shape)) {
+    CHECK_EQ(bounds.size(), 2);
+    auto it = bounds.begin();
+    this->bounds = {it[0], it[1]};
+  }
+  Spec(const std::vector<int>& shape, std::initializer_list<dtype> bounds)
+      : ShapeSpec(sizeof(dtype), shape) {
+    CHECK_EQ(bounds.size(), 2);
+    auto it = bounds.begin();
+    this->bounds = {it[0], it[1]};
+  }
+  Spec(std::initializer_list<int> shape, std::initializer_list<dtype> bounds)
+      : ShapeSpec(sizeof(dtype), std::vector<int>(shape)) {
+    CHECK_EQ(bounds.size(), 2);
+    auto it = bounds.begin();
+    this->bounds = {it[0], it[1]};
+  }
 
   /* init with elementwise bounds */
   Spec(std::vector<int>&& shape,

envpool/minigrid/minigrid_align_test.py

@@ -23,7 +23,7 @@
 from absl.testing import absltest
 
 import envpool.minigrid.registration  # noqa: F401
-from envpool.registration import make_gym
+from envpool.registration import make_gymnasium
 
 
 class _MiniGridEnvPoolAlignTest(absltest.TestCase):
@@ -48,7 +48,7 @@ def run_align_check(
     **kwargs: Any,
   ) -> None:
     env0 = gym.make(task_id)
-    env1 = make_gym(task_id, num_envs=num_envs, seed=0, **kwargs)
+    env1 = make_gymnasium(task_id, num_envs=num_envs, seed=0, **kwargs)
     obs_space0 = cast(Any, env0.observation_space)
     self.check_spec(
       obs_space0["direction"], env1.observation_space["direction"]

envpool/mujoco/dmc/mujoco_env.h

@@ -20,6 +20,7 @@
 #include <mjxmacro.h>
 #include <mujoco.h>
 
+#include <array>
 #include <memory>
 #include <random>
 #include <string>

envpool/pip.bzl

@@ -22,7 +22,7 @@ def workspace():
     if "pip_requirements" not in native.existing_rules().keys():
         pip_install(
             name = "pip_requirements",
-            python_interpreter = "python3",
+            python_interpreter_target = "@python3_10_x86_64-unknown-linux-gnu//:bin/python3",
             # default timeout value is 600, change it if you failed.
             # timeout = 3600,
             quiet = False,

envpool/python/dm_envpool.py

@@ -54,10 +54,12 @@ def __new__(cls: Any, name: str, parents: Tuple, attrs: Dict) -> Any:
       parents = (
         base, DMEnvPoolMixin, EnvPoolMixin, XlaMixin, dm_env.Environment
       )
-    except ImportError:
+    except (ImportError, AttributeError):
 
       def _xla(self: Any) -> None:
-        raise RuntimeError("XLA is disabled. To enable XLA please install jax.")
+        raise RuntimeError(
+          "XLA is unavailable. To enable XLA please install a compatible jax."
+        )
 
       attrs["xla"] = _xla
       parents = (base, DMEnvPoolMixin, EnvPoolMixin, dm_env.Environment)

envpool/python/gym_envpool.py

@@ -54,10 +54,12 @@ def __new__(cls: Any, name: str, parents: Tuple, attrs: Dict) -> Any:
       from .lax import XlaMixin
 
       parents = (base, GymEnvPoolMixin, EnvPoolMixin, XlaMixin, gym.Env)
-    except ImportError:
+    except (ImportError, AttributeError):
 
       def _xla(self: Any) -> None:
-        raise RuntimeError("XLA is disabled. To enable XLA please install jax.")
+        raise RuntimeError(
+          "XLA is unavailable. To enable XLA please install a compatible jax."
+        )
 
       attrs["xla"] = _xla
       parents = (base, GymEnvPoolMixin, EnvPoolMixin, gym.Env)

envpool/python/gymnasium_envpool.py

@@ -58,10 +58,12 @@ def __new__(cls: Any, name: str, parents: Tuple, attrs: Dict) -> Any:
       parents = (
         base, GymnasiumEnvPoolMixin, EnvPoolMixin, XlaMixin, gymnasium.Env
       )
-    except ImportError:
+    except (ImportError, AttributeError):
 
       def _xla(self: Any) -> None:
-        raise RuntimeError("XLA is disabled. To enable XLA please install jax.")
+        raise RuntimeError(
+          "XLA is unavailable. To enable XLA please install a compatible jax."
+        )
 
       attrs["xla"] = _xla
       parents = (base, GymnasiumEnvPoolMixin, EnvPoolMixin, gymnasium.Env)

envpool/python/xla_template.py

@@ -54,6 +54,11 @@ def _make_xla_function(
   specs: Tuple[Tuple[Any, ...], Tuple[Any, ...]],
   capsules: Tuple[Any, Any],
 ) -> Callable:
+  if not hasattr(_xla, "backend_specific_translations"):
+    raise RuntimeError(
+      "XLA is unavailable because this JAX version removed the legacy "
+      "backend translation API used by envpool."
+    )
   in_specs, out_specs = specs
   in_specs = _normalize_specs(in_specs)
   out_specs = _normalize_specs(out_specs)

envpool/vizdoom/vizdoom_pretrain_test.py

@@ -14,6 +14,7 @@
 """Test Vizdoom env by well-trained RL agents."""
 
 import os
+import shutil
 from typing import Optional, Tuple
 
 import numpy as np
@@ -38,6 +39,12 @@ class _VizdoomPretrainTest(absltest.TestCase):
   def get_path(self, path: str) -> str:
     return os.path.join("envpool", "vizdoom", "maps", path)
 
+  def cleanup_runtime_dir(self) -> None:
+    if os.path.isdir("_vizdoom"):
+      shutil.rmtree("_vizdoom")
+    elif os.path.exists("_vizdoom"):
+      os.remove("_vizdoom")
+
   def eval_c51(
     self,
     task: str,
@@ -60,6 +67,7 @@ def eval_c51(
       kwargs.update(cfg_path=cfg_path)
     if reward_config is not None:
       kwargs.update(reward_config=reward_config)
+    self.cleanup_runtime_dir()
     env = make_gym(task_id, **kwargs)
 
     state_shape = env.observation_space.shape
@@ -102,6 +110,8 @@ def eval_c51(
 
     logging.info(f"Mean reward of {task}: {reward.mean()} ± {reward.std()}")
     logging.info(f"Mean length of {task}: {length.mean()} ± {length.std()}")
+    env.close()
+    self.cleanup_runtime_dir()
     return reward, length
 
   def test_d1(self) -> None:

envpool/workspace0.bzl

@@ -249,6 +249,7 @@ def workspace():
         strip_prefix = "SDL2-2.28.4",
         urls = [
             "https://www.libsdl.org/release/SDL2-2.28.4.tar.gz",
+            "https://github.com/libsdl-org/SDL/releases/download/release-2.28.4/SDL2-2.28.4.tar.gz",
             "https://ml.cs.tsinghua.edu.cn/~jiayi/envpool/libsdl/SDL2-2.28.4.tar.gz",
         ],
         build_file = "//third_party/sdl2:sdl2.BUILD",

envpool/workspace1.bzl

@@ -18,15 +18,24 @@ load("@com_github_nelhage_rules_boost//:boost/boost.bzl", "boost_deps")
 load("@com_justbuchanan_rules_qt//:qt_configure.bzl", "qt_configure")
 load("@pybind11_bazel//:python_configure.bzl", "python_configure")
 load("@rules_foreign_cc//foreign_cc:repositories.bzl", "rules_foreign_cc_dependencies")
+load("@rules_python//python:repositories.bzl", "python_register_toolchains")
 
 def workspace():
     """Configure pip requirements."""
+    python_register_toolchains(
+        name = "python3_10",
+        python_version = "3.10",
+        ignore_root_user_error = True,
+    )
+
     python_configure(
         name = "local_config_python",
-        python_version = "3",
+        python_interpreter_target = "@python3_10_x86_64-unknown-linux-gnu//:bin/python3",
     )
 
-    rules_foreign_cc_dependencies()
+    rules_foreign_cc_dependencies(
+        register_built_pkgconfig_toolchain = False,
+    )
 
     boost_deps()
 

third_party/ale/ale.BUILD

@@ -40,6 +40,10 @@ template_rule(
 
 cc_library(
     name = "ale_interface",
+    copts = [
+        "-include",
+        "cstdint",
+    ],
     srcs = glob(
         [
             "src/**/*.h",

third_party/pip_requirements/requirements-dev.txt

@@ -1,11 +1,11 @@
 setuptools
 wheel
-numpy
+numpy<2
 dm-env
 gym>=0.26
 gymnasium>=0.26,!=0.27.0
 optree>=0.6.0
-jax[cpu]
+jax[cpu]<0.5
 absl-py
 packaging
 tqdm

third_party/pip_requirements/requirements-release.txt

@@ -1,9 +1,9 @@
 setuptools
 wheel
-numpy
+numpy<2
 dm-env
 gym>=0.26
 gymnasium>=0.26,!=0.27.0
 optree>=0.6.0
-jax[cpu]
+jax[cpu]<0.5
 packaging