223 lines
7.1 KiB
Python
223 lines
7.1 KiB
Python
"""Deterministic, headless evaluation for Unitree Go2 velocity policies."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sys
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from statistics import fmean, pstdev
|
|
from typing import Literal
|
|
|
|
TRAINER_ROOT = Path(__file__).resolve().parents[1]
|
|
SERVICE_ROOT = TRAINER_ROOT.parent
|
|
for source_root in (TRAINER_ROOT, SERVICE_ROOT):
|
|
if str(source_root) not in sys.path:
|
|
sys.path.insert(0, str(source_root))
|
|
|
|
import torch
|
|
import tyro
|
|
import warp as wp
|
|
|
|
if not hasattr(wp, "context"):
|
|
from warp._src import context as warp_context
|
|
|
|
wp.context = warp_context # type: ignore[attr-defined]
|
|
|
|
from mjlab.envs import ManagerBasedRlEnv
|
|
from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper
|
|
from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
|
|
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
|
|
from mjlab.utils.torch import configure_torch_backends
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
from tuning.schema import apply_reward_configuration, validate_configuration
|
|
|
|
SCENARIOS = (
|
|
(0.0, 0.0, 0.0),
|
|
(0.5, 0.0, 0.0),
|
|
(1.0, 0.0, 0.0),
|
|
(1.5, 0.0, 0.0),
|
|
(0.0, 0.5, 0.0),
|
|
(0.0, -0.5, 0.0),
|
|
(0.0, 0.0, 0.5),
|
|
(0.0, 0.0, -0.5),
|
|
(0.8, 0.25, 0.35),
|
|
)
|
|
METRIC_NAMES = (
|
|
"linear_velocity_rmse",
|
|
"angular_velocity_rmse",
|
|
"mean_action_acc",
|
|
"orientation_error",
|
|
"slip_velocity",
|
|
"mechanical_power",
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EvaluateConfig:
|
|
checkpoint: str
|
|
output: str
|
|
reward_config: str | None = None
|
|
num_envs: int = 256
|
|
steps_per_seed: int = 1000
|
|
seeds: tuple[int, ...] = field(default_factory=lambda: (101, 202, 303))
|
|
device: str | None = None
|
|
gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _fixed_commands(command: torch.Tensor) -> torch.Tensor:
|
|
values = torch.as_tensor(SCENARIOS, device=command.device, dtype=command.dtype)
|
|
indexes = torch.arange(command.shape[0], device=command.device) % values.shape[0]
|
|
command[:] = values[indexes]
|
|
return command
|
|
|
|
|
|
def _evaluate_seed(task_id: str, cfg: EvaluateConfig, seed: int) -> dict[str, float]:
|
|
torch.manual_seed(seed)
|
|
env_cfg = load_env_cfg(task_id, play=False)
|
|
agent_cfg = load_rl_cfg(task_id)
|
|
env_cfg.seed = seed
|
|
env_cfg.scene.num_envs = cfg.num_envs
|
|
env_cfg.curriculum = {}
|
|
env_cfg.observations["actor"].enable_corruption = False
|
|
env_cfg.events.pop("push_robot", None)
|
|
twist_cfg = env_cfg.commands["twist"]
|
|
assert isinstance(twist_cfg, UniformVelocityCommandCfg)
|
|
twist_cfg.heading_command = False
|
|
twist_cfg.ranges.heading = None
|
|
twist_cfg.rel_heading_envs = 0.0
|
|
twist_cfg.rel_standing_envs = 0.0
|
|
twist_cfg.resampling_time_range = (1.0e9, 1.0e9)
|
|
if cfg.reward_config:
|
|
reward_path = Path(cfg.reward_config).expanduser().resolve(strict=True)
|
|
with reward_path.open(encoding="utf-8") as stream:
|
|
apply_reward_configuration(env_cfg, validate_configuration(json.load(stream)))
|
|
|
|
env = ManagerBasedRlEnv(cfg=env_cfg, device=cfg.device or "cuda:0")
|
|
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)
|
|
runner.load(
|
|
str(Path(cfg.checkpoint).expanduser().resolve(strict=True)),
|
|
load_cfg={"actor": True},
|
|
strict=True,
|
|
map_location=str(wrapped.device),
|
|
)
|
|
policy = runner.get_inference_policy(device=str(wrapped.device))
|
|
twist = wrapped.unwrapped.command_manager.get_term("twist")
|
|
_fixed_commands(twist.command)
|
|
obs = wrapped.get_observations()
|
|
|
|
sums = {name: 0.0 for name in METRIC_NAMES}
|
|
samples = 0
|
|
terminations = 0
|
|
completions = 0
|
|
with torch.inference_mode():
|
|
for _ in range(cfg.steps_per_seed):
|
|
_fixed_commands(twist.command)
|
|
obs = wrapped.get_observations()
|
|
actions = policy(obs)
|
|
obs, _rewards, _dones, _extras = wrapped.step(actions)
|
|
manager = wrapped.unwrapped.metrics_manager
|
|
for index, name in enumerate(manager.active_terms):
|
|
if name in sums:
|
|
sums[name] += float(torch.sum(manager._step_values[:, index]).item())
|
|
samples += wrapped.num_envs
|
|
terminated = wrapped.unwrapped.reset_terminated
|
|
timed_out = wrapped.unwrapped.reset_time_outs
|
|
terminations += int(torch.count_nonzero(terminated).item())
|
|
completions += int(torch.count_nonzero(terminated | timed_out).item())
|
|
result = {name: sums[name] / max(samples, 1) for name in METRIC_NAMES}
|
|
result["fall_rate"] = terminations / max(completions, wrapped.num_envs)
|
|
return result
|
|
finally:
|
|
wrapped.close()
|
|
|
|
|
|
def run_evaluation(task_id: str, cfg: EvaluateConfig) -> dict:
|
|
configure_torch_backends()
|
|
selected = cfg.gpu_ids
|
|
if selected is None:
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
|
device = "cpu"
|
|
else:
|
|
if selected == "all":
|
|
selected = [0]
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected))
|
|
device = cfg.device or "cuda:0"
|
|
os.environ["MUJOCO_GL"] = "egl"
|
|
cfg = EvaluateConfig(**{**asdict(cfg), "device": device})
|
|
|
|
checkpoint = Path(cfg.checkpoint).expanduser().resolve(strict=True)
|
|
output = Path(cfg.output).expanduser().resolve()
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
per_seed = [_evaluate_seed(task_id, cfg, seed) for seed in cfg.seeds]
|
|
metrics = {
|
|
key: fmean(seed_metrics[key] for seed_metrics in per_seed)
|
|
for key in (*METRIC_NAMES, "fall_rate")
|
|
}
|
|
deviations = {
|
|
key: pstdev(seed_metrics[key] for seed_metrics in per_seed)
|
|
for key in (*METRIC_NAMES, "fall_rate")
|
|
}
|
|
result = {
|
|
"protocolVersion": 1,
|
|
"taskId": task_id,
|
|
"checkpoint": checkpoint.name,
|
|
"checkpointSha256": _sha256(checkpoint),
|
|
"seeds": list(cfg.seeds),
|
|
"numEnvs": cfg.num_envs,
|
|
"stepsPerSeed": cfg.steps_per_seed,
|
|
"scenarios": [list(value) for value in SCENARIOS],
|
|
"metrics": metrics,
|
|
"metricStd": deviations,
|
|
"seedMetrics": [
|
|
{"seed": seed, "metrics": values}
|
|
for seed, values in zip(cfg.seeds, per_seed, strict=True)
|
|
],
|
|
}
|
|
output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
writer = SummaryWriter(log_dir=str(output.parent / "evaluation-events"))
|
|
try:
|
|
for name, value in metrics.items():
|
|
writer.add_scalar(f"Evaluation/{name}", value, 0)
|
|
finally:
|
|
writer.close()
|
|
print("MUJOCO_EVALUATION " + json.dumps({"output": str(output), "metrics": metrics}))
|
|
return result
|
|
|
|
|
|
def main() -> None:
|
|
import mjlab.tasks # noqa: F401
|
|
import src.tasks # noqa: F401
|
|
|
|
chosen_task, remaining = tyro.cli(
|
|
tyro.extras.literal_type_from_choices(list_tasks()),
|
|
add_help=False,
|
|
return_unknown_args=True,
|
|
config=mjlab.TYRO_FLAGS,
|
|
)
|
|
args = tyro.cli(
|
|
EvaluateConfig,
|
|
args=remaining,
|
|
prog=sys.argv[0] + f" {chosen_task}",
|
|
config=mjlab.TYRO_FLAGS,
|
|
)
|
|
run_evaluation(chosen_task, args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|