180 lines
6 KiB
Python
180 lines
6 KiB
Python
"""Simulate a trained policy in the MuJoCo viewer.
|
|
|
|
Automatically extracts the training configuration (morphology, environment, etc.)
|
|
from the sidecar metadata YAML file to ensure simulation perfectly matches training.
|
|
Override simulation settings via CLI, e.g.:
|
|
uv run scripts/simulate.py \
|
|
simulation.morphology_override=configs/morphology/3_arms.yaml \
|
|
simulation.model_path=runs/.../final_model.flax
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
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.evaluation.checkpoint import load_metadata, metadata_to_configs
|
|
from brittle_star_project.evaluation.eval_env_builder import build_eval_env
|
|
from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer
|
|
from brittle_star_project.evaluation.video import (
|
|
record_episode,
|
|
create_evaluation_dir,
|
|
save_evaluation_metadata,
|
|
)
|
|
|
|
|
|
@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
|
|
def main(dict_cfg: DictConfig) -> None:
|
|
# 1. Hydra composes ONLY SimulationSettings
|
|
cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg))
|
|
sim_cfg = cfg.simulation
|
|
|
|
model_path_str = sim_cfg.model_path
|
|
if model_path_str is None:
|
|
raise ValueError(
|
|
"simulation.model_path must be set to a .flax checkpoint (e.g. final_model.flax)"
|
|
)
|
|
|
|
model_path = Path(hydra.utils.to_absolute_path(model_path_str))
|
|
if model_path.suffix != ".flax":
|
|
raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.")
|
|
|
|
# 2. Discover + load sidecar metadata YAML
|
|
metadata_override = None
|
|
if sim_cfg.metadata_path is not None:
|
|
metadata_override = Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path))
|
|
|
|
metadata = load_metadata(model_path, metadata_override)
|
|
|
|
# 3. Reconstruct typed configs from metadata
|
|
training = metadata_to_configs(metadata)
|
|
|
|
seed = int(cfg.experiment.seed)
|
|
|
|
# 4-7. Build evaluation environment and policy
|
|
override_path = None
|
|
if sim_cfg.morphology_override is not None:
|
|
override_path = Path(hydra.utils.to_absolute_path(sim_cfg.morphology_override))
|
|
|
|
bundle = build_eval_env(
|
|
model_path=model_path,
|
|
training=training,
|
|
metadata=metadata,
|
|
morphology_override_path=override_path,
|
|
)
|
|
|
|
env = bundle.env
|
|
policy = bundle.policy
|
|
action_low = bundle.action_low
|
|
action_high = bundle.action_high
|
|
action_mask = bundle.action_mask
|
|
|
|
state0 = env.reset(seed=seed)
|
|
|
|
# 8. Run simulation
|
|
headless = bool(sim_cfg.headless)
|
|
max_steps = sim_cfg.max_steps
|
|
|
|
if sim_cfg.record_video:
|
|
if max_steps is None:
|
|
raise ValueError("simulation.max_steps is required when simulation.record_video=true")
|
|
|
|
max_steps_i = int(max_steps)
|
|
if max_steps_i <= 0:
|
|
raise ValueError("simulation.max_steps must be > 0")
|
|
|
|
if sim_cfg.video_output_path is None:
|
|
eval_dir = create_evaluation_dir(model_path)
|
|
output_path = eval_dir / "simulation.mp4"
|
|
else:
|
|
output_path = Path(hydra.utils.to_absolute_path(sim_cfg.video_output_path))
|
|
eval_dir = output_path.parent
|
|
eval_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
result = record_episode(
|
|
env=env,
|
|
policy=policy,
|
|
seed=seed,
|
|
max_steps=max_steps_i,
|
|
action_low=action_low,
|
|
action_high=action_high,
|
|
action_mask=action_mask,
|
|
output_path=output_path,
|
|
camera_id=sim_cfg.camera_id,
|
|
width=sim_cfg.video_width,
|
|
height=sim_cfg.video_height,
|
|
fps=sim_cfg.video_fps,
|
|
)
|
|
|
|
save_evaluation_metadata(
|
|
eval_dir=eval_dir,
|
|
morphology_override_path=sim_cfg.morphology_override,
|
|
seed=seed,
|
|
max_steps=max_steps_i,
|
|
result=result,
|
|
)
|
|
final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
|
|
print(f"Video saved to {output_path}")
|
|
print(
|
|
"episode done: "
|
|
f"return={result.return_:.6f}, len={result.length}, "
|
|
f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}"
|
|
)
|
|
elif headless:
|
|
if max_steps is None:
|
|
raise ValueError("simulation.max_steps is required when simulation.headless=true")
|
|
|
|
max_steps_i = int(max_steps)
|
|
if max_steps_i <= 0:
|
|
raise ValueError("simulation.max_steps must be > 0")
|
|
|
|
result = rollout_headless(
|
|
env=env,
|
|
policy=policy,
|
|
seed=seed,
|
|
max_steps=max_steps_i,
|
|
action_low=action_low,
|
|
action_high=action_high,
|
|
action_mask=action_mask,
|
|
)
|
|
final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
|
|
print(
|
|
"episode done: "
|
|
f"return={result.return_:.6f}, len={result.length}, "
|
|
f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}"
|
|
)
|
|
else:
|
|
max_steps_val = None
|
|
if max_steps is not None:
|
|
max_steps_i = int(max_steps)
|
|
if max_steps_i <= 0:
|
|
raise ValueError("simulation.max_steps must be > 0")
|
|
max_steps_val = max_steps_i
|
|
|
|
model_dt = float(state0.mj_model.opt.timestep)
|
|
control_dt = model_dt * float(training.environment.num_physics_steps_per_control_step)
|
|
|
|
rollout_viewer(
|
|
env=env,
|
|
policy=policy,
|
|
seed=seed,
|
|
state=state0,
|
|
control_dt=control_dt,
|
|
max_steps=max_steps_val,
|
|
action_low=action_low,
|
|
action_high=action_high,
|
|
action_mask=action_mask,
|
|
)
|
|
|
|
env.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
register_configs()
|
|
main()
|