335 lines
11 KiB
Python
335 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import datetime
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import yaml
|
|
|
|
from brittle_star_project import BrittleStarEnv
|
|
from brittle_star_project.evaluation.policy import ControlPolicy
|
|
from brittle_star_project.evaluation.rollout import (
|
|
EpisodeResult,
|
|
_get_observations,
|
|
_get_xy_distance_to_target,
|
|
_target_reached,
|
|
_maybe_clip_action,
|
|
)
|
|
|
|
|
|
def create_evaluation_dir(model_path: Path) -> Path:
|
|
"""Create a unique timestamped directory for saving evaluation results."""
|
|
timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
eval_dir = model_path.parent / f"{model_path.stem}_evaluations" / f"eval_{timestamp}"
|
|
eval_dir.mkdir(parents=True, exist_ok=True)
|
|
return eval_dir
|
|
|
|
|
|
def _ensure_offscreen_size(model, width: int, height: int) -> None:
|
|
vis_global = getattr(getattr(model, "vis", None), "global_", None)
|
|
if vis_global is None:
|
|
return
|
|
vis_global.offwidth = int(max(width, vis_global.offwidth))
|
|
vis_global.offheight = int(max(height, vis_global.offheight))
|
|
|
|
|
|
def _apply_camera_overrides(
|
|
model,
|
|
*,
|
|
camera_fovy: dict[int, float] | None = None,
|
|
camera_xyz: tuple[
|
|
dict[int, float] | None,
|
|
dict[int, float] | None,
|
|
dict[int, float] | None,
|
|
] = (None, None, None),
|
|
) -> None:
|
|
if not camera_fovy and not (camera_xyz[0] or camera_xyz[1] or camera_xyz[2]):
|
|
return
|
|
|
|
ncam = int(getattr(model, "ncam", 0))
|
|
for cam_id, fovy in (camera_fovy or {}).items():
|
|
if cam_id < 0 or cam_id >= ncam:
|
|
raise ValueError(f"Camera id {cam_id} is out of range")
|
|
model.cam_fovy[cam_id] = float(fovy)
|
|
|
|
for cam_id, x in (camera_xyz[0] or {}).items():
|
|
if cam_id < 0 or cam_id >= ncam:
|
|
raise ValueError(f"Camera id {cam_id} is out of range")
|
|
model.cam_pos[cam_id][0] = float(x)
|
|
|
|
for cam_id, y in (camera_xyz[1] or {}).items():
|
|
if cam_id < 0 or cam_id >= ncam:
|
|
raise ValueError(f"Camera id {cam_id} is out of range")
|
|
model.cam_pos[cam_id][1] = float(y)
|
|
|
|
for cam_id, z in (camera_xyz[2] or {}).items():
|
|
if cam_id < 0 or cam_id >= ncam:
|
|
raise ValueError(f"Camera id {cam_id} is out of range")
|
|
model.cam_pos[cam_id][2] = float(z)
|
|
|
|
|
|
def hex_to_rgba(hex_color: str, alpha: float) -> np.ndarray:
|
|
color = hex_color.lstrip("#")
|
|
if len(color) != 6:
|
|
raise ValueError(f"Expected a 6-digit hex color, got {hex_color!r}")
|
|
red = int(color[0:2], 16) / 255.0
|
|
green = int(color[2:4], 16) / 255.0
|
|
blue = int(color[4:6], 16) / 255.0
|
|
return np.asarray([red, green, blue, float(alpha)], dtype=np.float32)
|
|
|
|
|
|
def save_evaluation_metadata(
|
|
eval_dir: Path,
|
|
*,
|
|
morphology_override_path: str | None,
|
|
seed: int,
|
|
max_steps: int | None,
|
|
result: EpisodeResult,
|
|
) -> None:
|
|
"""Save metadata about the evaluation run."""
|
|
metadata = {
|
|
"timestamp": datetime.datetime.now().isoformat(),
|
|
"morphology_override": morphology_override_path,
|
|
"seed": seed,
|
|
"max_steps": max_steps,
|
|
"result": {
|
|
"return": float(result.return_),
|
|
"length": int(result.length),
|
|
"reached_target": bool(result.reached_target),
|
|
"final_xy_dist": float(result.final_xy_dist)
|
|
if result.final_xy_dist is not None
|
|
else None,
|
|
},
|
|
}
|
|
with open(eval_dir / "evaluation_metadata.yaml", "w") as f:
|
|
yaml.safe_dump(metadata, f, sort_keys=False)
|
|
|
|
|
|
def record_episode(
|
|
*,
|
|
env: BrittleStarEnv,
|
|
policy: ControlPolicy,
|
|
seed: int,
|
|
max_steps: int,
|
|
action_low: np.ndarray | None,
|
|
action_high: np.ndarray | None,
|
|
action_mask: np.ndarray | None = None,
|
|
output_path: Path,
|
|
camera_id: int = 1,
|
|
fps: int = 60,
|
|
width: int = 640,
|
|
height: int = 480,
|
|
target_xy: tuple[float, float] | None = None,
|
|
) -> EpisodeResult:
|
|
"""Run an episode headlessly and record a video using MuJoCo's Renderer and imageio.
|
|
|
|
Args:
|
|
env: The environment.
|
|
policy: The policy agent.
|
|
seed: Random seed.
|
|
max_steps: Maximum number of steps.
|
|
action_low: Minimum action values.
|
|
action_high: Maximum action values.
|
|
action_mask: Boolean mask for the actions.
|
|
output_path: Where to save the .mp4 file.
|
|
camera_id: Camera index to use for rendering (1 is usually close-up).
|
|
fps: Frames per second for the video.
|
|
width: Video width.
|
|
height: Video height.
|
|
"""
|
|
try:
|
|
import imageio
|
|
import mujoco
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"Video recording requires 'imageio' and 'mujoco'. "
|
|
"Please install the evaluation dependencies: `uv pip install .[evaluation]`"
|
|
) from e
|
|
|
|
state = env.reset(seed=seed, target_position=target_xy)
|
|
model = state.mj_model
|
|
data = state.mj_data
|
|
|
|
_ensure_offscreen_size(model, width, height)
|
|
|
|
renderer = mujoco.Renderer(model, width=width, height=height)
|
|
ep_return = 0.0
|
|
observations = _get_observations(state)
|
|
prev_dist = _get_xy_distance_to_target(observations) if observations else None
|
|
initial_dist = prev_dist
|
|
reached_target = _target_reached(state=state)
|
|
|
|
frames = []
|
|
steps = 0
|
|
|
|
for _ in range(int(max_steps)):
|
|
# Capture frame
|
|
renderer.update_scene(data, camera=camera_id)
|
|
frames.append(renderer.render())
|
|
|
|
# Step environment
|
|
obs_dict = observations or {}
|
|
action = policy.act(observations=obs_dict)
|
|
if action_mask is not None:
|
|
action = action[action_mask]
|
|
action = _maybe_clip_action(action, action_low, action_high)
|
|
|
|
state = env.step(state=state, action=action)
|
|
steps += 1
|
|
|
|
observations = _get_observations(state)
|
|
cur_dist = _get_xy_distance_to_target(observations) if observations else None
|
|
if prev_dist is not None and cur_dist is not None:
|
|
ep_return += prev_dist - cur_dist
|
|
prev_dist = cur_dist
|
|
|
|
reached_target = _target_reached(state=state)
|
|
if reached_target:
|
|
break
|
|
|
|
# Capture final frame
|
|
renderer.update_scene(data, camera=camera_id)
|
|
frames.append(renderer.render())
|
|
renderer.close()
|
|
|
|
# Save video
|
|
imageio.mimsave(str(output_path), frames, fps=fps)
|
|
|
|
final_dist = _get_xy_distance_to_target(observations) if observations else None
|
|
return EpisodeResult(
|
|
return_=ep_return,
|
|
length=steps,
|
|
reached_target=reached_target,
|
|
final_xy_dist=final_dist,
|
|
initial_target_distance=initial_dist,
|
|
)
|
|
|
|
|
|
def record_episode_multi_camera(
|
|
*,
|
|
env: BrittleStarEnv,
|
|
policy: ControlPolicy,
|
|
seed: int,
|
|
max_steps: int,
|
|
action_low: np.ndarray | None,
|
|
action_high: np.ndarray | None,
|
|
output_paths: dict[int, Path],
|
|
action_mask: np.ndarray | None = None,
|
|
camera_ids: list[int] | None = None,
|
|
camera_fovy: dict[int, float] | None = None,
|
|
camera_xyz: tuple[
|
|
dict[int, float] | None,
|
|
dict[int, float] | None,
|
|
dict[int, float] | None,
|
|
] = (None, None, None),
|
|
target_xy: tuple[float, float] | None = None,
|
|
robot_color: str | None = None,
|
|
fps: int = 60,
|
|
width: int = 640,
|
|
height: int = 480,
|
|
) -> EpisodeResult:
|
|
"""Run one episode and render multiple camera views to separate files."""
|
|
try:
|
|
import imageio
|
|
import mujoco
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"Video recording requires 'imageio' and 'mujoco'. "
|
|
"Please install the evaluation dependencies: `uv pip install .[evaluation]`"
|
|
) from e
|
|
|
|
if camera_ids is None:
|
|
camera_ids = list(output_paths.keys())
|
|
|
|
for cam_id in camera_ids:
|
|
if cam_id not in output_paths:
|
|
raise ValueError(f"Missing output path for camera {cam_id}")
|
|
|
|
output_paths = {cam_id: output_paths[cam_id] for cam_id in camera_ids}
|
|
|
|
for path in output_paths.values():
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
state = env.reset(seed=seed, target_position=(target_xy[0], target_xy[1], 0.0))
|
|
model = state.mj_model
|
|
data = state.mj_data
|
|
|
|
_apply_camera_overrides(model, camera_fovy=camera_fovy, camera_xyz=camera_xyz)
|
|
|
|
if robot_color is not None:
|
|
robot_body_id = mujoco.mj_name2id(
|
|
model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
|
|
)
|
|
if robot_body_id < 0:
|
|
raise ValueError("Body 'BrittleStarMorphology/central_disk' not found in the model")
|
|
|
|
robot_rgba = hex_to_rgba(robot_color, 1.0)
|
|
body_parent = model.body_parentid
|
|
robot_body_ids = {int(robot_body_id)}
|
|
|
|
for body_id in range(1, int(model.nbody)):
|
|
current_body_id = int(body_id)
|
|
while current_body_id not in (-1, 0, int(robot_body_id)):
|
|
current_body_id = int(body_parent[current_body_id])
|
|
if current_body_id == int(robot_body_id):
|
|
robot_body_ids.add(body_id)
|
|
|
|
for geom_id in range(int(model.ngeom)):
|
|
if int(model.geom_bodyid[geom_id]) in robot_body_ids:
|
|
model.geom_rgba[geom_id][:] = robot_rgba
|
|
|
|
_ensure_offscreen_size(model, width, height)
|
|
|
|
renderer = mujoco.Renderer(model, width=width, height=height)
|
|
writers = {
|
|
cam_id: imageio.get_writer(str(path), fps=fps) for cam_id, path in output_paths.items()
|
|
}
|
|
|
|
ep_return = 0.0
|
|
observations = _get_observations(state)
|
|
prev_dist = _get_xy_distance_to_target(observations) if observations else None
|
|
initial_dist = prev_dist
|
|
reached_target = _target_reached(state=state)
|
|
|
|
steps = 0
|
|
try:
|
|
for _ in range(int(max_steps)):
|
|
for cam_id in camera_ids:
|
|
renderer.update_scene(data, camera=cam_id)
|
|
writers[cam_id].append_data(renderer.render())
|
|
|
|
obs_dict = observations or {}
|
|
action = policy.act(observations=obs_dict)
|
|
if action_mask is not None:
|
|
action = action[action_mask]
|
|
action = _maybe_clip_action(action, action_low, action_high)
|
|
|
|
state = env.step(state=state, action=action)
|
|
steps += 1
|
|
|
|
observations = _get_observations(state)
|
|
cur_dist = _get_xy_distance_to_target(observations) if observations else None
|
|
if prev_dist is not None and cur_dist is not None:
|
|
ep_return += prev_dist - cur_dist
|
|
prev_dist = cur_dist
|
|
|
|
reached_target = _target_reached(state=state)
|
|
if reached_target:
|
|
break
|
|
|
|
for cam_id in camera_ids:
|
|
renderer.update_scene(data, camera=cam_id)
|
|
writers[cam_id].append_data(renderer.render())
|
|
finally:
|
|
renderer.close()
|
|
for writer in writers.values():
|
|
writer.close()
|
|
|
|
final_dist = _get_xy_distance_to_target(observations) if observations else None
|
|
return EpisodeResult(
|
|
return_=ep_return,
|
|
length=steps,
|
|
reached_target=reached_target,
|
|
final_xy_dist=final_dist,
|
|
initial_target_distance=initial_dist,
|
|
)
|