141 lines
5 KiB
Python
141 lines
5 KiB
Python
#!/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()
|