Files
NODEDC_MISSION_CORE/simulation/ai-polygon/run_realtime.py
T

567 lines
24 KiB
Python

"""Native Worker-local physics/render/camera loop, independent of Core and AI latency."""
import argparse
import json
import math
import shutil
import sys
import time
import traceback
from datetime import UTC, datetime
from pathlib import Path
from local_state import StateChannel, read_json, write_json
from motion_control import CONTROL_PROFILE, DriveEnvelope
from realtime_ai import LatestInference
from rover_profile import PROFILE, DifferentialDrive, create_rover
parser = argparse.ArgumentParser()
parser.add_argument("--run", type=Path, required=True)
parser.add_argument("--source", type=Path, required=True)
parser.add_argument("--stream-address", required=True)
args, _ = parser.parse_known_args()
root = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(root / "src"))
run = read_json(args.run)
settings = run["world"]["settings"]
directory = args.run.parent
channel = StateChannel(directory)
result = {"outcome": "failed", "message": "Симуляция не завершена."}
app = ai = None
physics_callbacks = []
def startup_phase(phase):
record = {"phase": phase, "utc": datetime.now(UTC).isoformat(), "monotonic": time.monotonic()}
write_json(directory / "startup.json", record)
print(json.dumps({"startup": record}), flush=True)
try:
startup_phase("native-runtime")
from isaacsim import SimulationApp
app = SimulationApp(
{
"headless": True,
"hide_ui": True,
"multi_gpu": False,
"width": 1280,
"height": 720,
"window_width": 1280,
"window_height": 720,
"renderer": "RaytracedLighting",
"display_options": 0,
"extra_args": [
"--/app/player/useFixedTimeStepping=false",
"--/app/runLoops/main/manualModeEnabled=false",
"--/exts/isaacsim.core.throttling/enable_manualmode=false",
],
}
)
import carb.settings
import numpy as np
import omni.kit.app
import omni.replicator.core as rep
import omni.timeline
import omni.usd
from isaacsim.core.experimental.prims import Articulation
from isaacsim.core.experimental.utils.app import enable_extension
from isaacsim.core.simulation_manager import SimulationEvent, SimulationManager
from navigation.footprint import FRAME_DEADLINE_SECONDS
from navigation_client import ComposedInference
from omni.kit.loop import _loop as omni_loop
from omni.kit.viewport.utility import get_active_viewport
from pxr import Gf, UsdGeom, UsdLux
from terrain import RangeSensor, install_terrain
from k1link.simulation.ai_polygon.mission_policy import inclination
startup_phase("stream-and-gaussian-runtime")
config = carb.settings.get_settings()
prefix = "/exts/omni.kit.livestream.app/primaryStream/"
for key, value in {
"publicIp": args.stream_address,
"signalPort": 49100,
"streamPort": 47998,
"targetFps": 30,
"enableEventTracing": False,
}.items():
config.set(prefix + key, value)
enable_extension("omni.kit.livestream.app")
enable_extension("omni.kit.converter.gsplat")
from usd_convert_gsplat import read_ply, write_gaussian_splat_usd
converted = args.source.with_suffix(".usd")
if not converted.exists():
temporary = converted.with_suffix(".part.usd")
write_gaussian_splat_usd(
read_ply(str(args.source)),
str(temporary),
source_file=str(args.source),
prim_name="Gaussians",
up_axis="Z",
)
temporary.replace(converted)
stage = omni.usd.get_context().get_stage()
UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
UsdGeom.SetStageMetersPerUnit(stage, 1)
visual = UsdGeom.Xform.Define(stage, "/World/Scan")
visual.AddRotateXYZOp().Set(Gf.Vec3f(*settings["rotation_degrees"]))
visual.AddScaleOp().Set(Gf.Vec3f(settings["meters_per_unit"]))
stage.DefinePrim("/World/Scan/Gaussians").GetReferences().AddReference(str(converted))
startup_phase("terrain")
wheel_center_z, terrain_record = install_terrain(stage, run["terrain_manifest"], run["world"])
write_json(directory / "terrain.json", terrain_record)
UsdLux.DomeLight.Define(stage, "/World/Light").CreateIntensityAttr(500)
create_rover(
stage,
[*settings["spawn_xy"], wheel_center_z],
settings["heading_degrees"],
terrain_record["initial_ground_normal"],
)
write_json(directory / "robot-profile.json", PROFILE)
write_json(directory / "motion-control.json", CONTROL_PROFILE)
startup_phase("cameras")
def camera(path, aspect, focal_length=24):
value = UsdGeom.Camera.Define(stage, path)
value.CreateFocalLengthAttr(focal_length)
value.CreateHorizontalApertureAttr(36)
value.CreateVerticalApertureAttr(36 / aspect)
value.CreateClippingRangeAttr(Gf.Vec2f(0.05, 1000))
return value, value.AddTranslateOp(), value.AddOrientOp()
sensor, sensor_pos, sensor_rot = camera(
"/World/Sensor", 4 / 3, PROFILE["camera_focal_length_mm"]
)
focal_pixels = (
800 * PROFILE["camera_focal_length_mm"] / PROFILE["camera_horizontal_aperture_mm"]
)
observer, observer_pos, observer_rot = camera("/World/Observer", 16 / 9)
product = rep.create.render_product(str(sensor.GetPath()), (800, 600))
annotator = rep.AnnotatorRegistry.get_annotator("rgb")
annotator.attach(product)
config.set("/omni/replicator/captureOnPlay", True)
viewport = get_active_viewport()
if viewport is None:
raise RuntimeError("Streaming viewport is unavailable")
viewport.camera_path = str(observer.GetPath())
viewport.set_texture_resolution((1280, 720))
# Initialise the final camera before the first rendered frames. Starting
# WebRTC with an uninitialised observer and an already paused timeline can
# leave the native stream without its first image/offer.
initial = Gf.Vec3d(*settings["spawn_xy"], wheel_center_z + 0.12)
heading = math.radians(settings["heading_degrees"])
initial_eye = initial + Gf.Vec3d(-math.cos(heading) * 2.5, -math.sin(heading) * 2.5, 1.5)
observer_pos.Set(initial_eye)
observer_rot.Set(
Gf.Quatf(
Gf.Matrix4d()
.SetLookAt(initial_eye, initial + Gf.Vec3d(0.1, 0, 0.2), Gf.Vec3d(0, 0, 1))
.GetInverse()
.ExtractRotationQuat()
)
)
startup_phase("physics-setup")
SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
timeline = omni.timeline.get_timeline_interface()
# Render rate and physics rate are separate. Physics steps follow elapsed time.
config.set("/app/player/useFixedTimeStepping", False)
config.set("/exts/isaacsim.core.throttling/enable_manualmode", False)
omni_loop.acquire_loop_interface().set_manual_mode(False)
timeline.set_play_every_frame(False)
config.set("/persistent/simulation/minFrameRate", 1)
config.set("/app/runLoops/main/rateLimitEnabled", True)
config.set("/app/runLoops/main/rateLimitFrequency", 30)
startup_phase("initial-render-and-contact")
timeline.play()
for _ in range(30):
app.update()
articulation = Articulation("/World/Rover")
drive = DifferentialDrive(articulation)
wheel_indices = articulation.get_dof_indices(PROFILE["wheel_names"])
def body_transform():
positions, orientations = articulation.get_world_poses()
position, orientation = positions.numpy()[0], orientations.numpy()[0]
transform = Gf.Matrix4d(1)
transform.SetRotate(Gf.Quatd(float(orientation[0]), Gf.Vec3d(*map(float, orientation[1:]))))
transform.SetTranslateOnly(Gf.Vec3d(*map(float, position)))
return transform
settled = body_transform().ExtractTranslation()
write_json(
directory / "initial-contact.json",
{
"pose": list(settled),
"expected_xy": settings["spawn_xy"],
"spawn_contact_z": terrain_record["initial_contact_z"],
},
)
if math.hypot(settled[0] - settings["spawn_xy"][0], settled[1] - settings["spawn_xy"][1]) > 0.2:
raise RuntimeError("Rover start is unstable on this reconstructed surface")
timeline.pause()
app.update()
startup_phase("ready")
baseline = SimulationManager.get_num_physics_steps()
range_sensor = RangeSensor()
from k1link.simulation.ai_polygon.mission_policy import WaypointMission
# The episode owns the mission. Reconnecting a failed model client must not
# rewind its route cursor, recovery budget or latched terminal condition.
mission = WaypointMission(settings.get("route_xy", []))
ai = LatestInference(
lambda: ComposedInference(
root / "simulation/ai-polygon", run, directory / "camera", mission=mission
),
None,
directory / "camera",
)
ai.start()
control = {"control": "pause", "control_sequence": 0, "camera": "follow"}
start = last_report = last_sensor = last_control = time.monotonic()
render_frames = sensor_frames = sequence = 0
prior_metrics = (start, 0, 0, 0, baseline)
speed = 0.0
had_ai = False
unstable = False
terminal_at = None
envelope = DriveEnvelope(track_width=PROFILE["track_width_m"])
state = {"velocity": 0.0, "yaw_rate": 0.0, "stop_reason": "paused", "decision": None}
physics_rows = []
def control_step(dt, _context):
velocity, yaw_rate, reason, decision = ai.command(
time.monotonic(), frame_deadline=FRAME_DEADLINE_SECONDS
)
terminal = unstable or (
decision and decision["decision"]["reason"] in {"unstable", "stuck", "goal-reached"}
)
velocity, yaw_rate = envelope.step(
velocity, yaw_rate, dt, stop=bool(terminal) or reason != "none"
)
drive.command(velocity, yaw_rate)
state.update(
velocity=velocity,
yaw_rate=yaw_rate,
stop_reason="unstable" if unstable else reason,
decision=decision,
)
def update_cameras():
transform = body_transform()
position = transform.ExtractTranslation()
forward = transform.TransformDir(Gf.Vec3d(1, 0, 0)).GetNormalized()
yaw = math.atan2(forward[1], forward[0])
# Both sensor position and attitude follow the articulated body.
# Mounting height is measured from nominal wheel contact.
eye = transform.Transform(
Gf.Vec3d(
PROFILE["camera_forward_m"],
0,
settings["camera_height_m"] - PROFILE["body_contact_height_m"],
)
)
pitch = math.radians(PROFILE["camera_pitch_degrees"])
camera_axes = [
transform.TransformDir(Gf.Vec3d(*axis)).GetNormalized()
for axis in (
(math.cos(pitch), 0, math.sin(pitch)),
(0, 1, 0),
(-math.sin(pitch), 0, math.cos(pitch)),
)
]
def look_at(pos, rotation, origin, target, up=None):
if up is None:
up = Gf.Vec3d(0, 0, 1)
pos.Set(origin)
rotation.Set(
Gf.Quatf(
Gf.Matrix4d().SetLookAt(origin, target, up).GetInverse().ExtractRotationQuat()
)
)
look_at(
sensor_pos,
sensor_rot,
eye,
eye + camera_axes[0],
camera_axes[2],
)
mode = control.get("camera", "follow")
if mode == "camera":
viewport.camera_path = str(sensor.GetPath())
else:
viewport.camera_path = str(observer.GetPath())
offset = (
Gf.Vec3d(-math.cos(yaw) * 2.5, -math.sin(yaw) * 2.5, 1.5)
if mode == "follow"
else Gf.Vec3d(-3, -3, 5)
)
look_at(observer_pos, observer_rot, position + offset, position + Gf.Vec3d(0.1, 0, 0.2))
linear, angular = articulation.get_velocities()
q = transform.ExtractRotationQuat()
state.update(
transform=transform,
position=position,
yaw=yaw,
eye=eye,
camera_axes=camera_axes,
mode=mode,
pose=list(position) + list(q.GetImaginary()) + [q.GetReal()],
speed=float(np.linalg.norm(linear.numpy()[0, :2])),
body_velocity=linear.numpy()[0].tolist(),
body_angular_velocity=angular.numpy()[0].tolist(),
)
def physics_step(dt, _context):
global unstable
update_cameras()
unstable |= inclination(state["pose"]) >= PROFILE["stop_tilt_degrees"]
physics_rows.append(
{
"monotonic": time.monotonic(),
"dt": dt,
"pose": state["pose"],
"velocity": state["body_velocity"],
"angular_velocity": state["body_angular_velocity"],
"command": [state["velocity"], state["yaw_rate"]],
"stop_reason": state["stop_reason"],
}
)
callback_errors = []
def guarded(callback):
def invoke(dt, context):
try:
callback(dt, context)
except Exception:
# Native event dispatch logs Python exceptions and continues.
# A failed control/camera callback must instead end this run.
callback_errors.append(traceback.format_exc())
ai.enable(False)
drive.command(0.0, 0.0)
return invoke
update_cameras()
physics_callbacks.extend(
[
SimulationManager.register_callback(
guarded(control_step), SimulationEvent.PHYSICS_PRE_STEP
),
SimulationManager.register_callback(
guarded(physics_step), SimulationEvent.PHYSICS_POST_STEP
),
]
)
with (
(directory / "motion.jsonl").open("a", encoding="utf-8") as motion,
(directory / "physics-motion.jsonl").open("a", encoding="utf-8") as physics_motion,
):
while app.is_running():
now = time.monotonic()
if now - last_control >= 0.1:
control = channel.read("control", control)
last_control = now
if control["control"] == "stop":
result = {"outcome": "stopped", "message": "Симуляция и inference остановлены."}
break
if now - start >= run["request"]["duration_seconds"]:
result = {"outcome": "completed", "message": "Время прогона завершено."}
break
playing = control["control"] == "play"
ai.enable(playing)
had_ai |= playing
if playing and not timeline.is_playing():
timeline.play()
elif not playing and timeline.is_playing():
timeline.pause()
velocity, yaw_rate, stop_reason, decision = ai.command(
now, frame_deadline=FRAME_DEADLINE_SECONDS
)
transform = body_transform()
q = transform.ExtractRotationQuat()
pose = list(transform.ExtractTranslation()) + list(q.GetImaginary()) + [q.GetReal()]
unstable |= inclination(pose) >= PROFILE["stop_tilt_degrees"]
if unstable:
velocity, yaw_rate, stop_reason = 0.0, 0.0, "unstable"
terminal_reason = (
"unstable" if unstable else (decision["decision"]["reason"] if decision else None)
)
if terminal_reason in {"unstable", "stuck", "goal-reached"}:
velocity, yaw_rate = 0.0, 0.0
terminal_at = now if terminal_at is None else terminal_at
if now - terminal_at >= 2:
result = {
"outcome": "completed" if terminal_reason == "goal-reached" else "failed",
"message": {
"goal-reached": "Ровер достиг цели маршрута.",
"stuck": "Ровер не нашёл проезд после трёх попыток. Прогон завершён.",
"unstable": "Прогон завершён из-за опасного наклона ровера.",
}[terminal_reason],
}
break
# The physics tensor API applies live drive targets in radians/s;
# USD authoring attributes are only the initial scene configuration.
if not playing:
drive.command(*envelope.step(0.0, 0.0, 0.0, stop=True))
update_cameras()
state.update(
velocity=0.0, yaw_rate=0.0, stop_reason="paused", speed=0.0, decision=decision
)
app.update() # Never waits for inference, a network request, or an operator ACK.
if callback_errors:
raise RuntimeError(callback_errors[0])
render_frames += 1
captured_at = time.monotonic()
transform, position = state["transform"], state["position"]
eye, camera_axes = state["eye"], state["camera_axes"]
pose, yaw, speed, mode = state["pose"], state["yaw"], state["speed"], state["mode"]
velocity, yaw_rate = state["velocity"], state["yaw_rate"]
stop_reason, decision = state["stop_reason"], state["decision"]
if physics_rows:
physics_motion.writelines(
json.dumps(row, allow_nan=False) + "\n" for row in physics_rows
)
physics_rows.clear()
physics_motion.flush()
physics = SimulationManager.get_num_physics_steps()
sim_ns = round((physics - baseline) * 1e9 / 60)
if playing and captured_at - last_sensor >= 1 / 5:
rgb = annotator.get_data()
if isinstance(rgb, np.ndarray) and rgb.shape[:2] == (600, 800):
sensor_frames += 1
rear_eye = transform.Transform(
Gf.Vec3d(
PROFILE["rear_range_forward_m"],
0,
settings["camera_height_m"] - PROFILE["body_contact_height_m"],
)
)
points = range_sensor.capture(eye, transform, rear_eye)
quaternion = transform.ExtractRotationQuat()
imaginary = quaternion.GetImaginary()
# Columns are the camera's forward/left/up axes in world.
observation = {
"simulation_time_ns": sim_ns,
"points": points,
"pose": [float(v) for v in position]
+ [float(v) for v in imaginary]
+ [float(quaternion.GetReal())],
"calibration": {
"origin": list(eye),
"rotation": np.array(camera_axes).T.reshape(-1).tolist(),
"intrinsics": [focal_pixels, focal_pixels, 400.0, 300.0],
"body_contact_height_m": PROFILE["body_contact_height_m"],
"range_origins": [list(eye), list(rear_eye)],
},
}
ai.submit(
np.array(rgb[:, :, :3], copy=True, order="C"),
sensor_frames,
captured_at,
sim_ns,
observation,
)
last_sensor = captured_at
if captured_at - last_report >= 0.25:
prior_at, prior_render, prior_sensor, prior_ai, prior_physics = prior_metrics
elapsed = captured_at - prior_at
report = dict(
sequence=sequence,
control_sequence=control["control_sequence"],
state="running" if playing else "paused" if had_ai else "ready",
phase="models" if playing and not ai.ever_ready else "running",
simulation_time_ns=sim_ns,
wall_elapsed_seconds=captured_at - start,
physics_steps=physics - baseline,
render_frames=render_frames,
sensor_frames=sensor_frames,
inference_count=ai.count,
dropped_frames=ai.dropped,
rtf=(physics - prior_physics) / 60 / elapsed,
render_fps=(render_frames - prior_render) / elapsed,
sensor_fps=(sensor_frames - prior_sensor) / elapsed,
ai_hz=(ai.count - prior_ai) / elapsed,
inference_ms=decision["inference_ms"] if decision else None,
frame_age_ms=(captured_at - decision["captured_at"]) * 1000
if decision
else None,
command_age_ms=(captured_at - decision["completed_at"]) * 1000
if decision
else None,
pose_xy=[float(position[0]), float(position[1])],
pose_yaw=yaw,
speed_mps=speed,
applied_speed_mps=velocity,
applied_yaw_rate_rps=yaw_rate,
decision=decision["decision"] if decision else None,
stop_reason=stop_reason,
ai_ready=ai.ready,
stream_ready=True,
camera=mode,
)
channel.write("snapshot", report)
motion.write(
json.dumps(
{
**report,
"monotonic": captured_at,
"pose_xyz_xyzw": pose,
"tilt_degrees": inclination(pose),
"wheel_velocity_rps": articulation.get_dof_velocities(
dof_indices=wheel_indices
)
.numpy()
.tolist()[0],
"usd_position_xyz": list(
UsdGeom.XformCache()
.GetLocalToWorldTransform(
stage.GetPrimAtPath("/World/Rover/chassis")
)
.ExtractTranslation()
),
},
allow_nan=False,
)
+ "\n"
)
motion.flush()
last_report = captured_at
sequence += 1
if elapsed >= 1:
prior_metrics = (captured_at, render_frames, sensor_frames, ai.count, physics)
if shutil.disk_usage(directory).free < 1024**3:
raise RuntimeError("Worker evidence disk is full")
except Exception as exc:
message = str(exc)
placement_failure = any(
text in message
for text in (
"Configured spawn",
"Rover start",
"reconstructed ground",
"reconstructed support",
)
)
result = {
"outcome": "failed",
"message": "Не удалось устойчиво разместить ровер на грунте. Измените положение старта."
if placement_failure
else "Симуляция прервана из-за ошибки на Worker. Журнал сохранён.",
}
traceback.print_exc()
finally:
for callback_id in physics_callbacks:
SimulationManager.deregister_callback(callback_id)
channel.close()
if ai is not None:
ai.close()
write_json(directory / "result.json", result)
if app is not None:
app.close()