Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

scripts/demos/raycast_sensor.py

@@ -0,0 +1,252 @@
+"""Raycast sensor demo.
+
+Run with:
+  uv run mjpython scripts/demos/raycast_sensor.py [--viewer native|viser]  # macOS
+  uv run python scripts/demos/raycast_sensor.py [--viewer native|viser]    # Linux
+
+Examples:
+  # Grid pattern (default)
+  uv run python scripts/demos/raycast_sensor.py --pattern grid
+
+  # Pinhole camera pattern
+  uv run python scripts/demos/raycast_sensor.py --pattern pinhole
+
+  # With yaw alignment (ignores pitch/roll)
+  uv run python scripts/demos/raycast_sensor.py --alignment yaw
+
+If using the native viewer, you can launch in interactive mode with:
+  uv run mjpython scripts/demos/raycast_sensor.py --viewer native --interactive
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Literal
+
+import mujoco
+import numpy as np
+import torch
+import tyro
+
+import mjlab
+import mjlab.terrains as terrain_gen
+from mjlab.entity import EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
+from mjlab.rl import RslRlVecEnvWrapper
+from mjlab.scene import SceneCfg
+from mjlab.sensor import (
+  GridPatternCfg,
+  ObjRef,
+  PinholeCameraPatternCfg,
+  RayCastSensorCfg,
+)
+from mjlab.terrains.terrain_entity import TerrainEntityCfg
+from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
+from mjlab.utils.torch import configure_torch_backends
+from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
+
+
+def create_scanner_spec() -> mujoco.MjSpec:
+  spec = mujoco.MjSpec()
+  spec.modelname = "scanner"
+
+  mat = spec.add_material()
+  mat.name = "scanner_mat"
+  mat.rgba[:] = (1.0, 0.5, 0.0, 0.9)
+
+  scanner = spec.worldbody.add_body(mocap=True)
+  scanner.name = "scanner"
+  scanner.pos[:] = (0, 0, 2.0)
+
+  geom = scanner.add_geom()
+  geom.name = "scanner_geom"
+  geom.type = mujoco.mjtGeom.mjGEOM_BOX
+  geom.size[:] = (0.15, 0.15, 0.05)
+  geom.mass = 1.0
+  geom.material = "scanner_mat"
+
+  scanner.add_camera(name="scanner", fovy=58.0, resolution=(16, 12))
+
+  record_cam = scanner.add_camera(name="record_cam")
+  record_cam.pos[:] = (2, 0, 2)
+  record_cam.fovy = 40.0
+  record_cam.mode = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY
+  record_cam.targetbody = "scanner"
+
+  return spec
+
+
+def create_env_cfg(
+  pattern: Literal["grid", "pinhole"],
+  alignment: Literal["base", "yaw", "world"],
+) -> ManagerBasedRlEnvCfg:
+  custom_terrain_cfg = TerrainGeneratorCfg(
+    size=(4.0, 4.0),
+    border_width=0.5,
+    num_rows=1,
+    num_cols=4,
+    curriculum=True,
+    sub_terrains={
+      "pyramid_stairs_inv": terrain_gen.BoxInvertedPyramidStairsTerrainCfg(
+        proportion=0.25,
+        step_height_range=(0.1, 0.25),
+        step_width=0.3,
+        platform_width=1.5,
+        border_width=0.25,
+      ),
+      "hf_pyramid_slope_inv": terrain_gen.HfPyramidSlopedTerrainCfg(
+        proportion=0.25,
+        slope_range=(0.6, 1.5),
+        platform_width=1.5,
+        border_width=0.25,
+        inverted=True,
+      ),
+      "random_rough": terrain_gen.HfRandomUniformTerrainCfg(
+        proportion=0.25,
+        noise_range=(0.05, 0.15),
+        noise_step=0.02,
+        border_width=0.25,
+      ),
+      "wave_terrain": terrain_gen.HfWaveTerrainCfg(
+        proportion=0.25,
+        amplitude_range=(0.15, 0.25),
+        num_waves=3,
+        border_width=0.25,
+      ),
+    },
+    add_lights=True,
+  )
+
+  terrain_cfg = TerrainEntityCfg(
+    terrain_type="generator",
+    terrain_generator=custom_terrain_cfg,
+    num_envs=1,
+  )
+
+  scanner_entity_cfg = EntityCfg(
+    spec_fn=create_scanner_spec,
+    init_state=EntityCfg.InitialStateCfg(pos=(0.65, -0.4, 0.5)),
+  )
+
+  if pattern == "grid":
+    pattern_cfg = GridPatternCfg(
+      size=(0.6, 0.6),
+      resolution=0.1,
+      direction=(0.0, 0.0, -1.0),
+    )
+  else:
+    assert pattern == "pinhole"
+    pattern_cfg = PinholeCameraPatternCfg.from_mujoco_camera("scanner/scanner")
+
+  raycast_cfg = RayCastSensorCfg(
+    name="terrain_scan",
+    frame=ObjRef(type="body", name="scanner", entity="scanner"),
+    pattern=pattern_cfg,
+    ray_alignment=alignment,
+    max_distance=5.0,
+    exclude_parent_body=True,
+    debug_vis=True,
+    viz=RayCastSensorCfg.VizCfg(
+      hit_color=(0.0, 1.0, 0.0, 0.9),
+      miss_color=(1.0, 0.0, 0.0, 0.5),
+      show_rays=False,
+      show_normals=True,
+    ),
+  )
+
+  cfg = ManagerBasedRlEnvCfg(
+    decimation=10,
+    scene=SceneCfg(
+      num_envs=1,
+      env_spacing=0.0,
+      extent=2.0,
+      terrain=terrain_cfg,
+      entities={"scanner": scanner_entity_cfg},
+      sensors=(raycast_cfg,),
+    ),
+  )
+
+  cfg.viewer.body_name = "scanner"
+  cfg.viewer.distance = 12.0
+  cfg.viewer.elevation = -25.0
+  cfg.viewer.azimuth = 135.0
+
+  return cfg
+
+
+def main(
+  viewer: str = "auto",
+  interactive: bool = False,
+  pattern: Literal["grid", "pinhole"] = "grid",
+  alignment: Literal["base", "yaw", "world"] = "base",
+) -> None:
+  configure_torch_backends()
+
+  device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+  print("=" * 60)
+  print("Raycast Sensor Demo - 4 Terrain Types")
+  print(f"  Pattern: {pattern}")
+  print(f"  Alignment: {alignment}")
+  print("=" * 60)
+  print()
+
+  env_cfg = create_env_cfg(pattern, alignment)
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  env = RslRlVecEnvWrapper(env)
+
+  if viewer == "auto":
+    has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"))
+    resolved_viewer = "native" if has_display else "viser"
+  else:
+    resolved_viewer = viewer
+
+  use_auto_scan = (resolved_viewer == "viser") or (not interactive)
+
+  if use_auto_scan:
+
+    class AutoScanPolicy:
+      def __init__(self):
+        self.step_count = 0
+
+      def __call__(self, obs) -> torch.Tensor:
+        del obs
+        t = self.step_count * 0.005
+        y_period = 1000
+        y_normalized = (self.step_count % y_period) / y_period
+        y = -8.0 + 16.0 * y_normalized
+        x = 1.5 * np.sin(2 * np.pi * t * 0.3)
+        z = 1.0
+        env.unwrapped.sim.data.mocap_pos[0, 0, :] = torch.tensor(
+          [x, y, z], device=device, dtype=torch.float32
+        )
+        env.unwrapped.sim.data.mocap_quat[0, 0, :] = torch.tensor(
+          [1, 0, 0, 0], device=device, dtype=torch.float32
+        )
+        self.step_count += 1
+        return torch.zeros(env.unwrapped.action_space.shape, device=device)
+
+    policy = AutoScanPolicy()
+  else:
+
+    class PolicyZero:
+      def __call__(self, obs) -> torch.Tensor:
+        del obs
+        return torch.zeros(env.unwrapped.action_space.shape, device=device)
+
+    policy = PolicyZero()
+
+  if resolved_viewer == "native":
+    print("Launching native viewer...")
+    NativeMujocoViewer(env, policy).run()
+  elif resolved_viewer == "viser":
+    print("Launching viser viewer...")
+    ViserPlayViewer(env, policy).run()
+  else:
+    raise ValueError(f"Unknown viewer: {viewer}")
+
+  env.close()
+
+
+if __name__ == "__main__":
+  tyro.cli(main, config=mjlab.TYRO_FLAGS)