Enable physically emergent lateral tilt and yaw in the MuJoCo scaffold so side dynamics are visible in viewer, logs, and exports for future stability tuning.
108 lines
3.3 KiB
Python
108 lines
3.3 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_y_m: float
|
|
max_torso_y_m: 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_y = start.torso_y
|
|
max_torso_y = start.torso_y
|
|
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_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)
|
|
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_y_m=min_torso_y,
|
|
max_torso_y_m=max_torso_y,
|
|
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_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"- fallen: {result.fallen}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|