refactor: clean tools

scripts/analysis/dump_mjcf.py

@@ -1,186 +0,0 @@
-#!/usr/bin/env python3
-"""
-Dump MJCF XML for a brittle-star morphology using the project's Hydra configs.
-
-Usage examples:
-
-  # Use a named morphology config from configs/morphology (Hydra style)
-  uv run python scripts/analysis/dump_mjcf.py morphology=3_arms
-
-  # Use a morphology override YAML (same key as simulation.morphology_override)
-  uv run python scripts/analysis/dump_mjcf.py \
-    simulation.morphology_override=configs/morphology/3_arms.yaml
-
-Output path:
-    Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in
-    current directory.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-import sys
-import xml.etree.ElementTree as ET
-
-import hydra
-from omegaconf import DictConfig, OmegaConf
-
-from brittle_star_project.configs.main_config import BrittleStarConfig
-from brittle_star_project.configs.register_configs import register_configs
-from brittle_star_project.environment.env_config import MorphologyConfig
-from brittle_star_project.environment.factory import BrittleStarEnvFactory
-
-
-def try_serialize(obj):
-    """Try multiple common accessors to obtain an XML string from the morphology object."""
-    candidates = [
-        "to_xml_string",
-        "to_xml",
-        "to_string",
-        "to_mjcf",
-        "to_mjcf_string",
-        "get_mjcf",
-        "get_mjcf_str",
-        "get_mjcf_assets",
-        "export_to_xml_with_assets",
-        "get_xml",
-        "xml",
-        "mjcf",
-        "mjcf_model",
-        "mjcf_body",
-        "model",
-        "root",
-    ]
-
-    def normalize(out):
-        if out is None:
-            return None
-        # lxml element
-        try:
-            import lxml.etree as lxml_et
-
-            if isinstance(out, lxml_et._Element):
-                return lxml_et.tostring(out, encoding="unicode")
-        except Exception:
-            pass
-
-        if isinstance(out, ET.Element):
-            return ET.tostring(out, encoding="unicode")
-
-        if isinstance(out, bytes):
-            try:
-                return out.decode()
-            except Exception:
-                return None
-
-        if hasattr(out, "toxml") and callable(out.toxml):
-            try:
-                return out.toxml()
-            except Exception:
-                pass
-
-        try:
-            s = str(out)
-            if s.lstrip().startswith("<"):
-                return s
-            return s
-        except Exception:
-            return None
-
-    for name in candidates:
-        attr = getattr(obj, name, None)
-        if callable(attr):
-            try:
-                out = attr()
-            except Exception:
-                out = None
-            if out:
-                norm = normalize(out)
-                if norm:
-                    return norm
-        elif attr is not None:
-            norm = normalize(attr)
-            if norm:
-                return norm
-
-    if hasattr(obj, "mjcf"):
-        nested = getattr(obj, "mjcf")
-        if nested is not None:
-            return try_serialize(nested)
-
-    return None
-
-
-@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3")
-def main(dict_cfg: DictConfig) -> None:
-    # Compose typed config like the rest of the project
-    cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg))
-
-    # Check for a simulation morphology override (points to a YAML file)
-    sim_override = None
-    try:
-        sim_override = dict_cfg.get("simulation", {}).get("morphology_override", None)
-    except Exception:
-        sim_override = getattr(getattr(cfg, "simulation", None), "morphology_override", None)
-
-    if sim_override:
-        override_path = Path(hydra.utils.to_absolute_path(sim_override))
-        if not override_path.exists():
-            raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}")
-        import yaml
-
-        with open(override_path, "r") as f:
-            override_dict = yaml.safe_load(f)
-        env_morphology = OmegaConf.to_object(
-            OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict)
-        )
-    else:
-        env_morphology = cfg.morphology
-
-    # Ensure we have a MorphologyConfig instance
-    if isinstance(env_morphology, dict):
-        morph_cfg = MorphologyConfig(**env_morphology)
-    else:
-        morph_cfg = env_morphology
-
-    # Build morphology via project factory (same as runtime)
-    morph = BrittleStarEnvFactory.create_morphology(morph_cfg)
-
-    xml_text = try_serialize(morph)
-    if xml_text is None and hasattr(morph, "mjcf"):
-        xml_text = try_serialize(morph.mjcf)
-
-    if xml_text is None:
-        raise RuntimeError(
-            "Failed to serialize morphology to MJCF/XML. Inspect the `morph` object interactively."
-        )
-
-    # Prefer explicit CLI override `dump_out=...` if provided, otherwise choose a sensible default.
-    dump_out = None
-    try:
-        dump_out = dict_cfg.get("dump_out", None)
-    except Exception:
-        dump_out = None
-
-    if dump_out is None:
-        # If the user passed a morphology group on the CLI (e.g. morphology=3_arms),
-        # use a descriptive default path under `runs/morphologies/`.
-        morph_name = None
-        for a in sys.argv[1:]:
-            if a.startswith("morphology="):
-                morph_name = a.split("=", 1)[1]
-                break
-
-        default_out = f"runs/morphologies/{morph_name}.xml" if morph_name else "morphology.xml"
-        out_path = Path(hydra.utils.to_absolute_path(default_out))
-    else:
-        out_path = Path(hydra.utils.to_absolute_path(str(dump_out)))
-
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    out_path.write_text(xml_text)
-    print(f"Wrote MJCF XML to {out_path}")
-
-
-if __name__ == "__main__":
-    register_configs()
-    main()

scripts/tools/dump_mjcf.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""
+Dump MJCF XML for a brittle-star morphology using the project's Hydra configs.
+
+Usage examples:
+
+  # Use a named morphology config from configs/morphology (Hydra style)
+  uv run python scripts/analysis/dump_mjcf.py morphology=3_arms
+
+  # Use a morphology override YAML (same key as simulation.morphology_override)
+  uv run python scripts/analysis/dump_mjcf.py \
+    simulation.morphology_override=configs/morphology/3_arms.yaml
+
+Output path:
+    Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in
+    current directory or `runs/morphologies/<name>.xml`.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+import hydra
+import yaml
+from omegaconf import DictConfig, OmegaConf
+
+from brittle_star_project.configs.register_configs import register_configs
+from brittle_star_project.environment.env_config import MorphologyConfig
+from brittle_star_project.environment.factory import BrittleStarEnvFactory
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def extract_xml_string(obj: Any) -> Optional[str]:
+    """
+    Attempts to serialize the morphology object to an XML string by checking
+    common dm_control and internal API methods.
+    """
+    serialization_methods = [
+        "to_xml_string",
+        "to_xml",
+        "to_string",
+        "to_mjcf",
+        "to_mjcf_string",
+        "get_mjcf",
+        "get_mjcf_str",
+        "export_to_xml_string",
+    ]
+
+    # If the object itself has an 'mjcf' attribute, try to serialize that instead
+    target_obj = getattr(obj, "mjcf", obj)
+
+    for method_name in serialization_methods:
+        method = getattr(target_obj, method_name, None)
+        if callable(method):
+            try:
+                xml_data = method()
+                # Safely handle both string and byte responses
+                if isinstance(xml_data, str):
+                    return xml_data
+                elif isinstance(xml_data, bytes):
+                    return xml_data.decode("utf-8")
+            except Exception as e:
+                logger.debug(f"Method {method_name}() failed during serialization: {e}")
+
+    return None
+
+
+def resolve_output_path(cfg: DictConfig) -> Path:
+    """Determines the appropriate output path for the MJCF XML."""
+    dump_out = cfg.get("dump_out", None)
+    if dump_out is not None:
+        return Path(hydra.utils.to_absolute_path(str(dump_out)))
+
+    morph_name = "morphology"
+    for arg in sys.argv[1:]:
+        if arg.startswith("morphology="):
+            morph_name = arg.split("=", 1)[1]
+            break
+
+    default_out = (
+        f"runs/morphologies/{morph_name}.xml" if morph_name != "morphology" else "morphology.xml"
+    )
+    return Path(hydra.utils.to_absolute_path(default_out))
+
+
+@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3")
+def main(cfg: DictConfig) -> None:
+    """Main entry point to construct the morphology and dump its XML."""
+    logger.info("Initializing morphology construction...")
+
+    # Extract morphology config safely using dict `.get()` to avoid OmegaConf AttributeErrors
+    simulation_cfg = cfg.get("simulation", cfg)
+    override_path = simulation_cfg.get("morphology_override", None)
+
+    if override_path:
+        logger.info(f"Using morphology override: {override_path}")
+        with open(hydra.utils.to_absolute_path(override_path), "r") as f:
+            data = yaml.safe_load(f) or {}
+        morph_cfg = MorphologyConfig(**data)
+    else:
+        # Fallback to default simulation morphology, or an empty base config
+        morph_node = simulation_cfg.get("morphology", cfg.get("morphology", None))
+
+        if morph_node is not None:
+            # Convert OmegaConf node to dict and instantiate MorphologyConfig.
+            # This ensures any missing keys gracefully fall back to the dataclass defaults.
+            morph_dict = OmegaConf.to_container(morph_node, resolve=True)
+            if isinstance(morph_dict, dict):
+                # Filter to avoid unexpected kwargs if the dataclass is strictly defined
+                if dataclasses.is_dataclass(MorphologyConfig):
+                    valid_keys = {f.name for f in dataclasses.fields(MorphologyConfig)}
+                    morph_dict = {k: v for k, v in morph_dict.items() if k in valid_keys}
+                morph_cfg = MorphologyConfig(**morph_dict)
+            else:
+                morph_cfg = MorphologyConfig()
+        else:
+            morph_cfg = MorphologyConfig()
+
+    morphology = BrittleStarEnvFactory.create_morphology(morph_cfg)
+
+    xml_text = extract_xml_string(morphology)
+    if not xml_text:
+        raise RuntimeError("Failed to serialize morphology to MJCF/XML. ")
+
+    out_path = resolve_output_path(cfg)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with out_path.open("w", encoding="utf-8") as f:
+        f.write(xml_text)
+
+    logger.info(f"Successfully exported MJCF XML to: {out_path}")
+
+
+if __name__ == "__main__":
+    register_configs()
+    main()