Use function for applying transforms to multimesh Renderables.

PiperOrigin-RevId: 902712241
Change-Id: I5bfdab8bce12ffb1a1f030ae0a32a61e9b8157e1

src/experimental/filament/filament/renderable.cc

@@ -135,11 +135,23 @@ void Renderable::InitPartEntity(Part& part) {
 }
 
 void Renderable::SetTransform(const Trs& trs) {
-  transform_ = trs.ToTransform();
+  if (parts_.empty()) {
-  filament::TransformManager& tm = GetEngine()->getTransformManager();
+    transform_ = trs.ToTransform();
-  for (Part& part : parts_) {
+    return;
-    tm.setTransform(tm.getInstance(part.entity), transform_);
   }
+
+  filament::TransformManager& tm = GetEngine()->getTransformManager();
+  if (get_transform_fn_) {
+    for (int i = 0; i < parts_.size(); ++i) {
+      const mat4f& transform = get_transform_fn_(i, trs);
+      tm.setTransform(tm.getInstance(parts_[i].entity), transform);
+    }
+  } else {
+    for (Part& part : parts_) {
+      tm.setTransform(tm.getInstance(part.entity), trs.ToTransform());
+    }
+  }
+  transform_ = tm.getTransform(tm.getInstance(parts_[0].entity));
 }
 
 const mat4f& Renderable::GetTransform() const {
@@ -147,23 +159,18 @@ const mat4f& Renderable::GetTransform() const {
 }
 
 void Renderable::SetMeshes(std::span<const Mesh*> meshes,
-                           std::span<const mat4f> transforms) {
+                           GetTransformFn get_transform_fn) {
-  if (meshes.size() != transforms.size()) {
-    mju_error("Number of meshes does not match number of transforms.");
-  }
   if (!parts_.empty()) {
     mju_error("Cannot set meshes for renderable with multiple parts.");
   }
 
-  filament::TransformManager& tm = GetEngine()->getTransformManager();
+  get_transform_fn_ = get_transform_fn;
   for (int i = 0; i < meshes.size(); ++i) {
     Part& part = parts_.emplace_back();
     part.mesh = meshes[i];
     part.elem_offset = 0;
     part.elem_count = part.mesh->GetFilamentIndexBuffer()->getIndexCount();
     InitPartEntity(part);
-
-    tm.setTransform(tm.getInstance(part.entity), transforms[i]);
   }
 }
 

src/experimental/filament/filament/renderable.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_RENDERABLE_H_
 
 #include <cstdint>
+#include <functional>
 #include <span>
 #include <vector>
 
@@ -83,10 +84,13 @@ class Renderable {
   // Returns the current transform of the renderable.
   const filament::math::mat4f& GetTransform() const;
 
-  // Sets multiple meshes for a renderable. Each mesh is assigned a specific
+  // Sets multiple meshes for a renderable. Users can optionally provide a
-  // transform to allow for assembly of compound shapes.
+  // function that will be used to compute the transform for each (sub)mesh
+  // relative to the transform of the renderable itself. This allows users to
+  // construct compound (but rigid) objects from multiple meshes.
+  using GetTransformFn = std::function<filament::math::mat4f(int, const Trs&)>;
   void SetMeshes(std::span<const Mesh*> meshes,
-                 std::span<const filament::math::mat4f> transforms);
+                 GetTransformFn get_transform = nullptr);
 
   // Sets the layer mask for the managed filament Entities. Layer masks can be
   // used to show/hide the renderable in different views. Returns the previous
@@ -158,7 +162,7 @@ class Renderable {
   filament::Scene* assigned_scene_ = nullptr;
   std::vector<Part> parts_;
   filament::math::mat4f transform_;
-
+  GetTransformFn get_transform_fn_;
   std::uint8_t priority_ = kDefaultPriority;
   std::uint8_t layer_mask_ = kDefaultLayerMask;
   std::uint16_t blend_order_ = 0;

src/experimental/filament/filament/scene_geom_util.cc

@@ -103,7 +103,7 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
                               const mjvScene* scene,
                               ModelObjects* model_objects) {
   std::vector<const Mesh*> meshes;
-  std::vector<mat4f> transforms;
+  Renderable::GetTransformFn get_transforms;
 
   Trs trs = {
       .translation = ReadFloat3(geom.pos),
@@ -115,12 +115,12 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
     case mjGEOM_MESH:
       meshes.push_back(GetMesh(model_objects, geom.dataid));
       // Ignore size for meshes.
-      transforms.push_back(mat4f(trs.rotation, trs.translation));
+      trs.size = float3{1.0f, 1.0f, 1.0f};
       break;
     case mjGEOM_HFIELD:
       meshes.push_back(GetHeightField(model_objects, geom.dataid));
       // Ignore size for height fields.
-      transforms.push_back(mat4f(trs.rotation, trs.translation));
+      trs.size = float3{1.0f, 1.0f, 1.0f};
       break;
     case mjGEOM_PLANE: {
       meshes.push_back(GetShape(model_objects, ModelObjects::kPlane));
@@ -134,20 +134,16 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
       }
       // Planes only define an xy size, so set the z-dimension to 1.0f.
       trs.size.z = 1.0f;
-      transforms.push_back(trs.ToTransform());
       break;
     }
     case mjGEOM_SPHERE:
       meshes.push_back(GetShape(model_objects, ModelObjects::kSphere));
-      transforms.push_back(trs.ToTransform());
       break;
     case mjGEOM_ELLIPSOID:
       meshes.push_back(GetShape(model_objects, ModelObjects::kSphere));
-      transforms.push_back(trs.ToTransform());
       break;
     case mjGEOM_BOX:
       meshes.push_back(GetShape(model_objects, ModelObjects::kBox));
-      transforms.push_back(trs.ToTransform());
       break;
     case mjGEOM_CAPSULE: {
       // Capsules are a tube with two domes at the ends.
@@ -155,25 +151,31 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
       meshes.push_back(GetShape(model_objects, ModelObjects::kDome));
       meshes.push_back(GetShape(model_objects, ModelObjects::kDome));
 
-      transforms.push_back(trs.ToTransform());
+      get_transforms = [](int index, const Trs& trs) {
-
+        // We apply an inverse scale to the domes to counteract the capsule's
-      // We apply an inverse scale to the domes to counteract the capsule's
+        // overall scale so that the domes remain spherical in shape.
-      // overall scale so that the domes remain spherical in shape.
+        const float xz_size = 0.5f * (trs.size.x + trs.size.y);
-      const float xz_size = 0.5f * (trs.size.x + trs.size.y);
+        if (index == 0) {
-
+          return trs.ToTransform();
-      // Move the first dome to the top of the capsule.
+        } else if (index == 1) {
-      mat4f top = mat4f(trs.rotation, trs.translation);
+          // Move the first dome to the top of the capsule.
-      top *= mat4f::translation(float3{0, 0, trs.size.z});
+          mat4f top = mat4f(trs.rotation, trs.translation);
-      top *= mat4f::scaling(float3{trs.size.x, trs.size.y, xz_size});
+          top *= mat4f::translation(float3{0, 0, trs.size.z});
-      transforms.push_back(top);
+          top *= mat4f::scaling(float3{trs.size.x, trs.size.y, xz_size});
-
+          return top;
-      // Move the second dome to the bottom of the capsule and rotate it 180
+        } else if (index == 2) {
-      // degrees so that it's facing the right way.
+          // Move the second dome to the bottom of the capsule and rotate it 180
-      mat4f bottom = mat4f(trs.rotation, trs.translation);
+          // degrees so that it's facing the right way.
-      bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
+          mat4f bottom = mat4f(trs.rotation, trs.translation);
-      bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
-      bottom *= mat4f::scaling(float3{trs.size.x, trs.size.y, xz_size});
+          bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
-      transforms.push_back(bottom);
+          bottom *= mat4f::scaling(float3{trs.size.x, trs.size.y, xz_size});
+          return bottom;
+        } else {
+          mju_error("Invalid index for capsule geom: %d (expected [0,2])", index);
+          return trs.ToTransform();
+        }
+      };
       break;
     }
     case mjGEOM_CYLINDER: {
@@ -182,21 +184,28 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
       meshes.push_back(GetShape(model_objects, ModelObjects::kDisk));
       meshes.push_back(GetShape(model_objects, ModelObjects::kDisk));
 
-      transforms.push_back(trs.ToTransform());
+      get_transforms = [](int index, const Trs& trs) {
-
+        if (index == 0) {
-      // Move the first disk to the top of the cylinder.
+          return trs.ToTransform();
-      mat4f top = mat4f(trs.rotation, trs.translation);
+        } else if (index == 1) {
-      top *= mat4f::translation(float3{0, 0, trs.size.z});
+          // Move the first disk to the top of the cylinder.
-      top *= mat4f::scaling(trs.size);
+          mat4f top = mat4f(trs.rotation, trs.translation);
-      transforms.push_back(top);
+          top *= mat4f::translation(float3{0, 0, trs.size.z});
-
+          top *= mat4f::scaling(trs.size);
-      // Move the second disk to the bottom of the cylinder. Rotate the disk
+          return top;
-      // 180 degrees so that the normals point outwards.
+        } else if (index == 2) {
-      mat4f bottom = mat4f(trs.rotation, trs.translation);
+          // Move the second disk to the bottom of the cylinder. Rotate the disk
-      bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
+          // 180 degrees so that the normals point outwards.
-      bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          mat4f bottom = mat4f(trs.rotation, trs.translation);
-      bottom *= mat4f::scaling(trs.size);
+          bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
-      transforms.push_back(bottom);
+          bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          bottom *= mat4f::scaling(trs.size);
+          return bottom;
+        } else {
+          mju_error("Invalid index for cylinder geom: %d (expected [0,2])", index);
+          return trs.ToTransform();
+        }
+      };
       break;
     }
     case mjGEOM_ARROW: {
@@ -205,27 +214,33 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
       meshes.push_back(GetShape(model_objects, ModelObjects::kDisk));
       meshes.push_back(GetShape(model_objects, ModelObjects::kDisk));
 
-      mat4f base = mat4f(trs.rotation, trs.translation);
+      get_transforms = [](int index, const Trs& trs) {
-      base *= mat4f::scaling(float3{1, 1, kArrowScale});
+        mat4f base = mat4f(trs.rotation, trs.translation);
-      base *= mat4f::translation(float3{0, 0, trs.size.z});
+        base *= mat4f::scaling(float3{1, 1, kArrowScale});
-      transforms.push_back(base * mat4f::scaling(trs.size));
+        base *= mat4f::translation(float3{0, 0, trs.size.z});
-
+        if (index == 0) {
-      mat4f top = base;
+          return base * mat4f::scaling(trs.size);
-      top *= mat4f::translation(float3{0, 0, trs.size.z});
+        } else if (index == 1) {
-      top *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
+          mat4f top = base;
-      transforms.push_back(top * mat4f::scaling(trs.size));
+          top *= mat4f::translation(float3{0, 0, trs.size.z});
-
+          top *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
-      mat4f top_disk = base;
+          return top * mat4f::scaling(trs.size);
-      top_disk *= mat4f::translation(float3{0, 0, trs.size.z});
+        } else if (index == 2) {
-      top_disk *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          mat4f top_disk = base;
-      top_disk *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
+          top_disk *= mat4f::translation(float3{0, 0, trs.size.z});
-      transforms.push_back(top_disk * mat4f::scaling(trs.size));
+          top_disk *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
-
+          top_disk *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
-      mat4f bottom = base;
+          return top_disk * mat4f::scaling(trs.size);
-      bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
+        } else if (index == 3) {
-      bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          mat4f bottom = base;
-      transforms.push_back(bottom * mat4f::scaling(trs.size));
+          bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
-
+          bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          return bottom * mat4f::scaling(trs.size);
+        } else {
+          mju_error("Invalid index for arrow geom: %d (expected [0,3])", index);
+          return trs.ToTransform();
+        }
+      };
       break;
     }
     case mjGEOM_ARROW1: {
@@ -233,19 +248,26 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
       meshes.push_back(GetShape(model_objects, ModelObjects::kCone));
       meshes.push_back(GetShape(model_objects, ModelObjects::kDisk));
 
+      get_transforms = [](int index, const Trs& trs) {
       mat4f base = mat4f(trs.rotation, trs.translation);
       base *= mat4f::scaling(float3{1, 1, kArrowScale});
       base *= mat4f::translation(float3{0, 0, trs.size.z});
-      transforms.push_back(base * mat4f::scaling(trs.size));
+        if (index == 0) {
-
+          return base * mat4f::scaling(trs.size);
-      mat4f top = base;
+        } else if (index == 1) {
-      top *= mat4f::translation(float3{0, 0, trs.size.z});
+          mat4f top = base;
-      transforms.push_back(top * mat4f::scaling(trs.size));
+          top *= mat4f::translation(float3{0, 0, trs.size.z});
-
+          return top * mat4f::scaling(trs.size);
-      mat4f bottom = base;
+        } else if (index == 2) {
-      bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
+          mat4f bottom = base;
-      bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
-      transforms.push_back(bottom * mat4f::scaling(trs.size));
+          bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          return bottom * mat4f::scaling(trs.size);
+        } else {
+          mju_error("Invalid index for arrow1 geom: %d (expected [0,2])", index);
+          return trs.ToTransform();
+        }
+      };
       break;
     }
     case mjGEOM_ARROW2: {
@@ -255,55 +277,58 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
       meshes.push_back(GetShape(model_objects, ModelObjects::kDisk));
       meshes.push_back(GetShape(model_objects, ModelObjects::kDisk));
 
-      mat4f base = mat4f(trs.rotation, trs.translation);
+      get_transforms = [](int index, const Trs& trs) {
-      base *= mat4f::scaling(float3{1, 1, kArrowScale});
+        mat4f base = mat4f(trs.rotation, trs.translation);
-      base *= mat4f::translation(float3{0, 0, trs.size.z});
+        base *= mat4f::scaling(float3{1, 1, kArrowScale});
-      transforms.push_back(base * mat4f::scaling(trs.size));
+        base *= mat4f::translation(float3{0, 0, trs.size.z});
-
+        if (index == 0) {
-      mat4f top = base;
+          return base * mat4f::scaling(trs.size);
-      top *= mat4f::translation(float3{0, 0, trs.size.z});
+        } else if (index == 1) {
-      top *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
+          mat4f top = base;
-      transforms.push_back(top * mat4f::scaling(trs.size));
+          top *= mat4f::translation(float3{0, 0, trs.size.z});
-
+          top *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
-      mat4f bottom = base;
+          return top * mat4f::scaling(trs.size);
-      bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
+        } else if (index == 2) {
-      bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          mat4f bottom = base;
-      bottom *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
+          bottom *= mat4f::translation(float3{0, 0, -trs.size.z});
-      transforms.push_back(bottom * mat4f::scaling(trs.size));
+          bottom *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
-
+          bottom *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
-      mat4f top_disk = base;
+          return bottom * mat4f::scaling(trs.size);
-      top_disk *= mat4f::translation(float3{0, 0, trs.size.z});
+        } else if (index == 3) {
-      top_disk *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
+          mat4f top_disk = base;
-      top_disk *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
+          top_disk *= mat4f::translation(float3{0, 0, trs.size.z});
-      transforms.push_back(top_disk * mat4f::scaling(trs.size));
+          top_disk *= mat4f::rotation(std::numbers::pi, float3{1, 0, 0});
-
+          top_disk *= mat4f::scaling(float3{kArrowHeadSize, kArrowHeadSize, 1.0f});
-      mat4f bottom_disk = base;
+          return top_disk * mat4f::scaling(trs.size);
-      bottom_disk *= mat4f::translation(float3{0, 0, -trs.size.z});
+        } else if (index == 4) {
-      transforms.push_back(bottom_disk * mat4f::scaling(trs.size));
+          mat4f bottom_disk = base;
-
+          bottom_disk *= mat4f::translation(float3{0, 0, -trs.size.z});
+          return bottom_disk * mat4f::scaling(trs.size);
+        } else {
+          mju_error("Invalid index for arrow2 geom: %d (expected [0,4])", index);
+          return trs.ToTransform();
+        }
+      };
       break;
     }
     case mjGEOM_LINE:
       meshes.push_back(GetShape(model_objects, ModelObjects::kLine));
-      transforms.push_back(trs.ToTransform());
       break;
     case mjGEOM_LINEBOX:
       meshes.push_back(GetShape(model_objects, ModelObjects::kLineBox));
-      transforms.push_back(trs.ToTransform());
       break;
     case mjGEOM_TRIANGLE:
       meshes.push_back(GetShape(model_objects, ModelObjects::kTriangle));
-      transforms.push_back(trs.ToTransform());
       break;
     case mjGEOM_FLEX:
       meshes.push_back(GetSkinFlexMesh(model_objects, geom.objid));
       // Flexes are defined in global space.
-      transforms.push_back(mat4f());
+      trs = Trs();
       break;
     case mjGEOM_SKIN:
       meshes.push_back(GetSkinFlexMesh(model_objects, geom.objid));
       // Skins are defined in global space.
-      transforms.push_back(mat4f());
+      trs = Trs();
       break;
     case mjGEOM_NONE:
     case mjGEOM_LABEL:
@@ -315,7 +340,8 @@ static void PrepareGeomMeshes(Renderable& renderable, const mjvGeom& geom,
       break;
   }
 
-  renderable.SetMeshes(meshes, transforms);
+  renderable.SetMeshes(meshes, get_transforms);
+  renderable.SetTransform(trs);
 }
 
 static void UpdateGeomMaterial(Renderable& renderable, const mjvGeom& geom,