434 lines
14 KiB
Python
434 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import itertools
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
import flax
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
|
|
from brittle_star_project import (
|
|
Backend,
|
|
)
|
|
from brittle_star_project.environment import ArenaConfig, EnvConfig, MorphologyConfig, from_file
|
|
|
|
|
|
def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
|
|
"""Flatten the env's observation dict into a 1D vector.
|
|
|
|
concatenates values in the dict's iteration order and skips empty arrays.
|
|
"""
|
|
|
|
parts: list[jnp.ndarray] = []
|
|
for v in obs_dict.values():
|
|
arr = jnp.asarray(v)
|
|
if arr.size == 0:
|
|
continue
|
|
parts.append(arr.reshape((-1,)))
|
|
|
|
if not parts:
|
|
return jnp.zeros((0,), dtype=jnp.float32)
|
|
return jnp.concatenate(parts, axis=0)
|
|
|
|
|
|
# A minimal policy class to load a CleanRL/Flax checkpoint and run inference.
|
|
class CleanRLPPOPolicy:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
sensor_params: Any,
|
|
actor_params: Any,
|
|
action_dim: int,
|
|
) -> None:
|
|
from brittle_star_project.MLPs.mlps import Actor, GenericDenseLayersWithActivation
|
|
|
|
self._sensor = GenericDenseLayersWithActivation()
|
|
self._actor = Actor(action_dim=action_dim)
|
|
self._sensor_apply = jax.jit(self._sensor.apply)
|
|
self._actor_apply = jax.jit(self._actor.apply)
|
|
self._params = {
|
|
"sensor_params": sensor_params,
|
|
"actor_params": actor_params,
|
|
}
|
|
|
|
@staticmethod
|
|
def load(
|
|
path: Path,
|
|
*,
|
|
action_dim: int,
|
|
) -> "CleanRLPPOPolicy":
|
|
def _get_index(container: Any, idx: int) -> Any:
|
|
if isinstance(container, (list, tuple)):
|
|
return container[idx]
|
|
if isinstance(container, dict):
|
|
return container.get(idx, container.get(str(idx)))
|
|
raise KeyError(idx)
|
|
|
|
def _looks_like_indexed_dict(container: Any) -> bool:
|
|
return (
|
|
isinstance(container, dict)
|
|
and container
|
|
and all(str(k).isdigit() for k in container.keys())
|
|
)
|
|
|
|
def _parse_checkpoint(restored_obj: Any) -> tuple[Any, Any, Any, Any, Any]:
|
|
"""Extract checkpoint parts.
|
|
|
|
Returns (args_dict, sensor_params, actor_params, critic_params,
|
|
feature_extractor_params).
|
|
|
|
`PPOTrainer` saves:
|
|
flax.serialization.to_bytes(
|
|
[vars(args), [sensor, actor, critic, feature_extractor]]
|
|
)
|
|
|
|
`msgpack_restore()` may restore lists as dicts keyed by string
|
|
indices ("0", "1", ...), so we accept both shapes.
|
|
"""
|
|
|
|
args_part: Any | None = None
|
|
params_part: Any = restored_obj
|
|
|
|
if isinstance(restored_obj, (list, tuple)) and len(restored_obj) >= 2:
|
|
args_part = restored_obj[0]
|
|
params_part = restored_obj[1]
|
|
elif _looks_like_indexed_dict(restored_obj) and (
|
|
"0" in restored_obj or "1" in restored_obj
|
|
):
|
|
args_part = restored_obj.get("0", restored_obj.get(0))
|
|
params_part = restored_obj.get("1", restored_obj.get(1))
|
|
|
|
if _looks_like_indexed_dict(params_part):
|
|
sensor_params = _get_index(params_part, 0)
|
|
actor_params = _get_index(params_part, 1)
|
|
critic_params = _get_index(params_part, 2)
|
|
feature_extractor_params = _get_index(params_part, 3)
|
|
if sensor_params is None or actor_params is None:
|
|
raise ValueError("Missing required params in checkpoint")
|
|
return (
|
|
args_part,
|
|
sensor_params,
|
|
actor_params,
|
|
critic_params,
|
|
feature_extractor_params,
|
|
)
|
|
|
|
if isinstance(params_part, (list, tuple)) and len(params_part) >= 2:
|
|
sensor_params = params_part[0]
|
|
actor_params = params_part[1]
|
|
critic_params = params_part[2] if len(params_part) >= 3 else None
|
|
feature_extractor_params = params_part[3] if len(params_part) >= 4 else None
|
|
return (
|
|
args_part,
|
|
sensor_params,
|
|
actor_params,
|
|
critic_params,
|
|
feature_extractor_params,
|
|
)
|
|
|
|
raise ValueError(
|
|
f"Unexpected checkpoint structure in {path}. "
|
|
"Expected [args_dict, [sensor_params, actor_params, critic_params, "
|
|
"feature_extractor_params]] "
|
|
"or an equivalent dict-indexed variant."
|
|
)
|
|
|
|
payload = path.read_bytes()
|
|
restored = flax.serialization.msgpack_restore(payload)
|
|
_args_dict, sensor_params, actor_params, _critic_params, _feature_extractor_params = (
|
|
_parse_checkpoint(restored)
|
|
)
|
|
|
|
return CleanRLPPOPolicy(
|
|
sensor_params=sensor_params,
|
|
actor_params=actor_params,
|
|
action_dim=action_dim,
|
|
)
|
|
|
|
def act(self, *, observations: dict[str, Any]) -> np.ndarray:
|
|
obs = _flatten_obs_dict(observations)
|
|
hidden = self._sensor_apply(self._params["sensor_params"], obs)
|
|
mean, _log_std = self._actor_apply(self._params["actor_params"], hidden)
|
|
|
|
# Always evaluate with the actor mean.
|
|
# (Sampling adds exploration noise, which is useful for training but not for evaluation.)
|
|
return np.asarray(mean, dtype=np.float32).ravel()
|
|
|
|
|
|
def _get_observations(state: Any) -> dict[str, Any] | None:
|
|
return getattr(state, "observations", None)
|
|
|
|
|
|
def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None:
|
|
return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0])
|
|
|
|
|
|
def _target_reached(*, state: Any) -> bool:
|
|
return bool(getattr(state, "terminated", False))
|
|
|
|
|
|
def _rollout_one_episode_headless(
|
|
*,
|
|
env: Any,
|
|
policy: CleanRLPPOPolicy,
|
|
seed: int,
|
|
max_steps: int,
|
|
) -> tuple[float, int, bool, float | None]:
|
|
"""Run one rollout up to `max_steps`.
|
|
|
|
Returns (return, length, reached_target, final_xy_dist).
|
|
"""
|
|
state = env.reset(seed=seed)
|
|
|
|
ep_return = 0.0
|
|
|
|
observations = _get_observations(state)
|
|
prev_dist = _get_xy_distance_to_target(observations)
|
|
reached_target = _target_reached(state=state)
|
|
|
|
# NOTE: In the MJC backend, `state.reward` is always 0.0.
|
|
# To get a meaningful return, we compute a simple progress reward:
|
|
# r_t = d_{t-1} - d_t
|
|
# where d is `xy_distance_to_target`.
|
|
steps = 0
|
|
for _ in range(int(max_steps)):
|
|
action = policy.act(observations=observations)
|
|
|
|
nu = int(state.mj_model.nu)
|
|
if nu > 0 and action.shape != (nu,):
|
|
raise ValueError(f"Policy returned action shape {action.shape}, expected ({nu},)")
|
|
|
|
state = env.step(state=state, action=action)
|
|
steps += 1
|
|
|
|
observations = _get_observations(state)
|
|
cur_dist = _get_xy_distance_to_target(observations)
|
|
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
|
|
|
|
final_dist = _get_xy_distance_to_target(observations)
|
|
return ep_return, steps, reached_target, final_dist
|
|
|
|
|
|
def _run_one_episode_viewer(
|
|
*,
|
|
env: Any,
|
|
policy: CleanRLPPOPolicy,
|
|
seed: int,
|
|
state: Any,
|
|
control_dt: float,
|
|
max_steps: int | None,
|
|
) -> None:
|
|
import mujoco.viewer
|
|
|
|
model = state.mj_model
|
|
data = state.mj_data
|
|
|
|
seed = int(seed)
|
|
episode_return = 0.0
|
|
observations = _get_observations(state)
|
|
prev_dist = _get_xy_distance_to_target(observations)
|
|
reached_target = _target_reached(state=state)
|
|
|
|
steps = 0
|
|
# Use the viewer as a context manager to avoid GLX teardown races
|
|
# (e.g. GLXBadDrawable from X_GLXSwapBuffers after a window is destroyed).
|
|
with mujoco.viewer.launch_passive(model, data) as viewer:
|
|
step_iter = range(int(max_steps)) if max_steps is not None else itertools.count()
|
|
for _step_idx in step_iter:
|
|
if not viewer.is_running():
|
|
break
|
|
step_start = time.time()
|
|
|
|
action = policy.act(observations=observations)
|
|
if model.nu > 0 and action.shape != (int(model.nu),):
|
|
raise ValueError(
|
|
f"Policy returned action shape {action.shape}, expected ({int(model.nu)},)"
|
|
)
|
|
|
|
# The passive viewer runs a GUI thread; protect MuJoCo state mutation.
|
|
with viewer.lock():
|
|
state = env.step(state=state, action=action)
|
|
|
|
if not viewer.is_running():
|
|
break
|
|
viewer.sync()
|
|
|
|
steps += 1
|
|
|
|
observations = _get_observations(state)
|
|
cur_dist = _get_xy_distance_to_target(observations)
|
|
if prev_dist is not None and cur_dist is not None:
|
|
episode_return += prev_dist - cur_dist
|
|
prev_dist = cur_dist
|
|
|
|
reached_target = _target_reached(state=state)
|
|
if reached_target:
|
|
break
|
|
|
|
remaining = control_dt - (time.time() - step_start)
|
|
if remaining > 0:
|
|
time.sleep(remaining)
|
|
|
|
dist = _get_xy_distance_to_target(observations)
|
|
dist_str = "n/a" if dist is None else f"{dist:.3f}"
|
|
print(
|
|
"episode done: "
|
|
f"return={episode_return:.6f}, len={steps}, "
|
|
f"target_reached={reached_target}, final_xy_dist={dist_str}"
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(
|
|
description="Run a trained policy for exactly one episode (viewer or headless)."
|
|
)
|
|
p.add_argument(
|
|
"--config-path",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"Path to an environment JSON config (morphology/arena/env). "
|
|
"If omitted, uses the environment defaults. "
|
|
"Relative paths are resolved from the repository root."
|
|
),
|
|
)
|
|
p.add_argument(
|
|
"--model",
|
|
type=str,
|
|
required=True,
|
|
help=("Path to the Flax checkpoint saved by scripts/train.py (final_model.flax)."),
|
|
)
|
|
p.add_argument(
|
|
"--headless",
|
|
action="store_true",
|
|
help="Run without the MuJoCo viewer (still exactly one episode).",
|
|
)
|
|
p.add_argument(
|
|
"--max-steps",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"Number of control steps to run. "
|
|
"In --headless mode this is required and acts as a fixed horizon. "
|
|
"In viewer mode the default is infinite (run until window closed or target reached)."
|
|
),
|
|
)
|
|
p.add_argument(
|
|
"--backend",
|
|
choices=[b.value for b in Backend],
|
|
default=Backend.MJC.value,
|
|
)
|
|
p.add_argument("--seed", type=int, default=0)
|
|
return p.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
from brittle_star_project.environment import (
|
|
BrittleStarEnv,
|
|
BrittleStarEnvFactory,
|
|
)
|
|
|
|
args = parse_args()
|
|
|
|
if args.config_path is None:
|
|
morphology_cfg = MorphologyConfig()
|
|
arena_cfg = ArenaConfig()
|
|
env_cfg = EnvConfig()
|
|
else:
|
|
repo_root = Path(__file__).resolve().parents[1]
|
|
config_path = Path(args.config_path)
|
|
if not config_path.is_absolute():
|
|
config_path = repo_root / config_path
|
|
morphology_cfg, arena_cfg, env_cfg = from_file(str(config_path))
|
|
|
|
# ======= ENVIRONMENT SETUP =======
|
|
|
|
backend = Backend(args.backend)
|
|
|
|
factory = BrittleStarEnvFactory()
|
|
raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
|
|
env = BrittleStarEnv(
|
|
raw_env,
|
|
backend=backend,
|
|
config=env_cfg,
|
|
morphology_config=morphology_cfg,
|
|
)
|
|
|
|
seed_for_env = int(args.seed) if args.seed is not None else 0
|
|
state = env.reset(seed=seed_for_env)
|
|
|
|
# ======= MODEL SETUP =======
|
|
|
|
# Extract the number of actuators (nu) from the environment's model, so we can pass it to the
|
|
# policy/model.
|
|
nu = int(state.mj_model.nu)
|
|
|
|
model_path = Path(args.model)
|
|
if model_path.suffix != ".flax":
|
|
raise ValueError(
|
|
f"Expected the training artifact '.flax', got '{model_path.name}'."
|
|
)
|
|
|
|
policy = CleanRLPPOPolicy.load(
|
|
model_path,
|
|
action_dim=nu,
|
|
)
|
|
|
|
default_seed = seed_for_env
|
|
|
|
# ======= SIMULATION =======
|
|
|
|
if args.headless:
|
|
if args.max_steps is None:
|
|
raise ValueError("--max-steps is required in --headless mode")
|
|
max_steps = int(args.max_steps)
|
|
if max_steps <= 0:
|
|
raise ValueError("--max-steps must be > 0")
|
|
|
|
ep_seed = int(args.seed) if args.seed is not None else default_seed
|
|
ep_return, ep_len, reached_target, final_dist = _rollout_one_episode_headless(
|
|
env=env,
|
|
policy=policy,
|
|
seed=ep_seed,
|
|
max_steps=max_steps,
|
|
)
|
|
final_dist_str = "n/a" if final_dist is None else f"{final_dist:.3f}"
|
|
print(
|
|
"episode done: "
|
|
f"return={ep_return:.6f}, len={ep_len}, "
|
|
f"target_reached={reached_target}, final_xy_dist={final_dist_str}"
|
|
)
|
|
else:
|
|
max_steps: int | None
|
|
if args.max_steps is None:
|
|
max_steps = None
|
|
else:
|
|
max_steps = int(args.max_steps)
|
|
if max_steps <= 0:
|
|
raise ValueError("--max-steps must be > 0")
|
|
|
|
model_dt = float(state.mj_model.opt.timestep)
|
|
control_dt = model_dt * float(env_cfg.num_physics_steps_per_control_step)
|
|
_run_one_episode_viewer(
|
|
env=env,
|
|
policy=policy,
|
|
seed=int(args.seed) if args.seed is not None else default_seed,
|
|
state=state,
|
|
control_dt=control_dt,
|
|
max_steps=max_steps,
|
|
)
|
|
|
|
env.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|