walkersim/sim/mujoco/walker_sim/eval.py

from __future__ import annotations

import argparse
from dataclasses import dataclass

import mujoco
import numpy as np

from .controller import PDGaitConfig, PDGaitController
from .model_variants import describe_model_config, resolved_model_path
from .sim import load_model, reset_to_stand, snapshot


@dataclass
class EvalResult:
    seconds: float
    distance_m: float
    avg_speed_mps: float
    min_torso_y_m: float
    max_torso_y_m: float
    min_torso_z_m: float
    max_torso_z_m: float
    min_body_ab_z_m: float
    gait_cycles: float
    stability_window_ok: bool
    fallen: bool


def torso_hit_floor(model: mujoco.MjModel, data: mujoco.MjData) -> bool:
    torso_gid = model.geom("torso_geom").id
    floor_gid = model.geom("floor").id
    for i in range(data.ncon):
        c = data.contact[i]
        if (c.geom1 == torso_gid and c.geom2 == floor_gid) or (c.geom1 == floor_gid and c.geom2 == torso_gid):
            return True
    return False


def torso_point_world(data: mujoco.MjData, local_xyz: tuple[float, float, float]) -> np.ndarray:
    torso = data.body("torso")
    torso_pos = np.array(torso.xpos, dtype=np.float64)
    torso_rot = np.array(torso.xmat, dtype=np.float64).reshape(3, 3)
    local = np.array(local_xyz, dtype=np.float64)
    return torso_pos + torso_rot @ local


def upper_body_min_z(data: mujoco.MjData) -> float:
    # Torso side reference points used in playback (A/B on near and far sides).
    # If any of these approach the floor, treat as a practical fall/near-fall.
    points_local = [
        (-0.30, 0.08, 0.0),
        (0.30, 0.08, 0.0),
        (-0.30, -0.08, 0.0),
        (0.30, -0.08, 0.0),
    ]
    return min(float(torso_point_world(data, p)[2]) for p in points_local)


def upper_body_near_ground(data: mujoco.MjData, min_clearance_m: float) -> bool:
    return upper_body_min_z(data) < min_clearance_m


def run_eval(
    seconds: float,
    gait_hz: float,
    kp: float,
    kd: float,
    model_path: str | None = None,
    body_clearance_m: float = 0.08,
    foot_cm: float | None = None,
) -> EvalResult:
    gait_cycles = seconds * gait_hz
    stability_window_ok = gait_cycles >= 2.0

    with resolved_model_path(model_path, foot_cm) as resolved_path:
        model, data = load_model(resolved_path)
        reset_to_stand(model, data)
        controller = PDGaitController(model, data, PDGaitConfig(frequency_hz=gait_hz, kp=kp, kd=kd))

        start = snapshot(data)
        min_torso_y = start.torso_y
        max_torso_y = start.torso_y
        min_torso = start.torso_z
        max_torso = start.torso_z
        min_body_ab = upper_body_min_z(data)

        steps = int(seconds / model.opt.timestep)
        fallen = False

        for _ in range(steps):
            controller.step(data.time)
            mujoco.mj_step(model, data)
            s = snapshot(data)
            min_torso_y = min(min_torso_y, s.torso_y)
            max_torso_y = max(max_torso_y, s.torso_y)
            min_torso = min(min_torso, s.torso_z)
            max_torso = max(max_torso, s.torso_z)
            min_body_ab = min(min_body_ab, upper_body_min_z(data))
            if torso_hit_floor(model, data) or upper_body_near_ground(data, body_clearance_m):
                fallen = True

        end = snapshot(data)
        elapsed = max(end.t - start.t, 1e-9)
        distance = end.torso_x - start.torso_x

        return EvalResult(
            seconds=elapsed,
            distance_m=distance,
            avg_speed_mps=distance / elapsed,
            min_torso_y_m=min_torso_y,
            max_torso_y_m=max_torso_y,
            min_torso_z_m=min_torso,
            max_torso_z_m=max_torso,
            min_body_ab_z_m=min_body_ab,
            gait_cycles=gait_cycles,
            stability_window_ok=stability_window_ok,
            fallen=fallen,
        )


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Evaluate baseline MuJoCo walker")
    p.add_argument("--model", type=str, default=None, help="Path to MJCF model")
    p.add_argument("--seconds", type=float, default=10.0, help="Evaluation duration")
    p.add_argument("--gait-hz", type=float, default=0.9, help="Gait frequency")
    p.add_argument("--kp", type=float, default=24.0, help="PD proportional gain")
    p.add_argument("--kd", type=float, default=2.5, help="PD derivative gain")
    p.add_argument("--body-clearance-m", type=float, default=0.08, help="Min A/B torso-point floor clearance")
    p.add_argument("--foot-cm", type=float, default=None, help="Override lateral foot length in cm")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    print(f"effective_config: {describe_model_config(args.model, args.foot_cm)}")
    result = run_eval(
        seconds=args.seconds,
        gait_hz=args.gait_hz,
        kp=args.kp,
        kd=args.kd,
        model_path=args.model,
        body_clearance_m=args.body_clearance_m,
        foot_cm=args.foot_cm,
    )

    print("Walker evaluation")
    print(f"- duration_s:      {result.seconds:.3f}")
    print(f"- distance_m:      {result.distance_m:.3f}")
    print(f"- avg_speed_mps:   {result.avg_speed_mps:.3f}")
    print(f"- min_torso_y_m:   {result.min_torso_y_m:.3f}")
    print(f"- max_torso_y_m:   {result.max_torso_y_m:.3f}")
    print(f"- min_torso_z_m:   {result.min_torso_z_m:.3f}")
    print(f"- max_torso_z_m:   {result.max_torso_z_m:.3f}")
    print(f"- min_body_ab_z_m: {result.min_body_ab_z_m:.3f}")
    print(f"- gait_cycles:     {result.gait_cycles:.2f}")
    print(f"- stability_ok:    {result.stability_window_ok}")
    print(f"- fallen:          {result.fallen}")
    if not result.stability_window_ok:
        print("- warning:         evaluation window is short; use >=2 gait cycles for reliable fall detection")


if __name__ == "__main__":
    main()