242 lines
10 KiB
Python
242 lines
10 KiB
Python
"""Actual checkpoint inference, capturing first-terminal physics before mjlab resets."""
|
|
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from dataclasses import asdict
|
|
from pathlib import Path
|
|
from statistics import fmean
|
|
from types import SimpleNamespace
|
|
|
|
# Also executable as a fresh-interpreter seed worker (never fork a CUDA context).
|
|
for source in (Path(__file__).resolve().parents[1], Path(__file__).resolve().parents[2]):
|
|
if str(source) not in sys.path:
|
|
sys.path.insert(0, str(source))
|
|
|
|
import torch
|
|
from mjlab.envs import ManagerBasedRlEnv
|
|
from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper
|
|
from mjlab.sensor import ContactMatch, ContactSensorCfg
|
|
from mjlab.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls
|
|
from scripts.train import _load_reward_config, _load_task_config
|
|
from src.tasks.obstacle_avoidance.env_cfg import apply_obstacle_configuration
|
|
from task_config import OBSTACLE_TASK
|
|
from tuning.obstacle_scoring import METRICS, STEPS, protocol, score_trajectory
|
|
from tuning.schema import apply_reward_configuration
|
|
|
|
|
|
class FirstEpisodeRecorder:
|
|
"""The termination hook runs post-physics, before _reset_idx (including first terminal)."""
|
|
|
|
def __init__(self, env, layout):
|
|
self.env = env
|
|
self.layout = layout
|
|
self.samples = [[] for _ in range(env.num_envs)]
|
|
self.previous = torch.zeros((env.num_envs, 12), device=env.device)
|
|
self.calls = 0
|
|
self.compute = env.termination_manager.compute
|
|
env.termination_manager.compute = self.capture
|
|
|
|
def capture(self):
|
|
env = self.env
|
|
terminal = self.compute()
|
|
robot = env.scene["robot"].data
|
|
# Like mjlab termination itself, derived pose is one physics substep old.
|
|
position = robot.root_link_pos_w
|
|
local_xy = position[:, :2] - env.scene.env_origins[:, :2]
|
|
xy = local_xy + torch.tensor(self.layout["spawn"][:2], device=env.device)
|
|
boxes = self.layout["boxes"][1:] # Never include the support floor.
|
|
clearance = torch.full((env.num_envs,), 0.5, device=env.device)
|
|
if boxes:
|
|
centers = torch.tensor([b["pos"][:2] for b in boxes], device=env.device)
|
|
sizes = torch.tensor([b["size"][:2] for b in boxes], device=env.device)
|
|
clearance = (
|
|
((xy[:, None, :] - centers).abs() - sizes).clamp(min=0).norm(dim=2).amin(dim=1)
|
|
)
|
|
clearance = (clearance - 0.3).clamp(min=0) # Conservative body footprint radius.
|
|
forces = env.scene["evaluation_obstacles"].data.force_history
|
|
collision = forces.norm(dim=-1).flatten(1).amax(dim=1) > 1.0
|
|
else:
|
|
collision = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
|
nonfoot = env.scene["nonfoot_ground_touch"].data.force_history
|
|
collision |= nonfoot.norm(dim=-1).flatten(1).amax(dim=1) > 10.0
|
|
fall = (robot.projected_gravity_b[:, 2] > -0.3420201433) | (position[:, 2] < 0.12)
|
|
actions = env.action_manager.action
|
|
delta = (actions - self.previous).square().mean(dim=1)
|
|
self.previous.copy_(actions)
|
|
rays = env.scene["forward_scan"].data.distances
|
|
hits = (
|
|
((rays >= 0) & (rays <= env.scene["forward_scan"].cfg.max_distance)).float().mean(dim=1)
|
|
)
|
|
distance = env.command_manager.get_term("twist").errors()[1]
|
|
values = torch.stack((distance, clearance, delta, hits, collision, fall, terminal), dim=1)
|
|
if not torch.isfinite(values).all():
|
|
raise ValueError("Nonfinite rollout measurements")
|
|
keys = ("distance", "clearance", "action_delta", "ray_hit", "collision", "fall", "terminal")
|
|
for samples, row in zip(self.samples, values.cpu().tolist(), strict=True):
|
|
if not samples or not samples[-1]["terminal"]:
|
|
samples.append(dict(zip(keys, row, strict=True)))
|
|
self.calls += 1
|
|
return terminal
|
|
|
|
def metrics(self, horizon=STEPS):
|
|
if self.calls != horizon:
|
|
raise ValueError("Incomplete rollout horizon")
|
|
metrics = [score_trajectory(samples, horizon) for samples in self.samples]
|
|
return {key: fmean(item[key] for item in metrics) for key in METRICS}
|
|
|
|
|
|
def configure_seed(task_id, cfg, scenario, reward):
|
|
env_cfg = load_env_cfg(task_id, play=False)
|
|
env_cfg.seed = scenario["seed"]
|
|
env_cfg.scene.num_envs = cfg.num_envs
|
|
# Keep the benchmark's declared pair fixed; training itself randomizes every reset.
|
|
apply_obstacle_configuration(env_cfg, scenario["taskConfig"], randomize_navigation=False)
|
|
apply_reward_configuration(env_cfg, reward, task_id)
|
|
env_cfg.curriculum = {}
|
|
env_cfg.observations["actor"].enable_corruption = False
|
|
env_cfg.events.pop("push_robot", None)
|
|
# Primary geom names are literal compiled static terrain geoms, not a regex.
|
|
names = tuple(f"terrain_{i}" for i in range(1, len(scenario["terrain"]["boxes"])))
|
|
if names:
|
|
env_cfg.scene.sensors += (
|
|
ContactSensorCfg(
|
|
name="evaluation_obstacles",
|
|
primary=ContactMatch(mode="geom", pattern=names),
|
|
secondary=ContactMatch(mode="subtree", pattern="base_link", entity="robot"),
|
|
fields=("force",),
|
|
reduce="maxforce",
|
|
history_length=env_cfg.decimation,
|
|
),
|
|
)
|
|
return env_cfg
|
|
|
|
|
|
def evaluate_seed(task_id, cfg, scenario, reward):
|
|
torch.manual_seed(scenario["seed"])
|
|
env_cfg = configure_seed(task_id, cfg, scenario, reward)
|
|
agent_cfg = load_rl_cfg(task_id)
|
|
env = ManagerBasedRlEnv(env_cfg, device=cfg.device)
|
|
wrapped = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
|
|
try:
|
|
runner_cls = load_runner_cls(task_id) or MjlabOnPolicyRunner
|
|
runner = runner_cls(wrapped, asdict(agent_cfg), log_dir=None, device=wrapped.device)
|
|
# rsl_rl actor state_dict includes observation_normalizer buffers. Strict load
|
|
# rejects incompatible architecture/statistics; no fixture or zero-action fallback.
|
|
runner.load(
|
|
str(Path(cfg.checkpoint).resolve(strict=True)),
|
|
load_cfg={"actor": True},
|
|
strict=True,
|
|
map_location=str(wrapped.device),
|
|
)
|
|
policy = runner.get_inference_policy(device=str(wrapped.device))
|
|
obs, _ = env.reset(seed=scenario["seed"])
|
|
recorder = FirstEpisodeRecorder(env, scenario["terrain"])
|
|
with torch.inference_mode():
|
|
for _ in range(STEPS):
|
|
actions = policy(obs)
|
|
if not torch.isfinite(actions).all():
|
|
raise ValueError("Nonfinite policy actions")
|
|
obs, _, _, _ = wrapped.step(actions)
|
|
return {
|
|
"seed": scenario["seed"],
|
|
"metrics": recorder.metrics(),
|
|
"episodes": cfg.num_envs,
|
|
"rolloutSteps": STEPS,
|
|
}
|
|
finally:
|
|
wrapped.close()
|
|
|
|
|
|
def run_obstacle_evaluation(task_id, cfg):
|
|
from tuning.obstacle_scoring import SEEDS
|
|
|
|
if tuple(cfg.seeds) != SEEDS or cfg.steps_per_seed != STEPS:
|
|
raise ValueError("Obstacle evaluation seeds/horizon are fixed")
|
|
if not cfg.task_config or not cfg.reward_config:
|
|
raise ValueError("Obstacle evaluation requires task and parameter configurations")
|
|
raw = json.loads(Path(cfg.task_config).read_text())
|
|
custom = _load_task_config(task_id, cfg.task_config, raw["seed"])
|
|
reward = _load_reward_config(cfg.reward_config, None, OBSTACLE_TASK)
|
|
fixed = protocol(custom, cfg.num_envs)
|
|
seeds = evaluate_isolated_seeds(task_id, cfg, fixed, reward)
|
|
return {
|
|
"protocol": fixed,
|
|
"seedMetrics": seeds,
|
|
"metrics": {key: fmean(item["metrics"][key] for item in seeds) for key in METRICS},
|
|
}
|
|
|
|
|
|
def file_hash(path):
|
|
digest = hashlib.sha256()
|
|
with Path(path).open("rb") as stream:
|
|
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def evaluate_isolated_seeds(task_id, cfg, fixed, reward):
|
|
from tuning.obstacle_scoring import validate_evaluation
|
|
|
|
seeds = []
|
|
# Inherit the evaluation parent's process group: manager cancellation kills
|
|
# both parent and the active seed. subprocess.run kills/reaps on timeout too.
|
|
with tempfile.TemporaryDirectory(prefix="go2-obstacle-eval-") as directory:
|
|
root = Path(directory)
|
|
checkpoint = root / "checkpoint.pt"
|
|
source_hash = file_hash(cfg.checkpoint)
|
|
shutil.copyfile(cfg.checkpoint, checkpoint)
|
|
if file_hash(checkpoint) != source_hash:
|
|
raise ValueError("Checkpoint changed during snapshot")
|
|
for scenario in fixed["scenarios"]:
|
|
request = root / f"request-{scenario['seed']}.json"
|
|
output = root / f"result-{scenario['seed']}.json"
|
|
request.write_text(
|
|
json.dumps(
|
|
{
|
|
"taskId": task_id,
|
|
"config": {**asdict(cfg), "checkpoint": str(checkpoint)},
|
|
"scenario": scenario,
|
|
"reward": reward,
|
|
"checkpointSha256": source_hash,
|
|
},
|
|
allow_nan=False,
|
|
)
|
|
)
|
|
subprocess.run(
|
|
[sys.executable, str(Path(__file__).resolve()), str(request), str(output)],
|
|
check=True,
|
|
timeout=1200,
|
|
shell=False,
|
|
)
|
|
result = json.loads(output.read_text())
|
|
if (
|
|
result.pop("checkpointSha256", None) != source_hash
|
|
or file_hash(checkpoint) != source_hash
|
|
):
|
|
raise ValueError("Seed checkpoint identity mismatch")
|
|
seeds.append(result)
|
|
candidate = {
|
|
"protocol": fixed,
|
|
"seedMetrics": seeds,
|
|
"metrics": {key: fmean(item["metrics"][key] for item in seeds) for key in METRICS},
|
|
}
|
|
validate_evaluation(candidate, fixed)
|
|
return seeds
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import src.tasks # noqa: F401
|
|
|
|
request_path, output_path = map(Path, sys.argv[1:])
|
|
request = json.loads(request_path.read_text())
|
|
config = SimpleNamespace(**request["config"])
|
|
if file_hash(config.checkpoint) != request["checkpointSha256"]:
|
|
raise ValueError("Worker checkpoint identity mismatch")
|
|
result = evaluate_seed(request["taskId"], config, request["scenario"], request["reward"])
|
|
result["checkpointSha256"] = file_hash(config.checkpoint)
|
|
output_path.write_text(json.dumps(result, allow_nan=False))
|