207 lines
9.2 KiB
Python
207 lines
9.2 KiB
Python
"""One native Isaac episode. Receives pixels from physics; emits virtual commands only."""
|
|
|
|
import argparse
|
|
import base64
|
|
import io
|
|
import json
|
|
import math
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from core_client import CoreClient
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--run", type=Path, required=True)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--core", required=True)
|
|
parser.add_argument("--token-file", type=Path, required=True)
|
|
parser.add_argument("--instance", required=True)
|
|
args = parser.parse_args()
|
|
root = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(root / "src"))
|
|
run = json.loads(args.run.read_text())
|
|
settings = run["world"]["settings"]
|
|
client = CoreClient(args.core, args.token_file, args.instance)
|
|
result = {"outcome": "failed", "message": "Симуляция не завершена."}
|
|
app = model = robot = controller = None
|
|
try:
|
|
from isaacsim import SimulationApp
|
|
|
|
app = SimulationApp({"headless": True, "multi_gpu": False, "width": 800, "height": 600})
|
|
import numpy as np
|
|
|
|
# Isaac bundles the converter compatible with its own USD/ParticleField schema.
|
|
import omni.kit.app
|
|
import omni.replicator.core as rep
|
|
import omni.usd
|
|
from isaacsim.core.experimental.utils import app as app_utils
|
|
from isaacsim.core.simulation_manager import SimulationManager
|
|
from isaacsim.robot.experimental.wheeled_robots.controllers import DifferentialController
|
|
from isaacsim.robot.experimental.wheeled_robots.robots import WheeledRobot
|
|
from PIL import Image
|
|
from pxr import Gf, UsdGeom, UsdLux, UsdPhysics
|
|
|
|
from k1link.simulation.ai_polygon.inference import ModelInference
|
|
from k1link.simulation.ai_polygon.policy import RoadPolicy
|
|
|
|
omni.kit.app.get_app().get_extension_manager().set_extension_enabled_immediate(
|
|
"omni.kit.converter.gsplat", True
|
|
)
|
|
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))
|
|
# Only a hidden rigid ground plane. Vegetation in the scan is visual-only.
|
|
ground = UsdGeom.Plane.Define(stage, "/World/Ground")
|
|
ground.CreateAxisAttr("Z")
|
|
ground.CreateWidthAttr(20000)
|
|
ground.CreateLengthAttr(20000)
|
|
ground.AddTranslateOp().Set(Gf.Vec3d(0, 0, settings["ground_z"]))
|
|
UsdPhysics.CollisionAPI.Apply(ground.GetPrim())
|
|
ground.CreateVisibilityAttr(UsdGeom.Tokens.invisible)
|
|
UsdLux.DomeLight.Define(stage, "/World/Light").CreateIntensityAttr(500)
|
|
robot_asset = Path("D:/NDC_MISSIONCORE/runtime/simulation/assets/jetbot-6.1-v1/jetbot.usda")
|
|
if not robot_asset.is_file():
|
|
raise RuntimeError("Prepare the pinned stock Jetbot asset cache first")
|
|
heading = math.radians(settings["heading_degrees"])
|
|
robot = WheeledRobot(
|
|
paths="/World/Rover",
|
|
wheel_dof_names=["left_wheel_joint", "right_wheel_joint"],
|
|
usd_path=str(robot_asset),
|
|
positions=[*settings["spawn_xy"], settings["ground_z"] + 0.05],
|
|
orientations=[math.cos(heading / 2), 0, 0, math.sin(heading / 2)],
|
|
)
|
|
controller = DifferentialController(wheel_radius=0.03, wheel_base=0.1125)
|
|
camera = UsdGeom.Camera.Define(stage, "/World/Camera")
|
|
camera.CreateFocalLengthAttr(24)
|
|
camera.CreateHorizontalApertureAttr(36)
|
|
camera.CreateVerticalApertureAttr(27)
|
|
camera.CreateClippingRangeAttr(Gf.Vec2f(0.02, 1000))
|
|
camera_pos = camera.AddTranslateOp()
|
|
camera_rot = camera.AddOrientOp()
|
|
render_product = rep.create.render_product(str(camera.GetPath()), (800, 600))
|
|
annotator = rep.AnnotatorRegistry.get_annotator("rgb")
|
|
annotator.attach(render_product)
|
|
SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
|
|
app_utils.play()
|
|
app_utils.update_app(steps=20)
|
|
app_utils.pause()
|
|
baseline = SimulationManager.get_num_physics_steps()
|
|
profile = json.loads((root / "simulation/ai-polygon/models.worker-006.json").read_text())
|
|
model = ModelInference(
|
|
"http://127.0.0.1:18092", Path(profile["labels"]), "http://127.0.0.1:18091"
|
|
)
|
|
model.ready()
|
|
policy = RoadPolicy(settings["max_speed_mps"])
|
|
|
|
def position():
|
|
positions, rotations = robot.get_world_poses()
|
|
xyz, q = positions.numpy()[0], rotations.numpy()[0]
|
|
yaw = math.atan2(2 * (q[0] * q[3] + q[1] * q[2]), 1 - 2 * (q[2] ** 2 + q[3] ** 2))
|
|
return xyz, yaw
|
|
|
|
sequence = 0
|
|
last_cycle_ms = 0.0
|
|
while sequence < run["request"]["max_steps"]:
|
|
cycle_started = time.monotonic()
|
|
transport_ms = 0.0
|
|
transport_started = time.monotonic()
|
|
action = client.request(
|
|
"/worker/poll", {"instance_id": args.instance, "run_id": run["run_id"]}
|
|
)["action"]
|
|
transport_ms += (time.monotonic() - transport_started) * 1000
|
|
if action == "stop":
|
|
result = {"outcome": "stopped", "message": "Прогон остановлен."}
|
|
break
|
|
if action == "pause":
|
|
robot.apply_wheel_actions(controller.forward([0, 0]))
|
|
time.sleep(0.1)
|
|
continue
|
|
if action not in ("play", "step"):
|
|
raise RuntimeError("Core withdrew simulation ownership")
|
|
xyz, yaw = position()
|
|
# Pose is used solely to place the virtual sensor, never passed to policy.
|
|
eye = Gf.Vec3d(float(xyz[0]), float(xyz[1]), float(xyz[2] + settings["camera_height_m"]))
|
|
camera_pos.Set(eye)
|
|
# USD cameras look along local -Z with +Y up. Invert a world-to-camera
|
|
# look-at transform to keep the optical horizon level for every heading.
|
|
view = Gf.Matrix4d().SetLookAt(
|
|
eye, eye + Gf.Vec3d(math.cos(yaw), math.sin(yaw), 0), Gf.Vec3d(0, 0, 1)
|
|
)
|
|
camera_rot.Set(Gf.Quatf(view.GetInverse().ExtractRotationQuat()))
|
|
before = SimulationManager.get_num_physics_steps()
|
|
render_started = time.monotonic()
|
|
rep.orchestrator.step(rt_subframes=1, delta_time=0.0, pause_timeline=True)
|
|
if SimulationManager.get_num_physics_steps() != before:
|
|
raise RuntimeError("Rendering advanced the simulation clock")
|
|
render_ms = (time.monotonic() - render_started) * 1000
|
|
rgb = np.ascontiguousarray(annotator.get_data()[:, :, :3])
|
|
started = time.monotonic_ns()
|
|
road, boxes = model.infer(rgb)
|
|
decision = policy.decide(road, boxes)
|
|
inference_ms = (time.monotonic_ns() - started) / 1e6
|
|
image = io.BytesIO()
|
|
Image.fromarray(rgb).save(image, "JPEG", quality=85)
|
|
transport_started = time.monotonic()
|
|
client.request(
|
|
"/worker/runs/" + run["run_id"] + "/samples",
|
|
{
|
|
"sequence": sequence,
|
|
"simulation_time_ns": sequence * run["step_ns"],
|
|
"inference_ms": inference_ms,
|
|
"pose_xy": [float(xyz[0]), float(xyz[1])],
|
|
"decision": decision.model_dump(),
|
|
"image_jpeg_base64": base64.b64encode(image.getvalue()).decode(),
|
|
},
|
|
)
|
|
transport_ms += (time.monotonic() - transport_started) * 1000
|
|
robot.apply_wheel_actions(controller.forward([decision.speed_mps, decision.yaw_rate_rps]))
|
|
SimulationManager.step(steps=6)
|
|
if SimulationManager.get_num_physics_steps() != baseline + (sequence + 1) * 6:
|
|
raise RuntimeError("Physics clock left lockstep")
|
|
xyz, yaw = position()
|
|
transport_started = time.monotonic()
|
|
client.request(
|
|
"/worker/runs/" + run["run_id"] + "/applied",
|
|
{
|
|
"sequence": sequence,
|
|
"simulation_time_ns": (sequence + 1) * run["step_ns"],
|
|
"physics_steps": 6,
|
|
"pose_yaw": float(yaw),
|
|
"cycle_ms": last_cycle_ms or (time.monotonic() - cycle_started) * 1000,
|
|
"render_ms": render_ms,
|
|
"transport_ms": transport_ms,
|
|
"pose_xy": [float(xyz[0]), float(xyz[1])],
|
|
},
|
|
)
|
|
last_cycle_ms = (time.monotonic() - cycle_started) * 1000
|
|
sequence += 1
|
|
else:
|
|
result = {"outcome": "completed", "message": "Прогон завершён."}
|
|
except Exception as exc:
|
|
result = {"outcome": "failed", "message": "Прогон прерван. Подробности сохранены на Worker."}
|
|
print(type(exc).__name__ + ": " + str(exc), file=sys.stderr, flush=True)
|
|
finally:
|
|
if model is not None:
|
|
model.close()
|
|
args.run.with_name("result.json").write_text(json.dumps(result))
|
|
if app is not None:
|
|
app.close()
|