File size: 3,847 Bytes
b359f03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """Cross-evaluate every policy on every terrain -> transfer matrix.
python eval.py --policies flat rough hard --terrains flat rough hard
Each cell is one policy on one terrain: how often it falls, how long it
survives, how well it tracks the commanded velocity. Writes results/matrix.csv.
"""
import argparse, csv, json, os, re
from dataclasses import asdict
os.environ.setdefault("MUJOCO_GL", "egl")
os.environ["WANDB_MODE"] = "disabled"
import torch # noqa: E402
from huggingface_hub import HfApi, hf_hub_download # noqa: E402
from mjlab.envs import ManagerBasedRlEnv # noqa: E402
from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper # noqa: E402
from mjlab.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls # noqa: E402
from mjlab.utils.torch import configure_torch_backends # noqa: E402
from terrain import TERRAINS # noqa: E402
REPO = "mitanshugoel/go1-terrain"
TASK = "Mjlab-Velocity-Rough-Unitree-Go1"
p = argparse.ArgumentParser()
p.add_argument("--policies", nargs="+", default=["flat", "rough", "hard"])
p.add_argument("--terrains", nargs="+", default=["flat", "rough", "hard"])
p.add_argument("--envs", type=int, default=512)
p.add_argument("--steps", type=int, default=2000) # 2 episodes at 1000 steps
p.add_argument("--out", default="results/matrix.csv")
a = p.parse_args()
api = HfApi(token=os.environ["HF_TOKEN"])
files = api.list_repo_files(REPO, repo_type="model")
def latest_ckpt(policy):
cand = sorted((int(m.group(1)), f) for f in files
if (m := re.search(rf"^{policy}/run/.*model_(\d+)\.pt$", f)))
if not cand:
raise FileNotFoundError(f"no checkpoint for policy '{policy}'")
return hf_hub_download(REPO, cand[-1][1], repo_type="model",
token=os.environ["HF_TOKEN"], local_dir="/tmp/ck")
configure_torch_backends()
dev = "cuda:0"
rows = []
for terrain in a.terrains:
cfg = load_env_cfg(TASK, play=True)
agent = load_rl_cfg(TASK)
cfg.scene.num_envs = a.envs
cfg.scene.terrain.terrain_generator = TERRAINS[terrain]
cfg.curriculum = {} # fixed command range across all cells
env = ManagerBasedRlEnv(cfg=cfg, device=dev, render_mode=None)
env = RslRlVecEnvWrapper(env, clip_actions=agent.clip_actions)
runner = (load_runner_cls(TASK) or MjlabOnPolicyRunner)(env, asdict(agent), device=dev)
for policy in a.policies:
runner.load(latest_ckpt(policy), load_cfg={"actor": True}, strict=True,
map_location=dev)
pol = runner.get_inference_policy(device=dev)
obs = env.get_observations()
rew, dones, alive = [], 0, torch.zeros(a.envs, device=dev)
lengths = []
with torch.inference_mode():
for _ in range(a.steps):
obs, r, d, _ = env.step(pol(obs))
rew.append(r.mean().item())
alive += 1
if d.any():
lengths += alive[d.bool()].tolist()
alive[d.bool()] = 0
dones += int(d.sum())
row = dict(policy=policy, terrain=terrain,
mean_reward=round(sum(rew) / len(rew), 4),
terminations=dones,
mean_episode_len=round(sum(lengths) / len(lengths), 1) if lengths else a.steps,
envs=a.envs, steps=a.steps)
print(json.dumps(row), flush=True)
rows.append(row)
env.close()
del env, runner
torch.cuda.empty_cache()
os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
with open(a.out, "w", newline="") as f:
wr = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
wr.writeheader()
wr.writerows(rows)
print("\ntransfer matrix (mean episode length, higher = fewer falls)")
print(f"{'policy':>8} " + " ".join(f"{t:>10}" for t in a.terrains))
for pol in a.policies:
cells = {r["terrain"]: r["mean_episode_len"] for r in rows if r["policy"] == pol}
print(f"{pol:>8} " + " ".join(f"{cells.get(t, '-'):>10}" for t in a.terrains))
|