32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
1245 lines
39 KiB
Python
1245 lines
39 KiB
Python
"""Tests for raycast_sensor.py."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device, make_scene_and_sim
|
|
|
|
from mjlab.envs.mdp.observations import height_scan
|
|
from mjlab.sensor import (
|
|
GridPatternCfg,
|
|
ObjRef,
|
|
PinholeCameraPatternCfg,
|
|
RayCastData,
|
|
RayCastSensorCfg,
|
|
RingPatternCfg,
|
|
)
|
|
from mjlab.sensor.raycast_sensor import _geom_groups_to_vec6
|
|
from mjlab.sim import MujocoCfg, SimulationCfg
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def device():
|
|
return get_test_device()
|
|
|
|
|
|
def test_basic_raycast_hit_detection(robot_with_floor_xml, device):
|
|
"""Verify rays detect the ground plane and return correct distances."""
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="terrain_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(
|
|
size=(0.5, 0.5), resolution=0.25, direction=(0.0, 0.0, -1.0)
|
|
),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(
|
|
device, robot_with_floor_xml, sensors=(raycast_cfg,), num_envs=2
|
|
)
|
|
|
|
sensor = scene["terrain_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
assert isinstance(data, RayCastData)
|
|
assert data.distances.shape[0] == 2 # num_envs
|
|
assert data.distances.shape[1] == sensor.num_rays
|
|
assert data.normals_w.shape == (2, sensor.num_rays, 3)
|
|
|
|
# All rays should hit the floor (distance > 0).
|
|
assert torch.all(data.distances >= 0)
|
|
|
|
# Distance should be approximately 2m (body at z=2, floor at z=0).
|
|
assert torch.allclose(data.distances, torch.full_like(data.distances, 2.0), atol=0.1)
|
|
|
|
|
|
def test_raycast_normals_point_up(robot_with_floor_xml, device):
|
|
"""Verify surface normals point upward when hitting a horizontal floor."""
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="terrain_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(
|
|
size=(0.3, 0.3), resolution=0.15, direction=(0.0, 0.0, -1.0)
|
|
),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, robot_with_floor_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["terrain_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
# Normals should point up (+Z) for a horizontal floor.
|
|
expected = torch.zeros_like(data.normals_w)
|
|
expected[:, :, 2] = 1.0
|
|
assert torch.allclose(data.normals_w, expected)
|
|
|
|
|
|
def test_raycast_miss_returns_negative_one(device):
|
|
"""Verify rays that miss return distance of -1."""
|
|
no_floor_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="terrain_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(
|
|
size=(0.3, 0.3), resolution=0.15, direction=(0.0, 0.0, -1.0)
|
|
),
|
|
max_distance=10.0,
|
|
exclude_parent_body=True,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, no_floor_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["terrain_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
# All rays should miss (distance = -1).
|
|
assert torch.all(data.distances == -1)
|
|
|
|
|
|
def test_raycast_exclude_parent_body(robot_with_floor_xml, device):
|
|
"""Verify parent body is excluded from ray intersection when configured."""
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="terrain_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.1, 0.1), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=10.0,
|
|
exclude_parent_body=True,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, robot_with_floor_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["terrain_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
# Rays should hit the floor, not the parent body geom.
|
|
# Floor is at z=0, body is at z=2, so distance should be ~2m.
|
|
assert torch.allclose(data.distances, torch.full_like(data.distances, 2.0), atol=0.1)
|
|
|
|
|
|
def test_raycast_include_geom_groups(device):
|
|
"""Verify include_geom_groups filters which geoms are hit."""
|
|
groups_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0" group="0"/>
|
|
<geom name="platform" type="box" size="1 1 0.1" pos="0 0 1" group="1"/>
|
|
<body name="base" pos="0 0 3">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
# Only include group 0 (floor) - should skip the platform in group 1.
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="group_filter_test",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=10.0,
|
|
include_geom_groups=(0,),
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, groups_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["group_filter_test"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
# Should hit floor at z=0, not platform at z=1.1. Distance from z=3 to z=0 is 3m.
|
|
assert torch.allclose(data.distances, torch.full_like(data.distances, 3.0), atol=0.1)
|
|
|
|
# Now test with group 1 included - should hit platform instead.
|
|
raycast_cfg_group1 = RayCastSensorCfg(
|
|
name="group1_test",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=10.0,
|
|
include_geom_groups=(1,),
|
|
)
|
|
|
|
scene2, sim2 = make_scene_and_sim(device, groups_xml, sensors=(raycast_cfg_group1,))
|
|
|
|
sensor2 = scene2["group1_test"]
|
|
sim2.step()
|
|
sim2.sense()
|
|
data2 = sensor2.data
|
|
|
|
# Should hit platform at z=1.1. Distance from z=3 to z=1.1 is 1.9m.
|
|
assert torch.allclose(
|
|
data2.distances, torch.full_like(data2.distances, 1.9), atol=0.1
|
|
)
|
|
|
|
|
|
def test_raycast_frame_attachment_geom(device):
|
|
"""Verify rays can be attached to a geom frame."""
|
|
geom_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="sensor_mount" type="box" size="0.1 0.1 0.05" pos="0 0 -0.05"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="geom_scan",
|
|
frame=ObjRef(type="geom", name="sensor_mount", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.2, 0.2), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, geom_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["geom_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
assert isinstance(data, RayCastData)
|
|
# Geom is at z=1.95 (body at z=2, geom offset -0.05), floor at z=0.
|
|
assert torch.allclose(data.distances, torch.full_like(data.distances, 1.95), atol=0.1)
|
|
|
|
|
|
def test_raycast_frame_attachment_site(robot_with_floor_xml, device):
|
|
"""Verify rays can be attached to a site frame."""
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="site_scan",
|
|
frame=ObjRef(type="site", name="base_site", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.2, 0.2), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, robot_with_floor_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["site_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
assert isinstance(data, RayCastData)
|
|
# Site is at z=1.9 (body at z=2, site offset -0.1), floor at z=0.
|
|
assert torch.allclose(data.distances, torch.full_like(data.distances, 1.9), atol=0.1)
|
|
|
|
|
|
def test_raycast_grid_pattern_num_rays(device):
|
|
"""Verify grid pattern generates correct number of rays."""
|
|
simple_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1"/>
|
|
<body name="base" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
# Grid: size=(1.0, 0.5), resolution=0.5.
|
|
# X: from -0.5 to 0.5 step 0.5 -> 3 points.
|
|
# Y: from -0.25 to 0.25 step 0.5 -> 2 points.
|
|
# Total: 3 * 2 = 6 rays.
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="grid_test",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(1.0, 0.5), resolution=0.5),
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, simple_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["grid_test"]
|
|
assert sensor.num_rays == 6
|
|
|
|
|
|
def test_raycast_different_direction(device):
|
|
"""Verify rays work with non-default direction."""
|
|
wall_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="wall" type="box" size="0.1 5 5" pos="2 0 2"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="forward_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.2, 0.2), resolution=0.1, direction=(1.0, 0.0, 0.0)),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, wall_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["forward_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
# Wall is at x=1.9 (wall center at x=2, size 0.1), body at x=0.
|
|
# Distance should be ~1.9m.
|
|
assert torch.allclose(data.distances, torch.full_like(data.distances, 1.9), atol=0.1)
|
|
|
|
# Normal should point in -X direction (toward the body).
|
|
assert torch.allclose(
|
|
data.normals_w[:, :, 0], -torch.ones_like(data.normals_w[:, :, 0]), atol=0.01
|
|
)
|
|
|
|
|
|
def test_raycast_error_on_invalid_frame_type(device):
|
|
"""Verify ValueError is raised for invalid frame type."""
|
|
simple_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base"><geom type="sphere" size="0.1"/></body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="invalid",
|
|
frame=ObjRef(type="joint", name="some_joint", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.1, 0.1), resolution=0.1),
|
|
)
|
|
with pytest.raises(ValueError, match="must be 'body', 'site', or 'geom'"):
|
|
make_scene_and_sim(device, simple_xml, sensors=(raycast_cfg,))
|
|
|
|
|
|
def test_raycast_hit_pos_w_correctness(robot_with_floor_xml, device):
|
|
"""Verify hit_pos_w returns correct world-space hit positions."""
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="terrain_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.4, 0.4), resolution=0.2, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, robot_with_floor_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["terrain_scan"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
# All hit positions should be on the floor (z=0).
|
|
assert torch.allclose(
|
|
data.hit_pos_w[:, :, 2], torch.zeros_like(data.hit_pos_w[:, :, 2]), atol=0.01
|
|
)
|
|
|
|
# Hit positions X and Y should match the ray grid pattern offset from body origin.
|
|
# Body is at (0, 0, 2), grid is 0.4x0.4 with 0.2 resolution = 3x3 = 9 rays.
|
|
# X positions should be in range [-0.2, 0.2], Y positions in range [-0.2, 0.2].
|
|
assert torch.all(data.hit_pos_w[:, :, 0] >= -0.3)
|
|
assert torch.all(data.hit_pos_w[:, :, 0] <= 0.3)
|
|
assert torch.all(data.hit_pos_w[:, :, 1] >= -0.3)
|
|
assert torch.all(data.hit_pos_w[:, :, 1] <= 0.3)
|
|
|
|
|
|
def test_raycast_max_distance_clamping(device):
|
|
"""Verify hits beyond max_distance are reported as misses."""
|
|
far_floor_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 5">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="short_range",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.2, 0.2), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=3.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, far_floor_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["short_range"]
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
|
|
# All rays should miss (floor is beyond max_distance).
|
|
assert torch.all(data.distances == -1)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
|
|
def test_raycast_body_rotation_affects_rays(device):
|
|
"""Verify rays rotate with the body frame."""
|
|
rotated_body_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="rotated_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, 0)), njmax=20)
|
|
scene, sim = make_scene_and_sim(
|
|
device, rotated_body_xml, sensors=(raycast_cfg,), sim_cfg=sim_cfg
|
|
)
|
|
|
|
sensor = scene["rotated_scan"]
|
|
|
|
# First, verify baseline: unrotated body, rays hit floor at ~2m.
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data_unrotated = sensor.data
|
|
assert torch.allclose(
|
|
data_unrotated.distances, torch.full_like(data_unrotated.distances, 2.0), atol=0.1
|
|
)
|
|
|
|
# Now tilt body 45 degrees around X axis.
|
|
# Ray direction -Z in body frame becomes diagonal in world frame.
|
|
# Distance to floor should be 2 / cos(45) = 2 * sqrt(2) ≈ 2.83m.
|
|
angle = math.pi / 4
|
|
quat = [math.cos(angle / 2), math.sin(angle / 2), 0, 0] # w, x, y, z
|
|
sim.data.qpos[0, 3:7] = torch.tensor(quat, device=device)
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data_rotated = sensor.data
|
|
|
|
expected_distance = 2.0 / math.cos(angle) # ~2.83m
|
|
assert torch.allclose(
|
|
data_rotated.distances,
|
|
torch.full_like(data_rotated.distances, expected_distance),
|
|
atol=0.15,
|
|
), f"Expected ~{expected_distance:.2f}m, got {data_rotated.distances}"
|
|
|
|
|
|
# ============================================================================
|
|
# Pinhole Camera Pattern Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_pinhole_camera_pattern_num_rays(device):
|
|
"""Verify pinhole pattern generates width * height rays."""
|
|
simple_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1"/>
|
|
<body name="base" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="camera_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=PinholeCameraPatternCfg(width=16, height=12, fovy=74.0),
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, simple_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["camera_scan"]
|
|
assert sensor.num_rays == 16 * 12
|
|
|
|
|
|
def test_pinhole_from_intrinsic_matrix():
|
|
"""Verify from_intrinsic_matrix creates correct config."""
|
|
# Intrinsic matrix with fx=500, fy=500, cx=320, cy=240.
|
|
intrinsic = [500.0, 0, 320, 0, 500.0, 240, 0, 0, 1]
|
|
width, height = 640, 480
|
|
|
|
cfg = PinholeCameraPatternCfg.from_intrinsic_matrix(intrinsic, width, height)
|
|
|
|
# Expected vertical FOV: 2 * atan(480 / (2 * 500)) = 2 * atan(0.48) ≈ 51.3 degrees.
|
|
fy = intrinsic[4]
|
|
expected_fov = 2 * math.atan(height / (2 * fy)) * 180 / math.pi
|
|
assert abs(cfg.fovy - expected_fov) < 0.1
|
|
assert cfg.width == width
|
|
assert cfg.height == height
|
|
|
|
|
|
def test_pinhole_from_mujoco_camera(device):
|
|
"""Verify pinhole pattern can be created from MuJoCo camera."""
|
|
# XML with a camera that has explicit resolution, sensorsize, and focal.
|
|
camera_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
<camera name="depth_cam" pos="0 0 0" resolution="64 48"
|
|
sensorsize="0.00389 0.00292" focal="0.00193 0.00193"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
# Use from_mujoco_camera() to get params from MuJoCo camera.
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="camera_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=PinholeCameraPatternCfg.from_mujoco_camera("robot/depth_cam"),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, camera_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["camera_scan"]
|
|
# Should have 64 * 48 = 3072 rays.
|
|
assert sensor.num_rays == 64 * 48
|
|
|
|
# Verify rays work.
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
assert torch.all(data.distances >= 0) # Should hit floor
|
|
|
|
|
|
def test_pinhole_from_mujoco_camera_fovy_mode(device):
|
|
"""Verify pinhole pattern works with MuJoCo camera using fovy (not sensorsize)."""
|
|
# XML with a camera using fovy mode (no sensorsize/focal).
|
|
camera_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
<camera name="fovy_cam" pos="0 0 0" fovy="60" resolution="32 24"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="camera_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=PinholeCameraPatternCfg.from_mujoco_camera("robot/fovy_cam"),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, camera_xml, sensors=(raycast_cfg,))
|
|
|
|
sensor = scene["camera_scan"]
|
|
# Should have 32 * 24 = 768 rays.
|
|
assert sensor.num_rays == 32 * 24
|
|
|
|
# Verify rays work.
|
|
sim.step()
|
|
sim.sense()
|
|
data = sensor.data
|
|
assert torch.all(data.distances >= 0) # Should hit floor
|
|
|
|
|
|
# ============================================================================
|
|
# Ray Alignment Tests
|
|
# ============================================================================
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
|
|
def test_ray_alignment_yaw(device):
|
|
"""Verify yaw alignment ignores pitch/roll."""
|
|
rotated_body_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
# With yaw alignment, tilting the body should NOT affect ray direction.
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="yaw_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
ray_alignment="yaw",
|
|
max_distance=10.0,
|
|
)
|
|
|
|
sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, 0)), njmax=20)
|
|
scene, sim = make_scene_and_sim(
|
|
device, rotated_body_xml, sensors=(raycast_cfg,), sim_cfg=sim_cfg
|
|
)
|
|
|
|
sensor = scene["yaw_scan"]
|
|
|
|
# Baseline: unrotated.
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data_unrotated = sensor.data
|
|
baseline_dist = data_unrotated.distances.clone()
|
|
|
|
# Tilt body 45 degrees around X axis.
|
|
angle = math.pi / 4
|
|
quat = [math.cos(angle / 2), math.sin(angle / 2), 0, 0] # w, x, y, z
|
|
sim.data.qpos[0, 3:7] = torch.tensor(quat, device=device)
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data_tilted = sensor.data
|
|
|
|
# With yaw alignment, distance should remain ~2m (not change due to tilt).
|
|
assert torch.allclose(data_tilted.distances, baseline_dist, atol=0.1), (
|
|
f"Expected ~2m, got {data_tilted.distances}"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
|
|
def test_ray_alignment_world(device):
|
|
"""Verify world alignment keeps rays fixed."""
|
|
rotated_body_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
# With world alignment, rotating body should NOT affect ray direction.
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="world_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1, direction=(0.0, 0.0, -1.0)),
|
|
ray_alignment="world",
|
|
max_distance=10.0,
|
|
)
|
|
|
|
sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, 0)), njmax=20)
|
|
scene, sim = make_scene_and_sim(
|
|
device, rotated_body_xml, sensors=(raycast_cfg,), sim_cfg=sim_cfg
|
|
)
|
|
|
|
sensor = scene["world_scan"]
|
|
|
|
# Baseline: unrotated.
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data_unrotated = sensor.data
|
|
baseline_dist = data_unrotated.distances.clone()
|
|
|
|
# Rotate body 90 degrees around Z (yaw), then tilt 45 degrees around X.
|
|
# With world alignment, distance should still be ~2m.
|
|
yaw_angle = math.pi / 2
|
|
pitch_angle = math.pi / 4
|
|
# Compose quaternions: yaw then pitch.
|
|
cy, sy = math.cos(yaw_angle / 2), math.sin(yaw_angle / 2)
|
|
cp, sp = math.cos(pitch_angle / 2), math.sin(pitch_angle / 2)
|
|
# q_yaw = [cy, 0, 0, sy], q_pitch = [cp, sp, 0, 0]
|
|
# q = q_pitch * q_yaw
|
|
qw = cp * cy
|
|
qx = sp * cy
|
|
qy = sp * sy
|
|
qz = cp * sy
|
|
sim.data.qpos[0, 3:7] = torch.tensor([qw, qx, qy, qz], device=device)
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data_rotated = sensor.data
|
|
|
|
# With world alignment, distance should remain ~2m.
|
|
assert torch.allclose(data_rotated.distances, baseline_dist, atol=0.1), (
|
|
f"Expected ~2m, got {data_rotated.distances}"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
|
|
def test_ray_alignment_yaw_singularity(device):
|
|
"""Test yaw alignment handles 90 degree pitch singularity correctly.
|
|
|
|
With yaw alignment, rays should maintain their pattern regardless of
|
|
body pitch. At 90 degree pitch, the body's X-axis is vertical, making
|
|
yaw extraction ambiguous. The implementation uses Y-axis fallback to
|
|
produce a valid yaw rotation.
|
|
|
|
This test verifies that distances at 90 degree pitch match the
|
|
baseline (0 degree pitch).
|
|
"""
|
|
xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1" mass="1.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
# Use grid pattern with diagonal direction - has X component to
|
|
# expose singularity. Direction [1, 0, -1] points forward and down
|
|
# at 45 degrees.
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="yaw_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1, direction=(1.0, 0.0, -1.0)),
|
|
ray_alignment="yaw",
|
|
max_distance=10.0,
|
|
)
|
|
|
|
sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, 0)), njmax=20)
|
|
scene, sim = make_scene_and_sim(device, xml, sensors=(raycast_cfg,), sim_cfg=sim_cfg)
|
|
|
|
sensor = scene["yaw_scan"]
|
|
|
|
# Baseline: no rotation. Ray at 45 degrees from height 2m hits floor
|
|
# at x=2, z=0.
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
baseline_hit_pos = sensor.data.hit_pos_w.clone()
|
|
# Ray goes diagonally +X and -Z, starting from (0,0,2), should hit
|
|
# floor at (2, 0, 0).
|
|
assert torch.allclose(
|
|
baseline_hit_pos[0, 0, 0],
|
|
torch.tensor(2.0, device=device),
|
|
atol=0.1,
|
|
), f"Baseline X hit should be ~2, got {baseline_hit_pos[0, 0, 0]}"
|
|
assert torch.allclose(
|
|
baseline_hit_pos[0, 0, 2],
|
|
torch.tensor(0.0, device=device),
|
|
atol=0.1,
|
|
), f"Baseline Z hit should be ~0, got {baseline_hit_pos[0, 0, 2]}"
|
|
|
|
# Pitch 90 degrees around Y-axis. Body X-axis now points straight
|
|
# down (singularity).
|
|
angle = math.pi / 2
|
|
quat = [math.cos(angle / 2), 0, math.sin(angle / 2), 0] # w, x, y, z
|
|
sim.data.qpos[0, 3:7] = torch.tensor(quat, device=device)
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
|
|
singularity_hit_pos = sensor.data.hit_pos_w
|
|
|
|
# With yaw alignment, hit position should match baseline regardless
|
|
# of pitch. The ray should still go diagonally and hit at (2, 0, 0).
|
|
assert torch.allclose(singularity_hit_pos, baseline_hit_pos, atol=0.1), (
|
|
f"Yaw alignment failed at 90 degree pitch singularity.\n"
|
|
f"Baseline hit_pos: {baseline_hit_pos}\n"
|
|
f"Singularity hit_pos: {singularity_hit_pos}"
|
|
)
|
|
|
|
|
|
class _FakeEnv:
|
|
"""Minimal env-like object for testing observation functions."""
|
|
|
|
def __init__(self, scene):
|
|
self.scene = scene
|
|
|
|
|
|
def test_height_scan_hits(robot_with_floor_xml, device):
|
|
"""height_scan returns correct heights for hits."""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="terrain_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(
|
|
size=(0.3, 0.3),
|
|
resolution=0.15,
|
|
direction=(0.0, 0.0, -1.0),
|
|
),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(
|
|
device, robot_with_floor_xml, sensors=(raycast_cfg,), num_envs=2
|
|
)
|
|
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
env = _FakeEnv(scene)
|
|
heights = height_scan(env, "terrain_scan") # type: ignore[invalid-argument-type]
|
|
|
|
sensor = scene["terrain_scan"]
|
|
# Shape: [num_envs, num_rays].
|
|
assert heights.shape == (2, sensor.num_rays)
|
|
# Body at z=2, floor at z=0 → height ≈ 2.0.
|
|
assert torch.allclose(heights, torch.full_like(heights, 2.0), atol=0.1)
|
|
|
|
|
|
def test_height_scan_misses(device):
|
|
"""height_scan reports max_distance for rays that miss (no ground)."""
|
|
|
|
no_floor_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
raycast_cfg = RayCastSensorCfg(
|
|
name="terrain_scan",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=GridPatternCfg(
|
|
size=(0.2, 0.2),
|
|
resolution=0.1,
|
|
direction=(0.0, 0.0, -1.0),
|
|
),
|
|
max_distance=10.0,
|
|
exclude_parent_body=True,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, no_floor_xml, sensors=(raycast_cfg,))
|
|
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
env = _FakeEnv(scene)
|
|
heights = height_scan(env, "terrain_scan") # type: ignore[invalid-argument-type]
|
|
|
|
# Misses default to sensor max_distance.
|
|
assert torch.allclose(
|
|
heights, torch.full_like(heights, raycast_cfg.max_distance), atol=1e-5
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Multi-Frame Tests
|
|
# ============================================================================
|
|
|
|
MULTI_SITE_XML = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 3">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
|
|
<site name="site_top" pos="0 0 0.5"/>
|
|
<site name="site_mid" pos="0 0 0"/>
|
|
<site name="site_bot" pos="0 0 -0.5"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
def test_multi_frame_shapes(device):
|
|
"""Multi-frame sensor produces correct output shapes."""
|
|
cfg = RayCastSensorCfg(
|
|
name="multi",
|
|
frame=(
|
|
ObjRef(type="site", name="site_top", entity="robot"),
|
|
ObjRef(type="site", name="site_mid", entity="robot"),
|
|
ObjRef(type="site", name="site_bot", entity="robot"),
|
|
),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, MULTI_SITE_XML, (cfg,), num_envs=2)
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
sensor = scene["multi"]
|
|
data = sensor.data
|
|
|
|
assert sensor.num_frames == 3
|
|
assert sensor.num_rays_per_frame == 1
|
|
assert sensor.num_rays == 3
|
|
|
|
assert data.distances.shape == (2, 3)
|
|
assert data.frame_pos_w.shape == (2, 3, 3)
|
|
assert data.frame_quat_w.shape == (2, 3, 4)
|
|
# Backward compat: pos_w/quat_w equal first frame.
|
|
assert data.pos_w.shape == (2, 3)
|
|
assert data.quat_w.shape == (2, 4)
|
|
assert torch.allclose(data.pos_w, data.frame_pos_w[:, 0])
|
|
assert torch.allclose(data.quat_w, data.frame_quat_w[:, 0])
|
|
|
|
|
|
def test_multi_frame_heights(device):
|
|
"""Sites at different Z offsets produce proportional distances."""
|
|
cfg = RayCastSensorCfg(
|
|
name="multi",
|
|
frame=(
|
|
ObjRef(type="site", name="site_top", entity="robot"),
|
|
ObjRef(type="site", name="site_mid", entity="robot"),
|
|
ObjRef(type="site", name="site_bot", entity="robot"),
|
|
),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, MULTI_SITE_XML, (cfg,))
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
distances = scene["multi"].data.distances[0]
|
|
# site_top at z=3.5, site_mid at z=3, site_bot at z=2.5.
|
|
assert distances[0].item() == pytest.approx(3.5, abs=0.1)
|
|
assert distances[1].item() == pytest.approx(3.0, abs=0.1)
|
|
assert distances[2].item() == pytest.approx(2.5, abs=0.1)
|
|
|
|
|
|
def test_multi_frame_body_exclusion(device):
|
|
"""Each frame excludes only its own parent body, not others.
|
|
|
|
body_a has a platform geom directly below site_b. Frame B's rays
|
|
should skip body_b's own geom but HIT body_a's platform. Frame A's
|
|
rays should skip body_a and hit the floor.
|
|
"""
|
|
body_a_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
|
|
<body name="body_a" pos="0 0 1">
|
|
<freejoint name="free_a"/>
|
|
<geom name="geom_a" type="box" size="2 2 0.1" mass="5.0"/>
|
|
<site name="site_a" pos="0 0 0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
body_b_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="body_b" pos="0 0 3">
|
|
<freejoint name="free_b"/>
|
|
<geom name="geom_b" type="box" size="0.5 0.5 0.5" mass="5.0"/>
|
|
<site name="site_b" pos="0 0 0"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
cfg = RayCastSensorCfg(
|
|
name="multi",
|
|
frame=(
|
|
ObjRef(type="site", name="site_a", entity="body_a"),
|
|
ObjRef(type="site", name="site_b", entity="body_b"),
|
|
),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
|
|
max_distance=10.0,
|
|
exclude_parent_body=True,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(
|
|
device, {"body_a": body_a_xml, "body_b": body_b_xml}, (cfg,)
|
|
)
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
distances = scene["multi"].data.distances[0]
|
|
# Frame A (site_a at z=1): excludes body_a, hits floor at z=0 -> dist ~1.0.
|
|
assert distances[0].item() == pytest.approx(1.0, abs=0.15)
|
|
# Frame B (site_b at z=3): excludes body_b, hits body_a's platform
|
|
# at z=1.1 -> dist ~1.9.
|
|
assert distances[1].item() == pytest.approx(1.9, abs=0.15)
|
|
|
|
|
|
def test_multi_frame_height_scan(device):
|
|
"""height_scan works correctly with multi-frame sensors."""
|
|
cfg = RayCastSensorCfg(
|
|
name="multi",
|
|
frame=(
|
|
ObjRef(type="site", name="site_top", entity="robot"),
|
|
ObjRef(type="site", name="site_mid", entity="robot"),
|
|
ObjRef(type="site", name="site_bot", entity="robot"),
|
|
),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, MULTI_SITE_XML, (cfg,), num_envs=2)
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
env = _FakeEnv(scene)
|
|
heights = height_scan(env, "multi") # type: ignore[invalid-argument-type]
|
|
|
|
assert heights.shape == (2, 3)
|
|
# Each frame's height = frame_z - hit_z.
|
|
# site_top at z=3.5, site_mid at z=3, site_bot at z=2.5; floor at z=0.
|
|
assert heights[0, 0].item() == pytest.approx(3.5, abs=0.1)
|
|
assert heights[0, 1].item() == pytest.approx(3.0, abs=0.1)
|
|
assert heights[0, 2].item() == pytest.approx(2.5, abs=0.1)
|
|
|
|
|
|
# ============================================================================
|
|
# Ring Pattern Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_ring_pattern_num_rays(device):
|
|
"""Ring with 8 samples + center = 9 rays."""
|
|
cfg = RayCastSensorCfg(
|
|
name="ring",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=RingPatternCfg.single_ring(radius=0.1, num_samples=8, include_center=True),
|
|
)
|
|
|
|
simple_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1"/>
|
|
<body name="base" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
scene, sim = make_scene_and_sim(device, simple_xml, (cfg,))
|
|
sensor = scene["ring"]
|
|
assert sensor.num_rays == 9
|
|
assert sensor.num_rays_per_frame == 9
|
|
|
|
|
|
def test_ring_pattern_concentric(device):
|
|
"""Concentric rings: 1 center + 4 + 6 + 8 = 19 rays."""
|
|
cfg = RayCastSensorCfg(
|
|
name="ring",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=RingPatternCfg(
|
|
rings=(
|
|
RingPatternCfg.Ring(radius=0.05, num_samples=4),
|
|
RingPatternCfg.Ring(radius=0.10, num_samples=6),
|
|
RingPatternCfg.Ring(radius=0.20, num_samples=8),
|
|
),
|
|
include_center=True,
|
|
),
|
|
)
|
|
|
|
simple_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1"/>
|
|
<body name="base" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
scene, sim = make_scene_and_sim(device, simple_xml, (cfg,))
|
|
sensor = scene["ring"]
|
|
assert sensor.num_rays == 19
|
|
|
|
|
|
def test_ring_pattern_hits_floor(device):
|
|
"""Ring pattern hits floor and ring offsets lie on a circle."""
|
|
radius = 0.1
|
|
cfg = RayCastSensorCfg(
|
|
name="ring",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=RingPatternCfg.single_ring(
|
|
radius=radius, num_samples=4, include_center=True
|
|
),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
simple_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1"/>
|
|
<body name="base" pos="0 0 2">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.05"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
scene, sim = make_scene_and_sim(device, simple_xml, (cfg,))
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
data = scene["ring"].data
|
|
# All 5 rays (1 center + 4 ring) should hit the floor.
|
|
assert torch.all(data.distances >= 0)
|
|
assert torch.allclose(data.distances, torch.full_like(data.distances, 2.0), atol=0.1)
|
|
|
|
# Verify ring geometry: non-center hits should lie on a circle of
|
|
# the given radius around the body's XY position.
|
|
hit_xy = data.hit_pos_w[0, :, :2] # [5, 2]
|
|
body_xy = data.pos_w[0, :2] # [2]
|
|
dists_from_center = (hit_xy - body_xy).norm(dim=1)
|
|
# Center ray (index 0) should be at ~0 offset.
|
|
assert dists_from_center[0].item() == pytest.approx(0.0, abs=0.02)
|
|
# Ring rays (indices 1-4) should be at ~radius offset.
|
|
for i in range(1, 5):
|
|
assert dists_from_center[i].item() == pytest.approx(radius, abs=0.02)
|
|
|
|
|
|
def test_ring_pattern_no_center(device):
|
|
"""Ring with include_center=False omits the center ray."""
|
|
cfg = RayCastSensorCfg(
|
|
name="ring",
|
|
frame=ObjRef(type="body", name="base", entity="robot"),
|
|
pattern=RingPatternCfg.single_ring(radius=0.1, num_samples=4, include_center=False),
|
|
)
|
|
simple_xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="10 10 0.1"/>
|
|
<body name="base" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="sphere" size="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
scene, sim = make_scene_and_sim(device, simple_xml, (cfg,))
|
|
assert scene["ring"].num_rays == 4
|
|
|
|
|
|
def test_multi_frame_multiple_rays_per_frame(device):
|
|
"""Multi-frame with grid pattern: verifies reshape ordering."""
|
|
cfg = RayCastSensorCfg(
|
|
name="multi",
|
|
frame=(
|
|
ObjRef(type="site", name="site_top", entity="robot"),
|
|
ObjRef(type="site", name="site_bot", entity="robot"),
|
|
),
|
|
pattern=GridPatternCfg(size=(0.2, 0.2), resolution=0.1),
|
|
max_distance=10.0,
|
|
)
|
|
|
|
scene, sim = make_scene_and_sim(device, MULTI_SITE_XML, (cfg,))
|
|
sim.step()
|
|
sim.sense()
|
|
|
|
sensor = scene["multi"]
|
|
N = sensor.num_rays_per_frame
|
|
assert N == 9 # 3x3 grid
|
|
assert sensor.num_rays == 18 # 2 frames * 9
|
|
|
|
data = sensor.data
|
|
# First frame (site_top, z=3.5) rays should all be ~3.5.
|
|
top_dists = data.distances[0, :N]
|
|
assert torch.allclose(top_dists, torch.full_like(top_dists, 3.5), atol=0.1)
|
|
# Second frame (site_bot, z=2.5) rays should all be ~2.5.
|
|
bot_dists = data.distances[0, N:]
|
|
assert torch.allclose(bot_dists, torch.full_like(bot_dists, 2.5), atol=0.1)
|
|
|
|
# height_scan should also produce correct per-frame heights.
|
|
env = _FakeEnv(scene)
|
|
heights = height_scan(env, "multi") # type: ignore[invalid-argument-type]
|
|
assert heights.shape == (1, 18)
|
|
top_heights = heights[0, :N]
|
|
bot_heights = heights[0, N:]
|
|
assert torch.allclose(top_heights, torch.full_like(top_heights, 3.5), atol=0.1)
|
|
assert torch.allclose(bot_heights, torch.full_like(bot_heights, 2.5), atol=0.1)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
|
|
def test_site_origin_is_physical_with_world_alignment(device):
|
|
"""ray_alignment only controls direction, not origin position.
|
|
|
|
When a calf body pitches, the foot site swings to its physical
|
|
position (site_xpos). Rays still point -Z (world-aligned), giving
|
|
the correct distance from the actual foot to the ground.
|
|
"""
|
|
xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="50 50 0.1" pos="0 0 0"/>
|
|
<body name="calf" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="calf_geom" type="sphere" size="0.05" mass="1.0"/>
|
|
<site name="foot" pos="0 0 -0.5"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
cfg = RayCastSensorCfg(
|
|
name="scan",
|
|
frame=ObjRef(type="site", name="foot", entity="robot"),
|
|
pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
|
|
ray_alignment="world",
|
|
max_distance=5.0,
|
|
)
|
|
|
|
sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, 0)), njmax=20)
|
|
scene, sim = make_scene_and_sim(device, xml, (cfg,), sim_cfg=sim_cfg)
|
|
|
|
# Baseline: body upright, site at z=0.5.
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data = scene["scan"].data
|
|
assert data.pos_w[0, 2].item() == pytest.approx(0.5, abs=0.05)
|
|
assert data.distances[0, 0].item() == pytest.approx(0.5, abs=0.05)
|
|
|
|
# Pitch calf 90 degrees. Site swings to (0, -0.5, 1.0).
|
|
# Ray points -Z (world), distance to floor = 1.0.
|
|
angle = math.pi / 2
|
|
quat = [math.cos(angle / 2), math.sin(angle / 2), 0, 0]
|
|
sim.data.qpos[0, 3:7] = torch.tensor(quat, device=device)
|
|
sim.step()
|
|
scene.update(dt=sim.cfg.mujoco.timestep)
|
|
sim.sense()
|
|
data = scene["scan"].data
|
|
|
|
assert data.pos_w[0, 2].item() == pytest.approx(1.0, abs=0.1)
|
|
assert data.distances[0, 0].item() == pytest.approx(1.0, abs=0.1)
|
|
|
|
|
|
@pytest.mark.parametrize("groups", [(6,), (-1,), (0, 7)])
|
|
def test_include_geom_groups_out_of_range_raises(groups):
|
|
"""Out-of-range geom groups should raise rather than silently exclude."""
|
|
with pytest.raises(ValueError, match="include_geom_groups must be in"):
|
|
_geom_groups_to_vec6(groups)
|