203 lines
9.2 KiB
Python
203 lines
9.2 KiB
Python
"""Metric collision proxy and occlusion-aware virtual range sensor on Worker."""
|
|
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import struct
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def load_glb(path):
|
|
data = Path(path).read_bytes()
|
|
magic, version, size = struct.unpack_from("<III", data)
|
|
if (magic, version, size) != (0x46546C67, 2, len(data)):
|
|
raise ValueError("Expected a complete glTF 2 binary collision mesh")
|
|
offset, document, binary = 12, None, None
|
|
while offset < len(data):
|
|
length, kind = struct.unpack_from("<II", data, offset)
|
|
chunk = data[offset + 8 : offset + 8 + length]
|
|
if kind == 0x4E4F534A:
|
|
document = json.loads(chunk)
|
|
elif kind == 0x004E4942:
|
|
binary = chunk
|
|
offset += 8 + length
|
|
if document is None or binary is None:
|
|
raise ValueError("Missing mesh buffers")
|
|
|
|
def accessor(index):
|
|
spec = document["accessors"][index]
|
|
view = document["bufferViews"][spec["bufferView"]]
|
|
if "sparse" in spec or view.get("buffer", 0) != 0:
|
|
raise ValueError("Unsupported sparse/external collision buffer")
|
|
dtype = {5123: "<u2", 5125: "<u4", 5126: "<f4"}[spec["componentType"]]
|
|
width = {"SCALAR": 1, "VEC3": 3}[spec["type"]]
|
|
item = np.dtype(dtype).itemsize
|
|
start = view.get("byteOffset", 0) + spec.get("byteOffset", 0)
|
|
return np.ndarray(
|
|
(spec["count"], width),
|
|
dtype=dtype,
|
|
buffer=binary,
|
|
offset=start,
|
|
strides=(view.get("byteStride", width * item), item),
|
|
).copy()
|
|
|
|
vertices, faces, base = [], [], 0
|
|
# SplatTransform collision exports bake coordinates; reject transforms so
|
|
# we never silently misregister contact geometry against the Gaussian view.
|
|
for node in document.get("nodes", []):
|
|
if any(k in node for k in ("matrix", "translation", "rotation", "scale")):
|
|
raise ValueError("Collision node transform must be baked")
|
|
for mesh in document["meshes"]:
|
|
for primitive in mesh["primitives"]:
|
|
if primitive.get("mode", 4) != 4:
|
|
raise ValueError("Collision mesh must contain triangles")
|
|
points = accessor(primitive["attributes"]["POSITION"])
|
|
triangles = accessor(primitive["indices"]).reshape(-1, 3).astype(np.int32)
|
|
if not np.isfinite(points).all() or triangles.max() >= len(points):
|
|
raise ValueError("Invalid collision vertices")
|
|
vertices.append(points[:, [0, 2, 1]] * np.array([1, -1, 1], np.float32))
|
|
faces.append(triangles + base)
|
|
base += len(points)
|
|
return np.concatenate(vertices), np.concatenate(faces)
|
|
|
|
|
|
def ground_intersections(vertices, faces, x, y):
|
|
"""Vertical mesh intersections, used only to place the initial rigid body."""
|
|
triangles = vertices[faces]
|
|
low, high = triangles[:, :, :2].min(axis=1), triangles[:, :, :2].max(axis=1)
|
|
candidate = triangles[
|
|
(low[:, 0] <= x) & (high[:, 0] >= x) & (low[:, 1] <= y) & (high[:, 1] >= y)
|
|
]
|
|
if not len(candidate):
|
|
return np.empty(0)
|
|
a, b, c = candidate[:, 0], candidate[:, 1], candidate[:, 2]
|
|
den = (b[:, 1] - c[:, 1]) * (a[:, 0] - c[:, 0]) + (c[:, 0] - b[:, 0]) * (a[:, 1] - c[:, 1])
|
|
valid = np.abs(den) > 1e-8
|
|
a, b, c, den = a[valid], b[valid], c[valid], den[valid]
|
|
u = ((b[:, 1] - c[:, 1]) * (x - c[:, 0]) + (c[:, 0] - b[:, 0]) * (y - c[:, 1])) / den
|
|
v = ((c[:, 1] - a[:, 1]) * (x - c[:, 0]) + (a[:, 0] - c[:, 0]) * (y - c[:, 1])) / den
|
|
inside = (u >= -1e-6) & (v >= -1e-6) & (u + v <= 1 + 1e-6)
|
|
return (u * a[:, 2] + v * b[:, 2] + (1 - u - v) * c[:, 2])[inside]
|
|
|
|
|
|
def install_terrain(stage, manifest_path, world):
|
|
from pxr import UsdGeom, UsdPhysics, UsdShade
|
|
from rover_profile import PROFILE
|
|
|
|
from k1link.simulation.ai_polygon.terrain_contract import terrain_matches
|
|
|
|
manifest = json.loads(Path(manifest_path).read_text(encoding="utf-8-sig"))
|
|
if not terrain_matches(manifest, world):
|
|
raise ValueError("Terrain does not match the world calibration")
|
|
path = Path(manifest["collider"])
|
|
if hashlib.sha256(path.read_bytes()).hexdigest() != manifest["collider_sha256"]:
|
|
raise ValueError("Collision mesh identity changed")
|
|
points, triangles = load_glb(path)
|
|
settings = world["settings"]
|
|
heights = ground_intersections(points, triangles, *settings["spawn_xy"])
|
|
candidates = heights[np.abs(heights - settings["ground_z"]) < 0.8]
|
|
if not len(candidates):
|
|
raise ValueError("No reconstructed ground at the configured spawn")
|
|
# Place the whole footprint above nearby ground, without embedding a wheel
|
|
# in a stone beside the centre ray. This is setup, never a planner input.
|
|
angle = math.radians(settings["heading_degrees"])
|
|
support, offsets = [], []
|
|
for dx in (-0.5, 0, 0.5):
|
|
for dy in (-0.5, 0, 0.5):
|
|
x = settings["spawn_xy"][0] + dx * math.cos(angle) - dy * math.sin(angle)
|
|
y = settings["spawn_xy"][1] + dx * math.sin(angle) + dy * math.cos(angle)
|
|
intersections = ground_intersections(points, triangles, x, y)
|
|
nearby = intersections[np.abs(intersections - settings["ground_z"]) < 0.8]
|
|
if not len(nearby):
|
|
raise ValueError("Rover footprint has no reconstructed support")
|
|
support.append(float(nearby.max()))
|
|
offsets.append([x - settings["spawn_xy"][0], y - settings["spawn_xy"][1], 1])
|
|
plane = np.linalg.lstsq(np.asarray(offsets), np.asarray(support), rcond=None)[0]
|
|
residual = np.asarray(support) - np.asarray(offsets) @ plane
|
|
normal = np.array([-plane[0], -plane[1], 1.0])
|
|
normal /= np.linalg.norm(normal)
|
|
if np.max(np.abs(residual)) > 0.12 or normal[2] < math.cos(math.radians(25)):
|
|
raise ValueError("Configured spawn is too uneven for the metre-wide rover")
|
|
from spawn_clearance import obstructing_triangles
|
|
|
|
if obstructing_triangles(
|
|
points,
|
|
triangles,
|
|
settings["spawn_xy"],
|
|
settings["heading_degrees"],
|
|
plane,
|
|
step=PROFILE["max_step_m"],
|
|
):
|
|
raise ValueError("Configured spawn contains an obstacle inside the rover footprint")
|
|
# Align to the local support plane before gravity settles the suspensionless
|
|
# lab chassis. Wheel centres start 4 cm above the highest residual contact.
|
|
wheel_center_z = float(
|
|
plane[2] + (PROFILE["wheel_radius_m"] + 0.04) / normal[2] + max(0, residual.max())
|
|
)
|
|
mesh = UsdGeom.Mesh.Define(stage, "/World/Terrain")
|
|
mesh.CreatePointsAttr(points.tolist())
|
|
mesh.CreateFaceVertexCountsAttr([3] * len(triangles))
|
|
mesh.CreateFaceVertexIndicesAttr(triangles.reshape(-1).tolist())
|
|
mesh.CreateSubdivisionSchemeAttr(UsdGeom.Tokens.none)
|
|
mesh.CreateDoubleSidedAttr(True)
|
|
mesh.CreateVisibilityAttr(UsdGeom.Tokens.invisible)
|
|
UsdPhysics.CollisionAPI.Apply(mesh.GetPrim())
|
|
UsdPhysics.MeshCollisionAPI.Apply(mesh.GetPrim()).CreateApproximationAttr("none")
|
|
material = UsdShade.Material.Define(stage, "/World/Materials/Terrain")
|
|
surface = UsdPhysics.MaterialAPI.Apply(material.GetPrim())
|
|
surface.CreateStaticFrictionAttr(0.9)
|
|
surface.CreateDynamicFrictionAttr(0.8)
|
|
surface.CreateRestitutionAttr(0)
|
|
UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim()).Bind(
|
|
material, UsdShade.Tokens.weakerThanDescendants, "physics"
|
|
)
|
|
return wheel_center_z, dict(
|
|
manifest,
|
|
vertex_count=len(points),
|
|
triangle_count=len(triangles),
|
|
initial_contact_z=float(plane[2]),
|
|
initial_ground_normal=normal.tolist(),
|
|
support_residual_m=float(np.max(np.abs(residual))),
|
|
)
|
|
|
|
|
|
class RangeSensor:
|
|
"""Front 360-degree LiDAR plus rear near-field fan; first physical hit.
|
|
|
|
Separate mounts observe the ground beyond each bumper without seeing
|
|
through the chassis. The rear fan adds 660 rays to the 2160 front rays.
|
|
"""
|
|
|
|
def __init__(self):
|
|
import omni.physx
|
|
|
|
self.query = omni.physx.get_physx_scene_query_interface()
|
|
elevations = (-80, -75, -70, -60, -45, -35, -28, -22, -18, -14, -10, -5, 0, 5, 15)
|
|
self.directions = [
|
|
(math.cos(e) * math.cos(a), math.cos(e) * math.sin(a), math.sin(e))
|
|
for e in np.radians(elevations)
|
|
for a in np.radians(np.arange(-180, 180, 2.5))
|
|
]
|
|
self.rear_directions = [
|
|
(math.cos(e) * math.cos(a), math.cos(e) * math.sin(a), math.sin(e))
|
|
for e in np.radians(np.arange(-80, 16, 5))
|
|
for a in np.radians(np.arange(100, 261, 5))
|
|
]
|
|
|
|
def capture(self, origin, rotation, rear_origin=None):
|
|
from pxr import Gf
|
|
|
|
points = []
|
|
mounts = [(origin, self.directions)]
|
|
if rear_origin is not None:
|
|
mounts.append((rear_origin, self.rear_directions))
|
|
for mount, directions in mounts:
|
|
for direction in directions:
|
|
ray = rotation.TransformDir(Gf.Vec3d(*direction)).GetNormalized()
|
|
hit = self.query.raycast_closest(tuple(mount), tuple(ray), 8.0, bothSides=True)
|
|
if hit["hit"] and not str(hit.get("rigidBody", "")).startswith("/World/Rover"):
|
|
points.append(hit["position"])
|
|
return np.asarray(points, dtype=np.float32).reshape(-1, 3)
|