Pixel observations for dmc

.bazelrc

@@ -3,7 +3,7 @@ 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/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()