Trim MyoSuite oracle test cross-products

Author: Trinkle23897

BUILD

@@ -18,6 +18,14 @@ config_setting(
     ],
 )
 
+config_setting(
+    name = "linux_arm64",
+    constraint_values = [
+        "@platforms//cpu:arm64",
+        "@platforms//os:linux",
+    ],
+)
+
 config_setting(
     name = "windows",
     constraint_values = ["@platforms//os:windows"],

README.md

@@ -10,7 +10,6 @@
 
 - [x] [Atari games](https://envpool.readthedocs.io/en/latest/env/atari.html)
 - [x] [MuJoCo (Gymnasium)](https://envpool.readthedocs.io/en/latest/env/mujoco_gym.html)
-- [x] [MyoSuite](https://envpool.readthedocs.io/en/latest/env/myosuite.html)
 - [x] [Classic control RL envs](https://envpool.readthedocs.io/en/latest/env/classic_control.html): CartPole, MountainCar, Pendulum, Acrobot
 - [x] [Toy text RL envs](https://envpool.readthedocs.io/en/latest/env/toy_text.html): Catch, FrozenLake, Taxi, NChain, CliffWalking, Blackjack
 - [x] [ViZDoom single player](https://envpool.readthedocs.io/en/latest/env/vizdoom.html)
@@ -21,6 +20,7 @@
 - [x] [Minigrid](https://envpool.readthedocs.io/en/latest/env/minigrid.html)
 - [x] [Highway](https://envpool.readthedocs.io/en/latest/env/highway.html)
 - [x] [MetaWorld](https://envpool.readthedocs.io/en/latest/env/metaworld.html)
+- [x] [MyoSuite](https://envpool.readthedocs.io/en/latest/env/myosuite.html)
 
 Here are EnvPool's several highlights:
 

docs/_static/render_samples/myosuite_myobase_official_compare.png

@@ -13,8 +13,6 @@
 # limitations under the License.
 """Entry point for all envs' registration."""
 
-import importlib
-
 import envpool.atari.registration  # noqa: F401
 import envpool.box2d.registration  # noqa: F401
 import envpool.classic_control.registration  # noqa: F401
@@ -24,9 +22,8 @@
 import envpool.mujoco.dmc.registration  # noqa: F401
 import envpool.mujoco.gym.registration  # noqa: F401
 import envpool.mujoco.metaworld.registration  # noqa: F401
+import envpool.mujoco.myosuite.registration  # noqa: F401
 import envpool.mujoco.robotics.registration  # noqa: F401
 import envpool.procgen.registration  # noqa: F401
 import envpool.toy_text.registration  # noqa: F401
 import envpool.vizdoom.registration  # noqa: F401
-
-importlib.import_module("envpool.mujoco.myosuite.registration")

docs/_static/render_samples/myosuite_myochallenge_official_compare.png

@@ -341,11 +341,14 @@ cc_library(
         "//:linux_x86_64": [
             # CI fastbuild otherwise compiles the header-only MyoSuite runtime
             # at -O0 while the pinned official oracle wheel is release-built.
+            # This is x86_64-only because the official Linux x86_64 MuJoCo
+            # wheel enables platform SIMD; linux_arm64 only needs coverage
+            # disabled for stable long oracle rollouts.
             "-O3",
             "-fno-profile-arcs",
             "-fno-test-coverage",
         ],
-        "//:linux": [
+        "//:linux_arm64": [
             "-fno-profile-arcs",
             "-fno-test-coverage",
         ],
@@ -540,7 +543,7 @@ py_test(
     srcs = ["myosuite/myosuite_oracle_align_test.py"],
     data = [":myosuite_oracle_probe"],
     imports = ["../.."],
-    shard_count = 4,
+    shard_count = 8,
     deps = [
         ":myosuite",
         ":myosuite_registration",

docs/_static/render_samples/myosuite_myodm_official_compare.png

@@ -18,6 +18,7 @@
 import importlib
 import json
 import os
+import platform
 import subprocess
 import sys
 import tempfile
@@ -40,8 +41,12 @@
 importlib.import_module("envpool.mujoco.myosuite.registration")
 
 _ROLLOUT_STEPS = 128
-_ORACLE_SPACE_BATCH_SIZE = 32
-_ROLLOUT_BATCH_SIZE = 2
+_ORACLE_SPACE_BATCH_SIZE = 64
+_ROLLOUT_BATCH_SIZE = 4
+# Keep the expensive 128-step oracle traces to a diagonal sample across
+# orthogonal task modifiers. Full registry/space/render coverage still checks
+# every official ID; this set keeps muscle-condition/player-side combinations
+# from growing as a cartesian product.
 _ROLLOUT_TASK_IDS = frozenset({
     "MyoHandAirplaneFixed-v0",
     "MyoHandAirplaneFly-v0",
@@ -63,19 +68,16 @@
     "myoChallengeOslRunFixed-v0",
     "myoChallengeRelocateP1-v0",
     "myoChallengeSoccerP1-v0",
-    "myoChallengeSoccerP2-v0",
     "myoChallengeTableTennisP0-v0",
     "myoFatiChallengeBimanual-v0",
-    "myoFatiChallengeSoccerP1-v0",
-    "myoFatiChallengeSoccerP2-v0",
-    "myoSarcChallengeBimanual-v0",
-    "myoSarcChallengeSoccerP1-v0",
     "myoSarcChallengeSoccerP2-v0",
 })
 _BITWISE_ROLLOUT_TASK_IDS = frozenset({
     "myoFingerReachFixed-v0",
     "myoFingerPoseFixed-v0",
 })
+_LINUX_AARCH64_FINGER_ROLLOUT_RTOL = 1e-5
+_LINUX_AARCH64_FINGER_ROLLOUT_ATOL = 1e-7
 _EXPECTED_ORACLE_NUMPY2_BROKEN_IDS: frozenset[str] = frozenset()
 _SYNC_STATE_KEYS = (
     "qpos0",
@@ -145,6 +147,32 @@
 }
 
 
+def _assert_bitwise_rollout_obs(
+    actual: np.ndarray,
+    desired: np.ndarray,
+    *,
+    label: str,
+) -> None:
+    if sys.platform.startswith("linux") and platform.machine().lower() in {
+        "aarch64",
+        "arm64",
+    }:
+        # Linux aarch64 accumulates small float32 differences in these long
+        # MuJoCo finger traces after tens of steps. Keep the residual scoped to
+        # that platform and far below a semantically meaningful trajectory drift.
+        try:
+            np.testing.assert_allclose(
+                actual,
+                desired,
+                rtol=_LINUX_AARCH64_FINGER_ROLLOUT_RTOL,
+                atol=_LINUX_AARCH64_FINGER_ROLLOUT_ATOL,
+            )
+        except AssertionError as exc:
+            raise AssertionError(f"{label}\n{exc}") from exc
+        return
+    np.testing.assert_array_equal(actual, desired, err_msg=label)
+
+
 def _oracle_task_ids() -> tuple[str, ...]:
     return tuple(task["id"] for task in MYOSUITE_TASKS)
 
@@ -439,12 +467,12 @@ def test_oracle_rollout_surface(self) -> None:
                             envpool_obs.shape, (1, task["obs_dim"])
                         )
                         if task_id in _BITWISE_ROLLOUT_TASK_IDS:
-                            np.testing.assert_array_equal(
+                            _assert_bitwise_rollout_obs(
                                 envpool_obs[0].astype(np.float32),
                                 np.asarray(
                                     oracle_task["obs"][0], dtype=np.float32
                                 ),
-                                err_msg=f"{task_id} reset obs",
+                                label=f"{task_id} reset obs",
                             )
 
                         for step_id, action in enumerate(
@@ -467,10 +495,10 @@ def test_oracle_rollout_surface(self) -> None:
                                     oracle_task["rewards"][step_id],
                                     dtype=np.float32,
                                 )
-                                np.testing.assert_array_equal(
+                                _assert_bitwise_rollout_obs(
                                     envpool_step[0][0].astype(np.float32),
                                     oracle_obs,
-                                    err_msg=f"{task_id} step {step_id} obs",
+                                    label=f"{task_id} step {step_id} obs",
                                 )
                                 self.assertEqual(
                                     float(envpool_step[1][0]),

envpool/entry.py

@@ -44,6 +44,55 @@
 
 _TASK_IDS = tuple(str(task["id"]) for task in MYOSUITE_TASKS)
 _TASK_ID_SET = frozenset(_TASK_IDS)
+# Render traces are the expensive part of this suite. The full 398-ID surface is
+# covered by registry, space, reset/step, determinism, and generated docs. Keep
+# CI render checks to a diagonal set that catches wrong-camera/model/scene
+# regressions without retesting the full modifier cartesian product.
+_ORACLE_RENDER_REPRESENTATIVE_TASK_IDS = frozenset({
+    "MyoHandAirplaneFixed-v0",
+    "MyoHandAirplaneFly-v0",
+    "MyoHandCupPour-v0",
+    "MyoHandHammerUse-v0",
+    "MyoHandWatchRandom-v0",
+    "motorFingerPoseFixed-v0",
+    "myoArmReachFixed-v0",
+    "myoElbowPose1D6MExoFixed-v0",
+    "myoFingerReachFixed-v0",
+    "myoFingerPoseFixed-v0",
+    "myoHandReachFixed-v0",
+    "myoHandPoseFixed-v0",
+    "myoHandKeyTurnFixed-v0",
+    "myoHandObjHoldFixed-v0",
+    "myoHandPenTwirlFixed-v0",
+    "myoHandReorient8-v0",
+    "myoLegStandRandom-v0",
+    "myoLegWalk-v0",
+    "myoLegRoughTerrainWalk-v0",
+    "myoFatiArmReachFixed-v0",
+    "myoFatiHandReorient8-v0",
+    "myoFatiLegWalk-v0",
+    "myoSarcArmReachFixed-v0",
+    "myoSarcHandReorient8-v0",
+    "myoSarcLegWalk-v0",
+    "myoChallengeBaodingP1-v1",
+    "myoChallengeBimanual-v0",
+    "myoChallengeChaseTagP1-v0",
+    "myoChallengeDieReorientP1-v0",
+    "myoChallengeOslRunFixed-v0",
+    "myoChallengeRelocateP1-v0",
+    "myoChallengeSoccerP1-v0",
+    "myoChallengeTableTennisP0-v0",
+    "myoFatiChallengeBimanual-v0",
+    "myoSarcChallengeSoccerP2-v0",
+})
+_UNKNOWN_ORACLE_RENDER_REPRESENTATIVE_TASK_IDS = (
+    _ORACLE_RENDER_REPRESENTATIVE_TASK_IDS - _TASK_ID_SET
+)
+if _UNKNOWN_ORACLE_RENDER_REPRESENTATIVE_TASK_IDS:
+    raise ValueError(
+        "unknown MyoSuite oracle render representatives: "
+        f"{sorted(_UNKNOWN_ORACLE_RENDER_REPRESENTATIVE_TASK_IDS)}"
+    )
 
 
 def _render_task_allowlist_from_env() -> tuple[str, ...] | None:
@@ -74,16 +123,32 @@ def _filter_render_task_ids(task_ids: tuple[str, ...]) -> tuple[str, ...]:
     )
 
 
-_ORACLE_TRACE_TASK_IDS = tuple(
-    task_id
-    for task_id in _TASK_IDS
-    if task_id not in MYOSUITE_ORACLE_NUMPY2_BROKEN_IDS
-)
-_NATIVE_ONLY_RENDER_TASK_IDS = tuple(
-    task_id
-    for task_id in _TASK_IDS
-    if task_id in MYOSUITE_ORACLE_NUMPY2_BROKEN_IDS
-)
+def _native_render_task_ids() -> tuple[str, ...]:
+    if _RENDER_TASK_ALLOWLIST is not None:
+        return _filter_render_task_ids(_TASK_IDS)
+    return tuple(
+        task_id
+        for task_id in _TASK_IDS
+        if task_id in _ORACLE_RENDER_REPRESENTATIVE_TASK_IDS
+    )
+
+
+def _oracle_trace_task_ids() -> tuple[str, ...]:
+    if _RENDER_TASK_ALLOWLIST is None:
+        return tuple(
+            task_id
+            for task_id in _TASK_IDS
+            if task_id in _ORACLE_RENDER_REPRESENTATIVE_TASK_IDS
+            and task_id not in MYOSUITE_ORACLE_NUMPY2_BROKEN_IDS
+        )
+    return tuple(
+        task_id
+        for task_id in _filter_render_task_ids(_TASK_IDS)
+        if task_id not in MYOSUITE_ORACLE_NUMPY2_BROKEN_IDS
+    )
+
+
+_NATIVE_RENDER_TASK_IDS = _native_render_task_ids()
 _WIDTH = 64
 _HEIGHT = 48
 _ORACLE_RENDER_BATCH_SIZE = 8
@@ -199,10 +264,10 @@ def _task_batches(
 
 
 _SHARDED_ORACLE_TRACE_TASK_IDS = _render_shard_task_ids(
-    _filter_render_task_ids(_ORACLE_TRACE_TASK_IDS)
+    _oracle_trace_task_ids()
 )
-_SHARDED_NATIVE_ONLY_RENDER_TASK_IDS = _render_shard_task_ids(
-    _filter_render_task_ids(_NATIVE_ONLY_RENDER_TASK_IDS)
+_SHARDED_NATIVE_RENDER_TASK_IDS = _render_shard_task_ids(
+    _NATIVE_RENDER_TASK_IDS
 )
 
 
@@ -447,8 +512,8 @@ class MyoSuiteRenderTest(absltest.TestCase):
     """Validate native MyoSuite RGB rendering after reset and steps."""
 
     def test_reset_and_first_three_step_render(self) -> None:
-        """Render oracle-skip tasks through reset and the first three steps."""
-        for task_id in _SHARDED_NATIVE_ONLY_RENDER_TASK_IDS:
+        """Representative tasks render through reset and first three steps."""
+        for task_id in _SHARDED_NATIVE_RENDER_TASK_IDS:
             with self.subTest(task_id=task_id):
                 env = make_gymnasium(
                     task_id,

envpool/mujoco/BUILD

@@ -15,6 +15,7 @@
 """EnvPool pip requirements initialization, this is loaded in WORKSPACE."""
 
 load("@python_versions//:pip.bzl", "multi_pip_parse")
+load("//third_party/myosuite:oracle_workspace.bzl", "myosuite_oracle_pip_workspace")
 
 def workspace():
     """Configure pip requirements."""
@@ -41,25 +42,4 @@ def workspace():
             # extra_pip_args = ["--extra-index-url", "https://mirrors.aliyun.com/pypi/simple"],
         )
 
-    if "myosuite_oracle_requirements" not in native.existing_rules().keys():
-        # The official MyoSuite package is only a test/codegen oracle. Keep its
-        # dependency hub pinned to the same Python 3.12 interpreter for every
-        # EnvPool toolchain key so generic targets such as //:setup_py314 do not
-        # try to build oracle-only wheels for unsupported Python/platform pairs.
-        multi_pip_parse(
-            name = "myosuite_oracle_requirements",
-            default_version = "3.12",
-            python_interpreter_target = {
-                "3.11": "@python_versions_3_12_host//:python",
-                "3.12": "@python_versions_3_12_host//:python",
-                "3.13": "@python_versions_3_12_host//:python",
-                "3.14": "@python_versions_3_12_host//:python",
-            },
-            requirements_lock = {
-                "3.11": "@envpool//third_party/myosuite:oracle_requirements.txt",
-                "3.12": "@envpool//third_party/myosuite:oracle_requirements.txt",
-                "3.13": "@envpool//third_party/myosuite:oracle_requirements.txt",
-                "3.14": "@envpool//third_party/myosuite:oracle_requirements.txt",
-            },
-            quiet = False,
-        )
+    myosuite_oracle_pip_workspace()