RC_WheelLeg/05_software/train/rc_mjlab/mjlab/tests/conftest.py

"""Shared test fixtures and utilities."""

import os
from pathlib import Path

import mujoco
import pytest
import torch
import warp as wp

from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
from mjlab.scene import Scene, SceneCfg
from mjlab.sim.sim import Simulation, SimulationCfg


@pytest.fixture(scope="session", autouse=True)
def configure_test_environment():
  """Configure test environment settings automatically for all tests."""
  wp.config.quiet = True


def get_test_device() -> str:
  """Get device for testing, preferring CUDA if available.

  Can be overridden with FORCE_CPU=1 environment variable to test
  CPU-only behavior on GPU machines.
  """
  if os.environ.get("FORCE_CPU") == "1":
    return "cpu"
  return "cuda" if torch.cuda.is_available() else "cpu"


@pytest.fixture
def fixtures_dir() -> Path:
  """Path to test fixtures directory."""
  return Path(__file__).parent / "fixtures"


def load_fixture_xml(fixture_name: str) -> str:
  """Load XML content from fixture file.

  Args:
    fixture_name: Name of the fixture file (without .xml extension) or full path.

  Returns:
    XML content as string.
  """
  fixtures_path = Path(__file__).parent / "fixtures"
  if not fixture_name.endswith(".xml"):
    fixture_name = f"{fixture_name}.xml"
  fixture_file = fixtures_path / fixture_name
  return fixture_file.read_text()


def create_entity_with_actuator(xml_string: str, actuator_cfg):
  """Create entity with actuator from XML string.

  Args:
    xml_string: MuJoCo XML model string.
    actuator_cfg: Actuator configuration.

  Returns:
    Entity instance.
  """
  cfg = EntityCfg(
    spec_fn=lambda: mujoco.MjSpec.from_string(xml_string),
    articulation=EntityArticulationInfoCfg(actuators=(actuator_cfg,)),
  )
  return Entity(cfg)


def create_entity_from_fixture(fixture_name: str, actuator_cfg=None):
  """Create entity from fixture file.

  Args:
    fixture_name: Name of the fixture file (without .xml extension).
    actuator_cfg: Optional actuator configuration.

  Returns:
    Entity instance.
  """
  xml_string = load_fixture_xml(fixture_name)
  if actuator_cfg is not None:
    return create_entity_with_actuator(xml_string, actuator_cfg)
  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml_string))
  return Entity(cfg)


def initialize_entity(entity: Entity, device: str, num_envs: int = 1):
  """Initialize entity with simulation.

  Args:
    entity: Entity to initialize.
    device: Device to use ("cpu" or "cuda").
    num_envs: Number of environments.

  Returns:
    Tuple of (entity, simulation).
  """
  model = entity.compile()
  sim_cfg = SimulationCfg()
  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
  entity.initialize(model, sim.model, sim.data, device)
  return entity, sim


def make_scene_and_sim(
  device: str,
  xml: str | dict[str, str],
  sensors: tuple,
  num_envs: int = 1,
  sim_cfg: SimulationCfg | None = None,
) -> tuple[Scene, Simulation]:
  """Create a scene and simulation from inline XML with sensors wired up.

  ``xml`` may be a single XML string (registered as the ``robot`` entity) or a
  mapping of entity name to XML string for multi-entity scenes.
  """
  xml_by_entity = {"robot": xml} if isinstance(xml, str) else xml
  entities = {
    name: EntityCfg(spec_fn=lambda s=s: mujoco.MjSpec.from_string(s))
    for name, s in xml_by_entity.items()
  }
  scene_cfg = SceneCfg(
    num_envs=num_envs,
    env_spacing=5.0,
    entities=entities,
    sensors=sensors,
  )
  scene = Scene(scene_cfg, device)
  model = scene.compile()
  if sim_cfg is None:
    sim_cfg = SimulationCfg(njmax=20)
  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
  scene.initialize(sim.mj_model, sim.model, sim.data)
  if scene.sensor_context is not None:
    sim.set_sensor_context(scene.sensor_context)
  return scene, sim


# =============================================================================
# XML Fixture Loaders
# =============================================================================


@pytest.fixture
def fixed_base_box_xml() -> str:
  """Load fixed base box XML fixture."""
  return load_fixture_xml("fixed_base_box")


@pytest.fixture
def floating_base_box_xml() -> str:
  """Load floating base box XML fixture."""
  return load_fixture_xml("floating_base_box")


@pytest.fixture
def fixed_base_articulated_xml() -> str:
  """Load fixed base articulated robot XML fixture."""
  return load_fixture_xml("fixed_base_articulated")


@pytest.fixture
def floating_base_articulated_xml() -> str:
  """Load floating base articulated robot XML fixture."""
  return load_fixture_xml("floating_base_articulated")


@pytest.fixture
def biped_xml() -> str:
  """Load biped robot XML fixture with ground plane."""
  return load_fixture_xml("biped")


@pytest.fixture
def robot_with_floor_xml() -> str:
  """XML for a floating body above a ground plane."""
  return """
    <mujoco>
      <worldbody>
        <geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
        <body name="base" pos="0 0 2">
          <freejoint name="free_joint"/>
          <geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
          <site name="base_site" pos="0 0 -0.1"/>
        </body>
      </worldbody>
    </mujoco>
  """


@pytest.fixture
def falling_box_xml() -> str:
  """XML for a box that can fall onto a ground plane."""
  return """
    <mujoco>
      <worldbody>
        <body name="ground" pos="0 0 0">
          <geom name="ground_geom" type="plane" size="5 5 0.1" rgba="0.5 0.5 0.5 1"/>
        </body>
        <body name="box" pos="0 0 0.5">
          <freejoint name="box_joint"/>
          <geom name="box_geom" type="box" size="0.1 0.1 0.1" rgba="0.8 0.3 0.3 1"
            mass="1.0"/>
        </body>
      </worldbody>
    </mujoco>
  """