From a942a1679e51929db66d3c3e54c52b47038fee2a Mon Sep 17 00:00:00 2001
From: Kallinteris Andreas
 <30759571+Kallinteris-Andreas@users.noreply.github.com>
Date: Wed, 6 Dec 2023 11:45:59 +0200
Subject: [PATCH]  add `MuJoCo.tests.test_model_sensors()`  (#816)

---
 tests/envs/mujoco/test_mujoco_v5.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tests/envs/mujoco/test_mujoco_v5.py b/tests/envs/mujoco/test_mujoco_v5.py
index b68d8b5ff..2a8df3c2b 100644
--- a/tests/envs/mujoco/test_mujoco_v5.py
+++ b/tests/envs/mujoco/test_mujoco_v5.py
@@ -627,6 +627,27 @@ def test_model_object_count(version: str):
     assert env.model.ntendon == 0
 
 
+# note: fails with `mujoco-mjx==3.0.1`
+@pytest.mark.parametrize("version", ["v5", "v4", "v3", "v2"])
+def test_model_sensors(version: str):
+    """Verify that all the sensors of the model are loaded."""
+    env = gym.make(f"Ant-{version}").unwrapped
+    assert env.data.cfrc_ext.shape == (14, 6)
+
+    env = gym.make(f"Humanoid-{version}").unwrapped
+    assert env.data.cinert.shape == (14, 10)
+    assert env.data.cvel.shape == (14, 6)
+    assert env.data.qfrc_actuator.shape == (23,)
+    assert env.data.cfrc_ext.shape == (14, 6)
+
+    if version != "v3":  # HumanoidStandup v3 does not exist
+        env = gym.make(f"HumanoidStandup-{version}").unwrapped
+        assert env.data.cinert.shape == (14, 10)
+        assert env.data.cvel.shape == (14, 6)
+        assert env.data.qfrc_actuator.shape == (23,)
+        assert env.data.cfrc_ext.shape == (14, 6)
+
+
 def test_dt():
     """Assert that env.dt gets assigned correctly."""
     env_a = gym.make("Ant-v5", include_cfrc_ext_in_observation=False).unwrapped