82 lines
2.3 KiB
Python
82 lines
2.3 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_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_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()
|