refactor: clean tools

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()

scripts/tools/extract_observation_bounds.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Empirically extract observation bounds (focused on joint velocities).
+
+This script creates a MuJoCo environment using the project's factory and
+randomly samples actions to discover observed maxima for selected
+observation keys (joint_velocity, joint_position, joint_actuator_force).
+
+Usage:
+  python scripts/extract_observation_bounds.py \
+    --morphology configs/morphology/3_arms.yaml --num-steps 5000 --seed 42
+
+If `--morphology` is omitted the default `MorphologyConfig()` is used.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import yaml
+import numpy as np
+
+from brittle_star_project import BrittleStarEnvFactory, BrittleStarEnv, Backend
+from brittle_star_project.environment.env_config import (
+    MorphologyConfig,
+    ArenaConfig,
+    EnvConfig,
+)
+
+
+def load_morphology(path: str | None) -> MorphologyConfig:
+    if path is None:
+        return MorphologyConfig()
+    p = Path(path)
+    if not p.exists():
+        raise FileNotFoundError(f"Morphology file not found: {p}")
+    with open(p, "r") as f:
+        data = yaml.safe_load(f) or {}
+    return MorphologyConfig(**data)
+
+
+def _extract_observations(state):
+    # Under different backends the returned state may be a dict or an object
+    obs = getattr(state, "observations", None)
+    if obs is None and isinstance(state, dict):
+        obs = state.get("observations", state)
+    return obs
+
+
+def find_empirical_bounds(
+    morph_cfg: MorphologyConfig,
+    arena_cfg: ArenaConfig,
+    env_cfg: EnvConfig,
+    num_steps: int = 5000,
+    seed: int = 42,
+) -> None:
+    factory = BrittleStarEnvFactory()
+    raw_env = factory.create_environment(Backend.MJC, morph_cfg, arena_cfg, env_cfg)
+    env = BrittleStarEnv(raw_env, backend=Backend.MJC, config=env_cfg, morphology_config=morph_cfg)
+
+    # Initial reset
+    state = env.reset(seed=seed)
+
+    # Determine action bounds
+    action_space = getattr(raw_env, "action_space", None)
+    if action_space is None:
+        raise RuntimeError("Environment missing `action_space`; cannot sample actions.")
+
+    action_low = np.asarray(action_space.low, dtype=np.float32)
+    action_high = np.asarray(action_space.high, dtype=np.float32)
+    action_shape = action_low.shape
+
+    # Track maximum absolute observed values
+    tracked_keys = ["joint_velocity", "joint_position", "joint_actuator_force"]
+    max_observed = {k: 0.0 for k in tracked_keys}
+
+    # Include observation at reset
+    obs0 = _extract_observations(state)
+    if isinstance(obs0, dict):
+        for k in tracked_keys:
+            if k in obs0:
+                max_observed[k] = max(max_observed[k], float(np.max(np.abs(np.asarray(obs0[k])))))
+
+    rng = np.random.RandomState(seed)
+    for i in range(num_steps):
+        u = rng.uniform(size=action_shape)
+        action = action_low + (action_high - action_low) * u
+
+        # Provide a numpy RNG to the env step; wrapper will pass it if accepted.
+        step_out = env.step(state=state, action=action, rng=env.make_rng(seed + i + 1))
+
+        # Unpack next state from common return conventions
+        if hasattr(step_out, "state"):
+            next_state = step_out.state
+        elif isinstance(step_out, (tuple, list)) and len(step_out) >= 1:
+            next_state = step_out[0]
+        else:
+            next_state = step_out
+
+        obs = _extract_observations(next_state)
+        if isinstance(obs, dict):
+            for k in tracked_keys:
+                if k in obs:
+                    val = float(np.max(np.abs(np.asarray(obs[k]))))
+                    if val > max_observed[k]:
+                        max_observed[k] = val
+
+        state = next_state
+
+    # Print recommended bounds with a 20% safety margin
+    print("\n--- Recommended Observation Bounds (20% margin) ---")
+    for k, v in max_observed.items():
+        if v == 0.0:
+            print(f"{k}: observed max 0.0 (increase sampling or inspect env)")
+        else:
+            safe = v * 1.2
+            print(f"{k}: [-{safe:.6f}, {safe:.6f}]  (observed max: {v:.6f})")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--morphology", type=str, default=None, help="Path to morphology YAML (optional)"
+    )
+    parser.add_argument(
+        "--num-steps", type=int, default=5000, help="Number of random steps to sample"
+    )
+    parser.add_argument("--seed", type=int, default=42, help="RNG seed")
+    args = parser.parse_args()
+
+    morph_cfg = load_morphology(args.morphology)
+    arena_cfg = ArenaConfig()
+    env_cfg = EnvConfig()
+
+    find_empirical_bounds(morph_cfg, arena_cfg, env_cfg, num_steps=args.num_steps, seed=args.seed)
+
+
+if __name__ == "__main__":
+    main()