Add humanoid bench (#34)

* Add humanoid bench

* Update versions

* It works off the shelf

* Exclude all hooks

* Pre-commit fix

.pre-commit-config.yaml

@@ -27,3 +27,4 @@ repos:
         additional_dependencies:
           - "numpy"
           - "types-tabulate"
+        exclude: "safe_opax/benchmark_suites/humanoid_bench/.*"

pyproject.toml

@@ -8,7 +8,9 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.10,<3.12"
 equinox = "^0.11.1"
-jaxlib = {url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.23+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl", extras = ["cuda11_pip"]}
+jaxlib = { url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.23+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl", extras = [
+    "cuda11_pip",
+] }
 optax = "^0.1.7"
 hydra-core = "^1.3.2"
 jax = "^0.4.23"
@@ -22,15 +24,23 @@ dm-env = "^1.6"
 distrax = "^0.1.5"
 pillow = "^10.2.0"
 moviepy = "^1.0.3"
-safe-adaptation-gym = {git = "https://[email protected]/lasgroup/safe-adaptation-gym"}
-jmp = {git = "https://github.com/deepmind/jmp"}
+safe-adaptation-gym = { git = "https://[email protected]/lasgroup/safe-adaptation-gym" }
+jmp = { git = "https://github.com/deepmind/jmp" }
 tensorboard = "^2.16.2"
+torch = { version = "^2.4.0+cpu", source = "pytorch_cpu" }
+torchvision = { version = "^0.19.0+cpu", source = "pytorch_cpu" }
 
 
 [[tool.poetry.source]]
 name = "PyPI"
 priority = "primary"
 
+
+[[tool.poetry.source]]
+name = "pytorch_cpu"
+url = "https://download.pytorch.org/whl/cpu"
+priority = "explicit"
+
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.7.0"
 pre-commit = "^3.5.0"
@@ -74,6 +84,7 @@ exclude = [
     "dist",
     "node_modules",
     "venv",
+    "safe_opax/benchmark_suites/humanoid_bench",
 ]
 
 [tool.ruff.format]
@@ -90,3 +101,8 @@ disallow_untyped_defs = false
 disallow_subclassing_any = false
 disallow_untyped_calls = false
 strict_optional = true
+exclude = "safe_opax/benchmark_suites/humanoid_bench"
+
+[[tool.mypy.overrides]]
+module = ["safe_opax.benchmark_suites.humanoid_bench.*"]
+ignore_errors = true

safe_opax/benchmark_suites/__init__.py

@@ -15,6 +15,10 @@ def make(cfg: DictConfig) -> EnvironmentFactory:
     elif domain_name == "safe_adaptation_gym":
         from safe_opax.benchmark_suites.safe_adaptation_gym import make
 
+        make_env = make(cfg)
+    elif domain_name == "humanoid_bench":
+        from safe_opax.benchmark_suites.humanoid_bench import make
+
         make_env = make(cfg)
     else:
         raise NotImplementedError(f"Environment {domain_name} not implemented")

safe_opax/benchmark_suites/humanoid_bench/README.md

@@ -0,0 +1,4 @@
+This code is based on https://github.com/carlosferrazza/humanoid-bench
+
+
+See: https://humanoid-bench.github.io/

safe_opax/benchmark_suites/humanoid_bench/__init__.py

@@ -0,0 +1,28 @@
+from omegaconf import DictConfig
+
+from safe_opax.benchmark_suites.utils import get_domain_and_task
+from safe_opax.rl.types import EnvironmentFactory
+from safe_opax.rl.wrappers import ImageObservation
+
+
+
+
+def make(cfg: DictConfig) -> EnvironmentFactory:
+    def make_env():
+        from .env import HumanoidEnv
+
+        _, task_cfg = get_domain_and_task(cfg)
+        env_name = "h1hand-pole-v0"
+        env = HumanoidEnv(robot="h1hand", control="pos", task="pole")
+        if task_cfg.image_observation.enabled:
+            env = ImageObservation(
+                env,
+                task_cfg.image_observation.image_size,
+                task_cfg.image_observation.image_format
+            )
+        else:
+            from gymnasium.wrappers.flatten_observation import FlattenObservation
+
+            env = FlattenObservation(env)
+        return env
+    return make_env

safe_opax/benchmark_suites/humanoid_bench/assets/common/floor.xml

@@ -0,0 +1,6 @@
+<mujocoinclude>
+    <worldbody>
+        <light pos="0 0 10" dir="0 0 -1" directional="true" ambient="0.1 0.1 0.1" diffuse="0.5 0.5 0.5" specular="0.3 0.3 0.3"/>
+        <geom name="floor" size="20 20 1" type="plane" conaffinity="1" material="MatPlane"/>
+    </worldbody>
+</mujocoinclude>

safe_opax/benchmark_suites/humanoid_bench/assets/common/visual.xml

@@ -0,0 +1,13 @@
+<mujocoinclude>
+    <visual>
+        <map force="0.1" fogend="10" shadowclip="5" />
+        <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
+        <!--<quality shadowsize="16384" offsamples="16"/>-->
+        <quality shadowsize="16384"/>
+    </visual>
+    <asset>
+        <texture type="skybox" builtin="gradient" rgb1=".9 .9 .9" rgb2=".9 .9 .9" width="100" height="100"/>
+        <texture builtin="checker" height="100" width="100" name="texplane" rgb1="0.9 0.9 0.9" rgb2="0.9 0.9 0.9" type="2d" mark="edge" markrgb="0.7 0.7 0.7"/>
+        <material name="MatPlane" reflectance="0" shininess="1" specular="0.5" texrepeat="20 20" texture="texplane"/>
+    </asset>
+</mujocoinclude>

safe_opax/benchmark_suites/humanoid_bench/assets/digit/base-0_collision_0.xml

@@ -0,0 +1,20 @@
+<mujoco model="base-0_collision">
+<geom class="collision" type="mesh" mesh="base-0_part_0" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_1" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_2" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_3" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_4" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_5" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_6" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_7" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_8" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_9" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_10" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_11" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_12" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_13" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_14" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_15" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_16" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+<geom class="collision" type="mesh" mesh="base-0_part_17" pos='0 0 0' euler='-0 0 -0' friction='0.7 0.01 0.005'  mass='15.028392'/>
+</mujoco>

safe_opax/benchmark_suites/humanoid_bench/assets/digit/base-0_collision_assets.xml

@@ -0,0 +1,20 @@
+<mujoco model="base-0_collision_assets">
+<mesh  file="base-0_split/base-0_part_0.obj"  name="base-0_part_0"/>
+<mesh  file="base-0_split/base-0_part_1.obj"  name="base-0_part_1"/>
+<mesh  file="base-0_split/base-0_part_2.obj"  name="base-0_part_2"/>
+<mesh  file="base-0_split/base-0_part_3.obj"  name="base-0_part_3"/>
+<mesh  file="base-0_split/base-0_part_4.obj"  name="base-0_part_4"/>
+<mesh  file="base-0_split/base-0_part_5.obj"  name="base-0_part_5"/>
+<mesh  file="base-0_split/base-0_part_6.obj"  name="base-0_part_6"/>
+<mesh  file="base-0_split/base-0_part_7.obj"  name="base-0_part_7"/>
+<mesh  file="base-0_split/base-0_part_8.obj"  name="base-0_part_8"/>
+<mesh  file="base-0_split/base-0_part_9.obj"  name="base-0_part_9"/>
+<mesh  file="base-0_split/base-0_part_10.obj"  name="base-0_part_10"/>
+<mesh  file="base-0_split/base-0_part_11.obj"  name="base-0_part_11"/>
+<mesh  file="base-0_split/base-0_part_12.obj"  name="base-0_part_12"/>
+<mesh  file="base-0_split/base-0_part_13.obj"  name="base-0_part_13"/>
+<mesh  file="base-0_split/base-0_part_14.obj"  name="base-0_part_14"/>
+<mesh  file="base-0_split/base-0_part_15.obj"  name="base-0_part_15"/>
+<mesh  file="base-0_split/base-0_part_16.obj"  name="base-0_part_16"/>
+<mesh  file="base-0_split/base-0_part_17.obj"  name="base-0_part_17"/>
+</mujoco>

safe_opax/benchmark_suites/humanoid_bench/assets/digit/base-1_collision_0.xml

@@ -0,0 +1,7 @@
+<mujoco model="base-1_collision">
+<geom class="collision" type="mesh" mesh="base-1_part_0" pos='-0.001 0.091 0' euler='0 -1.5708 0' friction='0.7 0.01 0.005'/>
+<geom class="collision" type="mesh" mesh="base-1_part_1" pos='-0.001 0.091 0' euler='0 -1.5708 0' friction='0.7 0.01 0.005'/>
+<geom class="collision" type="mesh" mesh="base-1_part_2" pos='-0.001 0.091 0' euler='0 -1.5708 0' friction='0.7 0.01 0.005'/>
+<geom class="collision" type="mesh" mesh="base-1_part_3" pos='-0.001 0.091 0' euler='0 -1.5708 0' friction='0.7 0.01 0.005'/>
+<geom class="collision" type="mesh" mesh="base-1_part_4" pos='-0.001 0.091 0' euler='0 -1.5708 0' friction='0.7 0.01 0.005'/>
+</mujoco>

safe_opax/benchmark_suites/humanoid_bench/assets/digit/base-1_collision_assets.xml

@@ -0,0 +1,7 @@
+<mujoco model="base-1_collision_assets">
+<mesh  file="base-1_split/base-1_part_0.obj"  name="base-1_part_0"/>
+<mesh  file="base-1_split/base-1_part_1.obj"  name="base-1_part_1"/>
+<mesh  file="base-1_split/base-1_part_2.obj"  name="base-1_part_2"/>
+<mesh  file="base-1_split/base-1_part_3.obj"  name="base-1_part_3"/>
+<mesh  file="base-1_split/base-1_part_4.obj"  name="base-1_part_4"/>
+</mujoco>