102 lines
3.9 KiB
Python
102 lines
3.9 KiB
Python
"""Worker-only contact replay on a retained episode's exact collision asset.
|
|
|
|
This isolates physical support from perception/planning. It never supplies a
|
|
route or collision truth to inference and never writes to the original run.
|
|
"""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import sys
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
from isaacsim import SimulationApp
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--run", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--mode", choices=("straight", "replay"), default="straight")
|
|
parser.add_argument("--seconds", type=float, default=40)
|
|
args = parser.parse_args()
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
|
|
|
|
app = SimulationApp({"headless": True, "hide_ui": True})
|
|
try:
|
|
import omni.timeline
|
|
import omni.usd
|
|
from isaacsim.core.experimental.prims import Articulation
|
|
from isaacsim.core.simulation_manager import SimulationManager
|
|
from pxr import UsdGeom
|
|
from rover_profile import PROFILE, DifferentialDrive, create_rover
|
|
from terrain import install_terrain
|
|
|
|
from k1link.simulation.ai_polygon.mission_policy import inclination
|
|
|
|
run = json.loads(args.run.read_text())
|
|
settings = run["world"]["settings"]
|
|
stage = omni.usd.get_context().get_stage()
|
|
UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
|
|
UsdGeom.SetStageMetersPerUnit(stage, 1)
|
|
height, terrain = install_terrain(stage, run["terrain_manifest"], run["world"])
|
|
create_rover(
|
|
stage,
|
|
[*settings["spawn_xy"], height],
|
|
settings["heading_degrees"],
|
|
terrain["initial_ground_normal"],
|
|
)
|
|
SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
|
|
omni.timeline.get_timeline_interface().play()
|
|
for _ in range(30):
|
|
app.update()
|
|
robot = Articulation("/World/Rover")
|
|
drive = DifferentialDrive(robot)
|
|
playback = []
|
|
if args.mode == "replay":
|
|
playback = [
|
|
json.loads(x) for x in (args.run.parent / "motion.jsonl").read_text().splitlines()
|
|
]
|
|
trace, cursor, stopped = [], 0, False
|
|
started = time.monotonic()
|
|
for step in range(round(args.seconds * 60)):
|
|
p, q = robot.get_world_poses()
|
|
p, q = p.numpy()[0].tolist(), q.numpy()[0].tolist()
|
|
pose = p + q[1:] + q[:1]
|
|
tilt = inclination(pose)
|
|
stopped |= tilt >= PROFILE["stop_tilt_degrees"]
|
|
v, w = settings["max_speed_mps"], 0.0
|
|
if playback:
|
|
while (
|
|
cursor + 1 < len(playback)
|
|
and playback[cursor + 1]["simulation_time_ns"] / 1e9 <= step / 60
|
|
):
|
|
cursor += 1
|
|
v, w = playback[cursor]["applied_speed_mps"], playback[cursor]["applied_yaw_rate_rps"]
|
|
drive.command(0 if stopped else v, 0 if stopped else w)
|
|
if step % 10 == 0:
|
|
trace.append(dict(time_s=step / 60, pose=pose, tilt_degrees=tilt, stopped=stopped))
|
|
SimulationManager.step(steps=1)
|
|
report = dict(
|
|
schema_version="missioncore.scene-contact-qualification/v1",
|
|
utc=datetime.now(UTC).isoformat(),
|
|
monotonic=time.monotonic(),
|
|
elapsed_wall_s=time.monotonic() - started,
|
|
run_sha256=hashlib.sha256(args.run.read_bytes()).hexdigest(),
|
|
mode=args.mode,
|
|
wheel_collision=PROFILE["wheel_collision"],
|
|
seconds=args.seconds,
|
|
profile=PROFILE,
|
|
terrain=terrain,
|
|
trace=trace,
|
|
displacement_m=math.dist(trace[0]["pose"][:2], trace[-1]["pose"][:2]),
|
|
max_tilt_degrees=max(x["tilt_degrees"] for x in trace),
|
|
stopped=stopped,
|
|
)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(report, indent=2), encoding="utf-8")
|
|
print(json.dumps({k: v for k, v in report.items() if k not in ("trace", "terrain", "profile")}))
|
|
finally:
|
|
app.close()
|