Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 892971818
Change-Id: Ibce0c41294a4fcf2825a246a9533e8927fdc7e82

Author: thowell

mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py

@@ -64,6 +64,7 @@
 from mujoco.mjx.third_party.mujoco_warp._src.render import render as render
 from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_depth as get_depth
 from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_rgb as get_rgb
+from mujoco.mjx.third_party.mujoco_warp._src.render_util import get_segmentation as get_segmentation
 from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_pos as energy_pos
 from mujoco.mjx.third_party.mujoco_warp._src.sensor import energy_vel as energy_vel
 from mujoco.mjx.third_party.mujoco_warp._src.sensor import sensor_acc as sensor_acc
@@ -92,6 +93,7 @@
 from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType as BiasType
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter as BroadphaseFilter
 from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType as BroadphaseType
+from mujoco.mjx.third_party.mujoco_warp._src.types import Callback as Callback
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType as ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Constraint as Constraint
 from mujoco.mjx.third_party.mujoco_warp._src.types import Contact as Contact