diff --git a/configs/example.json b/configs/example.json deleted file mode 100644 index 8b646db..0000000 --- a/configs/example.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "morphology": { - "num_arms": 2, - "num_segments_per_arm": 4, - "use_p_control": true, - "use_torque_control": false - } -} \ No newline at end of file diff --git a/configs/example.yaml b/configs/example.yaml new file mode 100644 index 0000000..14bf3d7 --- /dev/null +++ b/configs/example.yaml @@ -0,0 +1,5 @@ +morphology: + num_arms: 2 + num_segments_per_arm: 4 + use_p_control: true + use_torque_control: false diff --git a/scripts/simulate.py b/scripts/simulate.py index 28062fa..957f6fa 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -10,7 +10,7 @@ from brittle_star_project import ( SimulationConfig, simulate_policy, ) -from brittle_star_project.environment import from_json +from brittle_star_project.environment import from_file from brittle_star_project.rl import RLModel # imports concrete models via rl.__init__ from brittle_star_project.rl.base import get_rl_model_registry @@ -44,7 +44,7 @@ def parse_args() -> argparse.Namespace: def main() -> None: args = parse_args() - morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json") + morphology_cfg, arena_cfg, env_cfg = from_file("../configs/test.yaml") # ======= ENVIRONMENT SETUP ======= diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py index a693286..78083e9 100644 --- a/src/brittle_star_project/environment/env_config.py +++ b/src/brittle_star_project/environment/env_config.py @@ -1,7 +1,6 @@ from __future__ import annotations from dataclasses import dataclass, field -import json from .env_types import Task @@ -51,14 +50,11 @@ class EnvConfig: def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]: - """Load configurations from a JSON or YAML file.""" - with open(path, "r") as f: - if path.endswith(".yaml") or path.endswith(".yml"): - import yaml + """Load configurations from a YAML file.""" + import yaml - config_dict = yaml.safe_load(f) - else: - config_dict = json.load(f) + with open(path, "r") as f: + config_dict = yaml.safe_load(f) morphology = MorphologyConfig(**config_dict.get("morphology", {})) arena = ArenaConfig(**config_dict.get("arena", {})) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 71380b0..40229c5 100644 --- a/src/experiment_logger/unified_logger.py +++ b/src/experiment_logger/unified_logger.py @@ -7,9 +7,9 @@ This logger ensures all experimental data is preserved by writing to: """ import datetime -import json import logging import subprocess +import yaml import sys import time from pathlib import Path @@ -99,7 +99,7 @@ class UnifiedLogger: self.metrics_dir = self.run_dir / "metrics" self.metrics_dir.mkdir(exist_ok=True) - self.config_file = self.run_dir / "config.json" + self.config_file = self.run_dir / "config.yaml" # Setup standard Python logging mirror self.text_log_file = self.run_dir / "run.log" @@ -196,7 +196,7 @@ class UnifiedLogger: """Save configuration to disk.""" try: with open(self.config_file, "w") as f: - json.dump(self.config, f, indent=2) + yaml.dump(self.config, f, default_flow_style=False, indent=2, sort_keys=False) self.info(f"Config saved to {self.config_file}") except Exception as e: self.error(f"Error saving config: {e}") @@ -263,10 +263,10 @@ class UnifiedLogger: return try: - metrics_file = self.metrics_dir / "metrics.jsonl" + metrics_file = self.metrics_dir / "metrics.yaml" with open(metrics_file, "a") as f: for metric in self.metrics_buffer: - # Convert numpy/jax types to native Python types for JSON serialization + # Convert numpy/jax types to native Python types for YAML serialization serializable_metric = {} for k, v in metric.items(): if hasattr(v, "item"): # numpy/jax scalar @@ -275,7 +275,8 @@ class UnifiedLogger: serializable_metric[k] = v.tolist() else: serializable_metric[k] = v - f.write(json.dumps(serializable_metric) + "\n") + f.write("---\n") + yaml.dump(serializable_metric, f, default_flow_style=False) self.metrics_buffer.clear() except Exception as e: self.error(f"Error flushing metrics: {e}") @@ -298,9 +299,9 @@ class UnifiedLogger: # Save metadata if provided if metadata: - metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.json" + metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.yaml" with open(metadata_path, "w") as f: - json.dump(metadata, f, indent=2) + yaml.dump(metadata, f, default_flow_style=False) self.info(f"Checkpoint saved: {checkpoint_path}") @@ -334,9 +335,9 @@ class UnifiedLogger: f.write(flax.serialization.to_bytes(params)) if metadata: - metadata_path = self.run_dir / "final_model_metadata.json" + metadata_path = self.run_dir / "final_model_metadata.yaml" with open(metadata_path, "w") as f: - json.dump(metadata, f, indent=2) + yaml.dump(metadata, f, default_flow_style=False) self.info(f"Final model saved: {final_model_path}")