Pixel observations for dmc

.bazelrc

@@ -1,9 +1,10 @@
+build --action_env=DISPLAY=''
 build --action_env=BAZEL_LINKLIBS=-l%:libstdc++.a:-lm
 build --action_env=BAZEL_LINKOPTS=-static-libgcc
 build --action_env=CUDA_DIR=/usr/local/cuda
 build --incompatible_strict_action_env --cxxopt=-std=c++17 --host_cxxopt=-std=c++17 --client_env=BAZEL_CXXOPTS=-std=c++17
 build:debug --cxxopt=-DENVPOOL_TEST --compilation_mode=dbg -s
-build:test --cxxopt=-DENVPOOL_TEST --copt=-g0 --copt=-O3 --copt=-DNDEBUG --copt=-msse --copt=-msse2 --copt=-mmmx
+build:test --cxxopt=-DENVPOOL_TEST --copt=-g0 --copt=-O3 --copt=-DNDEBUG --copt=-msse --copt=-msse2 --copt=-mmmx --copt=-DMESA_EGL_NO_X11_HEADERS --copt=-DEGL_NO_X11 --copt=-DMJ_OSMESA
 build:release --copt=-g0 --copt=-O3 --copt=-DNDEBUG --copt=-msse --copt=-msse2 --copt=-mmmx
 
 build:clang-tidy --aspects @bazel_clang_tidy//clang_tidy:clang_tidy.bzl%clang_tidy_aspect

envpool/make_test.py

@@ -126,7 +126,7 @@ def test_make_mujoco_gym(self) -> None:
         "Walker2d-v4",
       ]
     )
-
+  # make test with GL loading error
   def test_make_mujoco_dmc(self) -> None:
     self.check_step(
       [

envpool/mujoco/BUILD

@@ -86,11 +86,15 @@ cc_library(
         "dmc/walker.h",
     ],
     data = [":gen_mujoco_dmc_xml"],
+    linkopts = [
+        "-lOSMesa",
+    ],
     deps = [
         "//envpool/core:async_envpool",
         "@mujoco//:mujoco_lib",
         "@pugixml",
     ],
+    alwayslink = 1,
 )
 
 pybind_extension(

envpool/mujoco/dmc/mujoco_dmc_suite_ext_render_align_test.py

@@ -0,0 +1,103 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for Mujoco dm_control suite align check."""
+
+from typing import Any, List, no_type_check
+
+import dm_env
+import numpy as np
+from absl import logging
+from absl.testing import absltest
+from dm_control import suite
+
+from envpool.mujoco.dmc import DmcHumanoidCMUDMEnvPool, DmcHumanoidCMUEnvSpec
+
+
+class _MujocoDmcSuiteExtAlignTest(absltest.TestCase):
+
+  @no_type_check
+  def run_space_check(self, env0: dm_env.Environment, env1: Any) -> None:
+    """Check observation_spec() and action_spec()."""
+    obs0, obs1 = env0.observation_spec(), env1.observation_spec()
+    for k in obs0:
+      self.assertTrue(hasattr(obs1, k))
+      np.testing.assert_allclose(obs0[k].shape, getattr(obs1, k).shape)
+    act0, act1 = env0.action_spec(), env1.action_spec()
+    np.testing.assert_allclose(act0.shape, act1.shape)
+    np.testing.assert_allclose(act0.minimum, act1.minimum)
+    np.testing.assert_allclose(act0.maximum, act1.maximum)
+
+  @no_type_check
+  def reset_state(
+    self, env: dm_env.Environment, ts: dm_env.TimeStep, domain: str, task: str
+  ) -> None:
+    # manually reset, mimic initialize_episode
+    with env.physics.reset_context():
+      env.physics.data.qpos = ts.observation.qpos0[0]
+      if domain in ["humanoid_CMU"]:
+        env.physics.after_reset()
+
+  def sample_action(self, action_spec: dm_env.specs.Array) -> np.ndarray:
+    return np.random.uniform(
+      low=action_spec.minimum,
+      high=action_spec.maximum,
+      size=action_spec.shape,
+    )
+
+  def run_align_check(
+    self, env0: dm_env.Environment, env1: Any, domain: str, task: str
+  ) -> None:
+    logging.info(f"align check for {domain} {task}")
+    obs_spec, action_spec = env0.observation_spec(), env0.action_spec()
+    for i in range(3):
+      np.random.seed(i)
+      env0.reset()
+      a = self.sample_action(action_spec)
+      ts = env1.reset(np.array([0]))
+      self.reset_state(env0, ts, domain, task)
+      logging.info(f"reset qpos {ts.observation.qpos0[0]}")
+      cnt = 0
+      done = False
+      while not done:
+        cnt += 1
+        a = self.sample_action(action_spec)
+        # logging.info(f"{cnt} {a}")
+        ts0 = env0.step(a)
+        ts1 = env1.step(np.array([a]), np.array([0]))
+        done = ts0.step_type == dm_env.StepType.LAST
+        o0, o1 = ts0.observation, ts1.observation
+        for k in obs_spec:
+          np.testing.assert_allclose(o0[k], getattr(o1, k)[0])
+        np.testing.assert_allclose(ts0.step_type, ts1.step_type[0])
+        np.testing.assert_allclose(ts0.reward, ts1.reward[0], atol=1e-8)
+        np.testing.assert_allclose(ts0.discount, ts1.discount[0])
+
+  def run_align_check_entry(
+    self, domain: str, tasks: List[str], spec_cls: Any, envpool_cls: Any
+  ) -> None:
+    for task in tasks:
+      env0 = suite.load(domain, task)
+      env1 = envpool_cls(spec_cls(spec_cls.gen_config(task_name=task)))
+      self.run_space_check(env0, env1)
+      self.run_align_check(env0, env1, domain, task)
+
+  def test_humanoid_CMU(self) -> None:
+    self.run_align_check_entry(
+      "humanoid_CMU", ["stand", "run"], DmcHumanoidCMUEnvSpec,
+      DmcHumanoidCMUDMEnvPool
+    )
+
+
+if __name__ == "__main__":
+  absltest.main()

envpool/mujoco/dmc/mujoco_dmc_suite_ext_render_deterministic_test.py

@@ -0,0 +1,81 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for Mujoco dm_control deterministic check."""
+
+from typing import Any, List, Optional
+
+import dm_env
+import numpy as np
+from absl.testing import absltest
+
+from envpool.mujoco.dmc import DmcHumanoidCMUDMEnvPool, DmcHumanoidCMUEnvSpec
+
+
+class _MujocoDmcSuiteExtDeterministicTest(absltest.TestCase):
+
+  def check(
+    self,
+    spec_cls: Any,
+    envpool_cls: Any,
+    task: str,
+    obs_keys: List[str],
+    blacklist: Optional[List[str]] = None,
+    num_envs: int = 4,
+  ) -> None:
+    np.random.seed(0)
+    env0 = envpool_cls(
+      spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=0, task_name=task))
+    )
+    env1 = envpool_cls(
+      spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=0, task_name=task))
+    )
+    env2 = envpool_cls(
+      spec_cls(spec_cls.gen_config(num_envs=num_envs, seed=1, task_name=task))
+    )
+    act_spec = env0.action_spec()
+    for t in range(3000):
+      action = np.array(
+        [
+          np.random.uniform(
+            low=act_spec.minimum, high=act_spec.maximum, size=act_spec.shape
+          ) for _ in range(num_envs)
+        ]
+      )
+      ts0 = env0.step(action)
+      obs0 = ts0.observation
+      obs1 = env1.step(action).observation
+      obs2 = env2.step(action).observation
+      for k in obs_keys:
+        o0 = getattr(obs0, k)
+        o1 = getattr(obs1, k)
+        o2 = getattr(obs2, k)
+        np.testing.assert_allclose(o0, o1)
+        if blacklist and k in blacklist:
+          continue
+        if np.abs(o0).sum() > 0 and ts0.step_type[0] != dm_env.StepType.FIRST:
+          self.assertFalse(np.allclose(o0, o2), (t, k, o0, o2))
+
+  def test_humanoid_CMU(self) -> None:
+    obs_keys = [
+      "joint_angles", "head_height", "extremities", "torso_vertical",
+      "com_velocity", "velocity"
+    ]
+    for task in ["stand", "run"]:
+      self.check(
+        DmcHumanoidCMUEnvSpec, DmcHumanoidCMUDMEnvPool, task, obs_keys
+      )
+
+
+if __name__ == "__main__":
+  absltest.main()