236 lines
8.5 KiB
Python
236 lines
8.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 playback JSON for web viewer")
|
|
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(
|
|
"--format",
|
|
type=str,
|
|
choices=("v2", "v1"),
|
|
default="v2",
|
|
help="Output schema version: v2=true 3D playback, v1=legacy 2D trace overlay",
|
|
)
|
|
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="world",
|
|
help="Coordinate frame for exported samples",
|
|
)
|
|
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 world_to_body_local_xyz(data: mujoco.MjData, world_xyz: np.ndarray) -> tuple[float, 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[1]), float(local[2])
|
|
|
|
|
|
def world_to_body_local_points(data: mujoco.MjData, points: dict[str, np.ndarray]) -> dict[str, dict[str, float]]:
|
|
out: dict[str, dict[str, float]] = {}
|
|
for key, xyz in points.items():
|
|
x, y, z = world_to_body_local_xyz(data, xyz)
|
|
out[key] = {"x": x, "y": y, "z": z}
|
|
return out
|
|
|
|
|
|
def numpy_points_to_dict(points: dict[str, np.ndarray]) -> dict[str, dict[str, float]]:
|
|
out: dict[str, dict[str, float]] = {}
|
|
for key, xyz in points.items():
|
|
out[key] = {
|
|
"x": float(xyz[0]),
|
|
"y": float(xyz[1]),
|
|
"z": float(xyz[2]),
|
|
}
|
|
return out
|
|
|
|
|
|
def torso_point_world(data: mujoco.MjData, local_xyz: tuple[float, float, float]) -> np.ndarray:
|
|
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 = np.array(local_xyz, dtype=np.float64)
|
|
return torso_pos + torso_rot @ local
|
|
|
|
|
|
def collect_world_points(data: mujoco.MjData) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], np.ndarray]:
|
|
near_points = {
|
|
"A": torso_point_world(data, (-0.30, 0.08, 0.0)),
|
|
"B": torso_point_world(data, (0.30, 0.08, 0.0)),
|
|
"O": np.array(data.body("near_crank").xpos, dtype=np.float64),
|
|
"C": np.array(data.body("near_leg_f").xpos, dtype=np.float64),
|
|
"D": site_xyz(data, "near_D_site"),
|
|
"E": site_xyz(data, "near_E_site"),
|
|
"F": site_xyz(data, "near_F_site"),
|
|
"G": site_xyz(data, "near_G_site"),
|
|
}
|
|
far_points = {
|
|
"A": torso_point_world(data, (-0.30, -0.08, 0.0)),
|
|
"B": torso_point_world(data, (0.30, -0.08, 0.0)),
|
|
"O": np.array(data.body("far_crank").xpos, dtype=np.float64),
|
|
"C": np.array(data.body("far_leg_f").xpos, dtype=np.float64),
|
|
"D": site_xyz(data, "far_D_site"),
|
|
"E": site_xyz(data, "far_E_site"),
|
|
"F": site_xyz(data, "far_F_site"),
|
|
"G": site_xyz(data, "far_G_site"),
|
|
}
|
|
torso = torso_xyz(data)
|
|
return near_points, far_points, torso
|
|
|
|
|
|
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_2d = []
|
|
frames = []
|
|
|
|
for _ in range(total_steps):
|
|
controller.step(data.time)
|
|
mujoco.mj_step(model, data)
|
|
|
|
if data.time >= next_sample:
|
|
near_world, far_world, torso_world = collect_world_points(data)
|
|
|
|
n_f_world = near_world["F"]
|
|
n_g_world = near_world["G"]
|
|
f_f_world = far_world["F"]
|
|
f_g_world = far_world["G"]
|
|
t_world = torso_world
|
|
|
|
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_2d.append(to_web_point(t_x, t_z, args.scale_cm, args.origin_x, args.origin_y, x_ref_m))
|
|
|
|
if args.frame == "body":
|
|
near_pts = world_to_body_local_points(data, near_world)
|
|
far_pts = world_to_body_local_points(data, far_world)
|
|
t_body_x, t_body_y, t_body_z = 0.0, 0.0, 0.0
|
|
torso_pt = {"x": t_body_x, "y": t_body_y, "z": t_body_z}
|
|
else:
|
|
near_pts = numpy_points_to_dict(near_world)
|
|
far_pts = numpy_points_to_dict(far_world)
|
|
torso_pt = {
|
|
"x": float(t_world[0]),
|
|
"y": float(t_world[1]),
|
|
"z": float(t_world[2]),
|
|
}
|
|
|
|
frames.append(
|
|
{
|
|
"t": float(data.time),
|
|
"near": near_pts,
|
|
"far": far_pts,
|
|
"torso": torso_pt,
|
|
}
|
|
)
|
|
next_sample += sample_dt
|
|
|
|
payload_v1 = {
|
|
"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_2d,
|
|
}
|
|
|
|
payload_v2 = {
|
|
"format": "walkersim-mujoco-playback-v2",
|
|
"units": "meters",
|
|
"frame": args.frame,
|
|
"sampleHz": float(args.sample_hz),
|
|
"durationSec": float(args.seconds),
|
|
"frames": frames,
|
|
"overlay2d": {
|
|
"origin": {"x": args.origin_x, "y": args.origin_y},
|
|
"nearF": near_f,
|
|
"nearG": near_g,
|
|
"farF": far_f,
|
|
"farG": far_g,
|
|
"torso": torso_2d,
|
|
},
|
|
}
|
|
|
|
payload = payload_v2 if args.format == "v2" else payload_v1
|
|
out_path.write_text(json.dumps(payload), encoding="utf-8")
|
|
print(f"Wrote MuJoCo web payload ({payload['format']}) to {out_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|