Add Capsule geom

crates/bevy_mujoco/src/lib.rs

@@ -314,36 +314,28 @@ fn spawn_geom(
     geom: &Geom,
 ) {
     let material = match geom.material {
-        Material::Rgba { rgba: [r, g, b, a] } => materials.add(Color::srgba_u8(
-            (r * 255.0) as u8,
-            (g * 255.0) as u8,
-            (b * 255.0) as u8,
-            (a * 255.0) as u8,
-        )),
+        Material::Rgba { rgba: [r, g, b, a] } => materials.add(Color::srgba(r, g, b, a)),
         Material::Pbr { material_index } => material_handles[&material_index].clone(),
     };
 
-    let mut geom_entity = entity_commands.spawn((
-        Transform::from_translation(Vec3::from(geom.pos)).with_rotation(bevy_quat(geom.quat)),
-        MeshMaterial3d(material),
-    ));
-
-    let mesh_handle = match &geom.geom_variant {
-        GeomVariant::Mesh { mesh_index } => mesh_handles[mesh_index].clone(),
-        GeomVariant::Sphere { radius } => meshes.add(Sphere::new(*radius)),
+    let (mesh_handle, alignment_rotation) = match geom.geom_variant {
+        GeomVariant::Mesh { mesh_index } => (mesh_handles[&mesh_index].clone(), Quat::IDENTITY),
+        GeomVariant::Sphere { radius } => (meshes.add(Sphere::new(radius)), Quat::IDENTITY),
         GeomVariant::Box {
             extent: [hx, hy, hz],
-        } => meshes.add(Cuboid::new(*hx, *hy, *hz)),
+        } => (
+            // MuJoCo box extent is half-lengths, whereas bevy cuboid takes full lengths
+            meshes.add(Cuboid::new(2. * hx, 2. * hy, 2. * hz)),
+            Quat::IDENTITY,
+        ),
         GeomVariant::Plane {
             normal: [nx, ny, nz],
         } => {
-            // The plane is supposed to be infinite but that seems mighty expensive.
-            // Use a finite size instead
             const SCALE: f32 = 100.0;
-
-            let mut mesh = Plane3d::new(Vec3::new(*nx, *ny, *nz), Vec2::splat(SCALE))
+            let mut mesh = Plane3d::new(Vec3::new(nx, ny, nz), Vec2::splat(SCALE))
                 .mesh()
                 .build();
+
             let uv = mesh
                 .attribute_mut(Mesh::ATTRIBUTE_UV_0)
                 .expect("Plane3d should generate UV attributes");
@@ -357,16 +349,35 @@ fn spawn_geom(
                 }
                 _ => panic!("expected UV coordinates to be Float32x2"),
             }
-
-            meshes.add(mesh)
+            (meshes.add(mesh), Quat::IDENTITY)
         }
         GeomVariant::Cylinder {
             radius,
             half_height,
-        } => meshes.add(Cylinder::new(*radius, *half_height)),
+        } => (
+            meshes.add(Cylinder::new(radius, 2. * half_height)),
+            Quat::from_rotation_x(FRAC_PI_2),
+        ),
+        GeomVariant::Capsule {
+            radius,
+            half_height,
+        } => (
+            meshes.add(Capsule3d::new(radius, 2. * half_height)),
+            Quat::from_rotation_x(FRAC_PI_2),
+        ),
     };
 
-    geom_entity.insert(Mesh3d(mesh_handle));
+    entity_commands
+        .spawn((
+            Transform::from_translation(Vec3::from(geom.pos)).with_rotation(bevy_quat(geom.quat)),
+        ))
+        .with_children(|parent| {
+            parent.spawn((
+                Mesh3d(mesh_handle),
+                MeshMaterial3d(material),
+                Transform::from_rotation(alignment_rotation),
+            ));
+        });
 }
 
 fn update_bodies(

crates/simulation_message/src/lib.rs

@@ -221,6 +221,7 @@ pub enum GeomVariant {
     Box { extent: [f32; 3] },
     Plane { normal: [f32; 3] },
     Cylinder { radius: f32, half_height: f32 },
+    Capsule { radius: f32, half_height: f32 },
 }
 
 #[pymethods]
@@ -313,6 +314,27 @@ impl Geom {
             },
         }
     }
+
+    #[staticmethod]
+    pub fn capsule(
+        index: usize,
+        radius: f32,
+        half_height: f32,
+        material: Material,
+        pos: [f32; 3],
+        quat: [f32; 4],
+    ) -> Self {
+        Self {
+            index,
+            material,
+            pos,
+            quat,
+            geom_variant: GeomVariant::Capsule {
+                radius,
+                half_height,
+            },
+        }
+    }
 }
 
 #[pyclass(frozen)]

tools/mujoco-simulator/mujoco-rust-server/mujoco_rust_server/mujoco_rust_server.pyi

@@ -101,6 +101,15 @@ class Geom:
         pos: list[float],
         quat: list[float],
     ) -> Geom: ...
+    @staticmethod
+    def capsule(
+        index: int,
+        radius: float,
+        half_height: float,
+        material: Material,
+        pos: list[float],
+        quat: list[float],
+    ) -> Geom: ...
 
 class Light:
     name: str | None

tools/mujoco-simulator/mujoco-simulator/K1/K1.xml

@@ -50,7 +50,7 @@
             reflectance="0"
         />
 
-            <material
+        <material
             name="black_metal_rough"
             rgba="0.1 0.1 0.1 1"
             specular="0.8"
@@ -193,6 +193,21 @@
             pos="2 1 0"
             rgba="0.8 0.2 0.2 1"
         />
+        <geom
+            name="static_capsule"
+            type="capsule"
+            size="0.2 0.5"
+            pos="-2 1 0"
+            rgba="0.8 0.2 0.2 1"
+        />
+
+        <geom
+            name="static_cylinder"
+            type="cylinder"
+            size="0.2 0.5"
+            pos="-2 -1 0"
+            rgba="0.8 0.2 0.2 1"
+        />
 
         <body name="Trunk" pos="0 0 0.6">
             <site name='imu' size='0.01' pos='0.0 0.0 0.0' />
@@ -920,4 +935,4 @@
         <accelerometer name="accelerometer" site="imu" noise="0.005" />
     </sensor>
 
-</mujoco>
+</mujoco>

tools/mujoco-simulator/mujoco-simulator/mujoco_simulator/scene.py

@@ -1,4 +1,5 @@
 import logging
+from operator import ge
 
 import mujoco
 import numpy as np
@@ -83,7 +84,19 @@ def resolve_geom(model: MjModel, geom_index: int) -> Geom | None:
             quat=quat,
         )
 
-    logging.warning("Unhandled mujoco geom type:", geom_type)
+    if geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
+        radius: float = model.geom_size[geom_index][0]
+        half_height: float = model.geom_size[geom_index][1]
+        return Geom.capsule(
+            index=geom_index,
+            radius=radius,
+            half_height=half_height,
+            material=material,
+            pos=pos,
+            quat=quat,
+        )
+
+    logging.warning(f"Unhandled mujoco geom type: {geom_type.name}")
 
     return None