897aad9cc4
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.
84 lines
2.4 KiB
Python
84 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
from pathlib import Path
|
|
|
|
import mujoco
|
|
|
|
from .controller import PDGaitController
|
|
from .sim import load_model, reset_to_stand, snapshot
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(description="Export MuJoCo walker trajectory CSV")
|
|
p.add_argument("--model", type=str, default=None, help="Path to MJCF model")
|
|
p.add_argument("--seconds", type=float, default=10.0, help="Duration to export")
|
|
p.add_argument("--out", type=str, required=True, help="Output CSV path")
|
|
p.add_argument("--sample-hz", type=float, default=120.0, help="Sampling frequency")
|
|
return p.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
out_path = Path(args.out)
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
model, data = load_model(args.model)
|
|
reset_to_stand(model, data)
|
|
controller = PDGaitController(model, data)
|
|
|
|
sample_dt = 1.0 / args.sample_hz
|
|
next_sample = 0.0
|
|
total_steps = int(args.seconds / model.opt.timestep)
|
|
rows = []
|
|
|
|
for _ in range(total_steps):
|
|
controller.step(data.time)
|
|
mujoco.mj_step(model, data)
|
|
if data.time >= next_sample:
|
|
rows.append(snapshot(data))
|
|
next_sample += sample_dt
|
|
|
|
with out_path.open("w", newline="", encoding="utf-8") as f:
|
|
writer = csv.writer(f)
|
|
writer.writerow(
|
|
[
|
|
"t",
|
|
"torso_x",
|
|
"torso_y",
|
|
"torso_z",
|
|
"near_f_x",
|
|
"near_f_z",
|
|
"near_g_x",
|
|
"near_g_z",
|
|
"far_f_x",
|
|
"far_f_z",
|
|
"far_g_x",
|
|
"far_g_z",
|
|
]
|
|
)
|
|
for r in rows:
|
|
writer.writerow(
|
|
[
|
|
f"{r.t:.6f}",
|
|
f"{r.torso_x:.6f}",
|
|
f"{r.torso_y:.6f}",
|
|
f"{r.torso_z:.6f}",
|
|
f"{r.near_f_x:.6f}",
|
|
f"{r.near_f_z:.6f}",
|
|
f"{r.near_g_x:.6f}",
|
|
f"{r.near_g_z:.6f}",
|
|
f"{r.far_f_x:.6f}",
|
|
f"{r.far_f_z:.6f}",
|
|
f"{r.far_g_x:.6f}",
|
|
f"{r.far_g_z:.6f}",
|
|
]
|
|
)
|
|
|
|
print(f"Wrote {len(rows)} samples to {out_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|