Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

tests/test_go1_constants.py

@@ -0,0 +1,114 @@
+"""Tests for go1_constants.py."""
+
+import re
+
+import mujoco
+import numpy as np
+import pytest
+
+from mjlab.asset_zoo.robots.unitree_go1 import go1_constants
+from mjlab.entity import Entity
+from mjlab.utils.string import resolve_expr
+
+
+@pytest.fixture(scope="module")
+def go1_entity() -> Entity:
+  return Entity(go1_constants.get_go1_robot_cfg())
+
+
+@pytest.fixture(scope="module")
+def go1_model(go1_entity: Entity) -> mujoco.MjModel:
+  return go1_entity.spec.compile()
+
+
+# fmt: off
+@pytest.mark.parametrize(
+  "actuator_config,stiffness,damping",
+  [
+    (go1_constants.GO1_HIP_ACTUATOR_CFG, go1_constants.STIFFNESS_HIP, go1_constants.DAMPING_HIP),
+    (go1_constants.GO1_KNEE_ACTUATOR_CFG, go1_constants.STIFFNESS_KNEE, go1_constants.DAMPING_KNEE),
+  ],
+)
+# fmt: on
+def test_actuator_parameters(go1_model, actuator_config, stiffness, damping):
+  for i in range(go1_model.nu):
+    actuator = go1_model.actuator(i)
+    actuator_name = actuator.name
+    matches = any(
+      re.match(pattern, actuator_name) for pattern in actuator_config.target_names_expr
+    )
+    if matches:
+      assert actuator.gainprm[0] == stiffness
+      assert actuator.biasprm[1] == -stiffness
+      assert actuator.biasprm[2] == -damping
+      assert actuator.forcerange[0] == -actuator_config.effort_limit
+      assert actuator.forcerange[1] == actuator_config.effort_limit
+
+
+def test_keyframe_joint_positions(go1_entity, go1_model) -> None:
+  """Test that keyframe joint positions match the configuration."""
+  key = go1_model.key("init_state")
+  expected_joint_pos = go1_constants.INIT_STATE.joint_pos
+  assert expected_joint_pos is not None
+  expected_values = resolve_expr(expected_joint_pos, go1_entity.joint_names, 0.0)
+  for joint_name, expected_value in zip(
+    go1_entity.joint_names, expected_values, strict=True
+  ):
+    joint = go1_model.joint(joint_name)
+    qpos_idx = joint.qposadr[0]
+    actual_value = key.qpos[qpos_idx]
+    np.testing.assert_allclose(
+      actual_value,
+      expected_value,
+      rtol=1e-5,
+      err_msg=f"Joint {joint_name} position mismatch: "
+      f"expected {expected_value}, got {actual_value}",
+    )
+
+
+def test_foot_collision_geoms(go1_model) -> None:
+  """Foot collision geoms should have specific properties."""
+  foot_pattern = r"^[FR][LR]_foot_collision$"
+  for i in range(go1_model.ngeom):
+    geom = go1_model.geom(i)
+    if re.match(foot_pattern, geom.name):
+      assert geom.condim == 6
+      assert geom.priority == 1
+      assert geom.friction[0] == 1.0
+
+
+def test_all_collision_geoms_enabled(go1_model) -> None:
+  """Every *_collision geom, including numbered ones, must stay collidable.
+
+  The go1 XML numbers its secondary collision capsules (FR_thigh_collision1,
+  ...), which a pattern anchored at "_collision" would silently drop, letting
+  disable_other_geoms zero their contype. Guard against that regression with
+  an oracle independent of the config's own regex resolution.
+  """
+  collision_geoms = [
+    go1_model.geom(i)
+    for i in range(go1_model.ngeom)
+    if "_collision" in go1_model.geom(i).name
+  ]
+  assert len(collision_geoms) == 30
+  for geom in collision_geoms:
+    assert geom.contype == 1, f"{geom.name} lost its contype"
+    assert geom.conaffinity == 1, f"{geom.name} lost its conaffinity"
+
+
+def test_collision_geom_count(go1_model) -> None:
+  """Go1 should have 4 foot collision geoms."""
+  foot_pattern = r"^[FR][LR]_foot_collision$"
+  foot_geoms = [
+    g.name for g in [go1_model.geom(i) for i in range(go1_model.ngeom)]
+    if re.match(foot_pattern, g.name)
+  ]
+  assert len(foot_geoms) == 4
+
+
+def test_go1_entity_creation(go1_entity) -> None:
+  """Test basic Go1 entity properties."""
+  assert go1_entity.num_actuators == 12
+  assert go1_entity.num_joints == 12
+  assert go1_entity.is_actuated
+  assert not go1_entity.is_fixed_base