98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from dataclasses import dataclass
|
|
|
|
import mujoco
|
|
|
|
from .controller import PDGaitConfig, PDGaitController
|
|
from .sim import load_model, reset_to_stand, snapshot
|
|
|
|
|
|
@dataclass
|
|
class EvalResult:
|
|
seconds: float
|
|
distance_m: float
|
|
avg_speed_mps: float
|
|
min_torso_z_m: float
|
|
max_torso_z_m: float
|
|
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 run_eval(seconds: float, gait_hz: float, kp: float, kd: float, model_path: str | None = None) -> EvalResult:
|
|
model, data = load_model(model_path)
|
|
reset_to_stand(model, data)
|
|
controller = PDGaitController(model, data, PDGaitConfig(frequency_hz=gait_hz, kp=kp, kd=kd))
|
|
|
|
start = snapshot(data)
|
|
min_torso = start.torso_z
|
|
max_torso = start.torso_z
|
|
|
|
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 = min(min_torso, s.torso_z)
|
|
max_torso = max(max_torso, s.torso_z)
|
|
if torso_hit_floor(model, data) or s.torso_z < 0.08:
|
|
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_z_m=min_torso,
|
|
max_torso_z_m=max_torso,
|
|
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")
|
|
return p.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
result = run_eval(
|
|
seconds=args.seconds,
|
|
gait_hz=args.gait_hz,
|
|
kp=args.kp,
|
|
kd=args.kd,
|
|
model_path=args.model,
|
|
)
|
|
|
|
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_z_m: {result.min_torso_z_m:.3f}")
|
|
print(f"- max_torso_z_m: {result.max_torso_z_m:.3f}")
|
|
print(f"- fallen: {result.fallen}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|