87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
import mujoco
|
|
|
|
from .controller import PDGaitConfig, PDGaitController
|
|
from .eval import run_eval
|
|
from .sim import load_model, reset_to_stand, snapshot
|
|
|
|
|
|
def run_headless(model, data, seconds: float, controller: PDGaitController, sample_hz: float) -> None:
|
|
sample_dt = 1.0 / sample_hz
|
|
next_sample = 0.0
|
|
steps = int(seconds / model.opt.timestep)
|
|
|
|
for _ in range(steps):
|
|
controller.step(data.time)
|
|
mujoco.mj_step(model, data)
|
|
|
|
if data.time >= next_sample:
|
|
s = snapshot(data)
|
|
print(
|
|
f"t={s.t:6.3f} "
|
|
f"cranks=({s.near_crank_rad:+.2f},{s.far_crank_rad:+.2f}) "
|
|
f"torso=({s.torso_x:+.3f},{s.torso_z:+.3f}) "
|
|
f"nearF=({s.near_f_x:+.3f},{s.near_f_z:+.3f}) "
|
|
f"nearG=({s.near_g_x:+.3f},{s.near_g_z:+.3f}) "
|
|
f"farF=({s.far_f_x:+.3f},{s.far_f_z:+.3f}) "
|
|
f"farG=({s.far_g_x:+.3f},{s.far_g_z:+.3f})"
|
|
)
|
|
next_sample += sample_dt
|
|
|
|
|
|
def run_viewer(model, data, seconds: float, controller: PDGaitController) -> None:
|
|
try:
|
|
from mujoco import viewer
|
|
except Exception as err: # pragma: no cover
|
|
raise RuntimeError("MuJoCo viewer is unavailable in this environment") from err
|
|
|
|
with viewer.launch_passive(model, data) as v:
|
|
end_t = data.time + seconds
|
|
while v.is_running() and data.time < end_t:
|
|
controller.step(data.time)
|
|
mujoco.mj_step(model, data)
|
|
v.sync()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(description="Run MuJoCo walker baseline")
|
|
p.add_argument("--model", type=str, default=None, help="Path to MJCF model")
|
|
p.add_argument("--seconds", type=float, default=10.0, help="Simulation duration")
|
|
p.add_argument("--viewer", action="store_true", help="Launch interactive viewer")
|
|
p.add_argument("--sample-hz", type=float, default=5.0, help="Headless print sample rate")
|
|
p.add_argument("--gait-hz", type=float, default=0.9, help="Controller 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()
|
|
model, data = load_model(args.model)
|
|
reset_to_stand(model, data)
|
|
|
|
cfg = PDGaitConfig(frequency_hz=args.gait_hz, kp=args.kp, kd=args.kd)
|
|
controller = PDGaitController(model, data, cfg)
|
|
|
|
if args.viewer:
|
|
run_viewer(model, data, args.seconds, controller)
|
|
else:
|
|
run_headless(model, data, args.seconds, controller, sample_hz=args.sample_hz)
|
|
|
|
result = run_eval(
|
|
seconds=args.seconds,
|
|
gait_hz=args.gait_hz,
|
|
kp=args.kp,
|
|
kd=args.kd,
|
|
model_path=args.model,
|
|
)
|
|
print("---")
|
|
print(f"summary: speed={result.avg_speed_mps:.3f} m/s, distance={result.distance_m:.3f} m, fallen={result.fallen}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|