124 lines
4.5 KiB
Python
124 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import mujoco
|
|
import numpy as np
|
|
|
|
from .controller import PDGaitController
|
|
from .geometry import RAW_POINTS, SCALE_TO_CM
|
|
from .sim import load_model, reset_to_stand, site_xyz, torso_xyz
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(description="Export MuJoCo trace JSON for web overlay")
|
|
p.add_argument("--model", type=str, default=None, help="Path to MJCF model")
|
|
p.add_argument("--seconds", type=float, default=10.0, help="Duration")
|
|
p.add_argument("--sample-hz", type=float, default=120.0, help="Sampling frequency")
|
|
p.add_argument("--out", type=str, default="../../artifacts/mujoco_trace.json", help="Output JSON path")
|
|
p.add_argument("--scale-cm", type=float, default=100.0, help="Meters to drawing units scale")
|
|
p.add_argument(
|
|
"--frame",
|
|
type=str,
|
|
choices=("body", "world"),
|
|
default="body",
|
|
help="Coordinate frame for exported traces (body matches kinematic overlay best)",
|
|
)
|
|
p.add_argument("--origin-x", type=float, default=RAW_POINTS["O"].x * SCALE_TO_CM, help="Web-space X origin")
|
|
p.add_argument("--origin-y", type=float, default=RAW_POINTS["O"].y * SCALE_TO_CM, help="Web-space Y origin")
|
|
return p.parse_args()
|
|
|
|
|
|
def to_web_point(
|
|
x_m: float,
|
|
z_m: float,
|
|
scale_cm: float,
|
|
origin_x: float,
|
|
origin_y: float,
|
|
x_ref_m: float,
|
|
) -> dict[str, float]:
|
|
return {
|
|
"x": origin_x + (x_m - x_ref_m) * scale_cm,
|
|
"y": origin_y - z_m * scale_cm,
|
|
}
|
|
|
|
|
|
def world_to_body_local_xy(data: mujoco.MjData, world_xyz: np.ndarray) -> tuple[float, float]:
|
|
torso = data.body("torso")
|
|
torso_pos = np.array(torso.xpos, dtype=np.float64)
|
|
torso_rot = np.array(torso.xmat, dtype=np.float64).reshape(3, 3)
|
|
local = torso_rot.T @ (world_xyz - torso_pos)
|
|
return float(local[0]), float(local[2])
|
|
|
|
|
|
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)
|
|
x_ref_m = 0.0 if args.frame == "body" else float(torso_xyz(data)[0])
|
|
|
|
sample_dt = 1.0 / args.sample_hz
|
|
next_sample = 0.0
|
|
total_steps = int(args.seconds / model.opt.timestep)
|
|
|
|
near_f = []
|
|
near_g = []
|
|
far_f = []
|
|
far_g = []
|
|
torso = []
|
|
|
|
for _ in range(total_steps):
|
|
controller.step(data.time)
|
|
mujoco.mj_step(model, data)
|
|
|
|
if data.time >= next_sample:
|
|
n_f_world = site_xyz(data, "near_F_site")
|
|
n_g_world = site_xyz(data, "near_G_site")
|
|
f_f_world = site_xyz(data, "far_F_site")
|
|
f_g_world = site_xyz(data, "far_G_site")
|
|
t_world = torso_xyz(data)
|
|
|
|
if args.frame == "body":
|
|
n_f_x, n_f_z = world_to_body_local_xy(data, n_f_world)
|
|
n_g_x, n_g_z = world_to_body_local_xy(data, n_g_world)
|
|
f_f_x, f_f_z = world_to_body_local_xy(data, f_f_world)
|
|
f_g_x, f_g_z = world_to_body_local_xy(data, f_g_world)
|
|
t_x, t_z = 0.0, 0.0
|
|
else:
|
|
n_f_x, n_f_z = float(n_f_world[0]), float(n_f_world[2])
|
|
n_g_x, n_g_z = float(n_g_world[0]), float(n_g_world[2])
|
|
f_f_x, f_f_z = float(f_f_world[0]), float(f_f_world[2])
|
|
f_g_x, f_g_z = float(f_g_world[0]), float(f_g_world[2])
|
|
t_x, t_z = float(t_world[0]), float(t_world[2])
|
|
|
|
near_f.append(to_web_point(n_f_x, n_f_z, args.scale_cm, args.origin_x, args.origin_y, x_ref_m))
|
|
near_g.append(to_web_point(n_g_x, n_g_z, args.scale_cm, args.origin_x, args.origin_y, x_ref_m))
|
|
far_f.append(to_web_point(f_f_x, f_f_z, args.scale_cm, args.origin_x, args.origin_y, x_ref_m))
|
|
far_g.append(to_web_point(f_g_x, f_g_z, args.scale_cm, args.origin_x, args.origin_y, x_ref_m))
|
|
torso.append(to_web_point(t_x, t_z, args.scale_cm, args.origin_x, args.origin_y, x_ref_m))
|
|
next_sample += sample_dt
|
|
|
|
payload = {
|
|
"format": "walkersim-mujoco-trace-v1",
|
|
"units": "cm-like drawing space",
|
|
"frame": args.frame,
|
|
"origin": {"x": args.origin_x, "y": args.origin_y},
|
|
"nearF": near_f,
|
|
"nearG": near_g,
|
|
"farF": far_f,
|
|
"farG": far_g,
|
|
"torso": torso,
|
|
}
|
|
out_path.write_text(json.dumps(payload), encoding="utf-8")
|
|
print(f"Wrote MuJoCo web trace JSON to {out_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|