Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

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

tests/test_xml_utils.py

@@ -0,0 +1,116 @@
+"""Tests for mjlab.utils.xml."""
+
+import zipfile
+
+import mujoco
+import numpy as np
+
+from mjlab.utils.xml import fix_spec_xml, strip_buffer_textures
+
+
+def test_strip_buffer_textures():
+  """Stripping buffer textures removes them and enables XML roundtrip."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="terrain")
+
+  tex = spec.add_texture(
+    name="hf_tex",
+    type=mujoco.mjtTexture.mjTEXTURE_2D,
+    width=2,
+    height=2,
+  )
+  tex.data = (np.ones((2, 2, 3), dtype=np.uint8) * 128).tobytes()
+
+  mat = spec.add_material(name="hf_mat")
+  mat.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "hf_tex"
+
+  hf = spec.add_hfield(name="hf", nrow=4, ncol=4, size=[1, 1, 0.1, 0.01])
+  hf.userdata = np.zeros(16).tolist()
+  body.add_geom(
+    type=mujoco.mjtGeom.mjGEOM_HFIELD,
+    hfieldname="hf",
+    material="hf_mat",
+  )
+
+  strip_buffer_textures(spec)
+
+  assert len(spec.textures) == 0
+  assert len(spec.materials) == 0
+  assert spec.geoms[0].material == ""
+
+  # XML roundtrip should succeed now.
+  model = mujoco.MjModel.from_xml_string(spec.to_xml())
+  assert model.nhfield == 1
+
+
+def test_fix_spec_xml():
+  """Removes empty defaults and optionally injects meshdir."""
+  xml = (
+    "<mujoco>\n"
+    "  <compiler/>\n"
+    "  <default/>\n"
+    '  <default class="foo/bar"/>\n'
+    "  <worldbody/>\n"
+    "</mujoco>\n"
+  )
+  result = fix_spec_xml(xml, meshdir="assets")
+
+  assert "<default/>" not in result
+  assert 'class="foo/bar"' not in result
+  assert 'meshdir="assets"' in result
+  assert "meshdir" not in fix_spec_xml(xml)
+
+
+def test_rough_terrain_write_xml_roundtrip(tmp_path):
+  """Export terrain with hfield via Entity.write_xml, reload as MjModel."""
+  from mjlab.terrains.heightfield_terrains import (
+    HfRandomUniformTerrainCfg,
+  )
+  from mjlab.terrains.terrain_entity import TerrainEntity, TerrainEntityCfg
+  from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
+
+  terrain = TerrainEntity(
+    TerrainEntityCfg(
+      terrain_type="generator",
+      terrain_generator=TerrainGeneratorCfg(
+        size=(2.0, 2.0),
+        num_rows=1,
+        num_cols=1,
+        sub_terrains={
+          "rough": HfRandomUniformTerrainCfg(
+            proportion=1.0,
+            noise_range=(-0.05, 0.05),
+            noise_step=0.005,
+          ),
+        },
+      ),
+    ),
+    device="cpu",
+  )
+  xml_path = tmp_path / "terrain.xml"
+  terrain.write_xml(xml_path)
+
+  model = mujoco.MjModel.from_xml_path(str(xml_path))
+  assert model.nhfield == 1
+
+
+def test_scene_write_zip(tmp_path):
+  """Scene.write produces a directory; with zip=True, a zip archive."""
+  from mjlab.scene.scene import Scene, SceneCfg
+
+  scene = Scene(SceneCfg(), device="cpu")
+
+  # Directory export.
+  dir_out = tmp_path / "scene_dir"
+  scene.write(dir_out)
+  assert (dir_out / "scene.xml").exists()
+  mujoco.MjModel.from_xml_path(str(dir_out / "scene.xml"))
+
+  # Zip export.
+  zip_out = tmp_path / "scene_pkg"
+  scene.write(zip_out, zip=True)
+  zip_path = zip_out.with_suffix(".zip")
+  assert zip_path.exists()
+  assert not zip_out.exists()
+  with zipfile.ZipFile(zip_path) as zf:
+    assert "scene.xml" in zf.namelist()