Track size/scale separately for renderables.

PiperOrigin-RevId: 918476985
Change-Id: I414b1abe2f9b3a3a7549ab5195971127d5598c36

src/experimental/filament/compat/scene_geom_util.cc

@@ -57,29 +57,21 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
   std::memcpy(position, &geom.pos, 3 * sizeof(float));
   float rotation[9];
   std::memcpy(rotation, &geom.mat, 9 * sizeof(float));
-  float size[3];
-  std::memcpy(size, &geom.size, 3 * sizeof(float));
 
   const mjtGeom geom_type = (mjtGeom)geom.type;
   switch (geom_type) {
     case mjGEOM_MESH:
     case mjGEOM_SDF:
       mjrf_setRenderableMesh(renderable, model_objs->GetMesh(geom.dataid), 0, 0);
-      // Ignore size for meshes.
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_HFIELD:
       mjrf_setRenderableMesh(renderable, model_objs->GetHeightField(geom.dataid), 0, 0);
-      // Ignore size for meshes.
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_PLANE: {
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
 
+      float size[3];
+      std::memcpy(size, &geom.size, 3 * sizeof(float));
       const bool is_infinite = !(size[0] > 0 && size[1] > 0);
       if (is_infinite) {
         // Infinite planes are scaled to match the tile size used by
@@ -90,40 +82,52 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       }
       // Planes only define an xy size, so set the z-dimension to 1.0f.
       size[2] = 1.0f;
+      mjrf_setRenderableSize(renderable, size);
       break;
     }
     case mjGEOM_SPHERE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ELLIPSOID:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_BOX:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_CAPSULE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_CYLINDER:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ARROW:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ARROW1:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ARROW2:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_LINE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_LINEBOX:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_TRIANGLE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_FLEX:
       mjrf_setRenderableMesh(renderable, model_objs->GetFlexSkinMesh(geom.objid), 0, 0);
@@ -133,9 +137,6 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       rotation[0] = 1.f;
       rotation[4] = 1.f;
       rotation[8] = 1.f;
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_SKIN:
       mjrf_setRenderableMesh(renderable, model_objs->GetFlexSkinMesh(geom.objid), 0, 0);
@@ -145,9 +146,6 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       rotation[0] = 1.f;
       rotation[4] = 1.f;
       rotation[8] = 1.f;
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_NONE:
     case mjGEOM_LABEL:
@@ -158,7 +156,7 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       break;
   }
 
-  mjrf_setRenderableTransform(renderable, position, rotation, size);
+  mjrf_setRenderableTransform(renderable, position, rotation);
 }
 
 static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,

src/experimental/filament/filament/renderable.cc

@@ -22,13 +22,13 @@
 #include <filament/RenderableManager.h>
 #include <filament/Scene.h>
 #include <filament/TransformManager.h>
+#include <math/mat3.h>
 #include <math/mat4.h>
 #include <math/vec2.h>
 #include <math/vec3.h>
 #include <math/vec4.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/builtins.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
@@ -40,6 +40,7 @@ namespace mujoco {
 using filament::math::float2;
 using filament::math::float3;
 using filament::math::float4;
+using filament::math::mat3f;
 using filament::math::mat4f;
 
 // An arbitrary scale factor for arrows.
@@ -149,21 +150,33 @@ void Renderable::InitPartEntity(Part& part) {
   }
 }
 
-void Renderable::SetTransform(const Trs& trs) {
+void Renderable::SetTransform(const float3& position, const mat3f& rotation) {
+  trs_.translation = position;
+  trs_.rotation = rotation;
+  UpdateTransform();
+}
+
+void Renderable::SetSize(const float3& size) {
+  trs_.size = size;
+  UpdateTransform();
+}
+
+void Renderable::UpdateTransform() {
   if (parts_.empty()) {
-    transform_ = trs.ToTransform();
+    transform_ = trs_.ToTransform();
     return;
   }
 
   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);
+      const mat4f transform = get_transform_fn_(i, trs_);
       tm.setTransform(tm.getInstance(parts_[i].entity), transform);
     }
   } else {
+    const mat4f transform = trs_.ToTransform();
     for (Part& part : parts_) {
-      tm.setTransform(tm.getInstance(part.entity), trs.ToTransform());
+      tm.setTransform(tm.getInstance(part.entity), transform);
     }
   }
   transform_ = tm.getTransform(tm.getInstance(parts_[0].entity));

src/experimental/filament/filament/renderable.h

@@ -21,10 +21,11 @@
 
 #include <filament/Engine.h>
 #include <filament/Scene.h>
+#include <math/mat3.h>
 #include <math/mat4.h>
+#include <math/vec3.h>
 #include <utils/Entity.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
 #include "experimental/filament/render_context_filament.h"
@@ -53,8 +54,14 @@ class Renderable : public mjrRenderable {
   // Sets the mesh of this renderable to a built-in mesh based on the geom type.
   void SetGeomMesh(mjtGeom type, int nstack, int nslice, int nquad);
 
-  // Sets the transform of this renderable.
-  void SetTransform(const Trs& trs);
+  // Sets the position and rotation of this renderable.
+  void SetTransform(const filament::math::float3& position,
+                    const filament::math::mat3f& rotation);
+
+  // Sets the size of this renderable. Note: this is effectively the same as a
+  // scale for most renderables. However, for e.g. capsules, the spherical ends
+  // are not scaled and remain fixed in size.
+  void SetSize(const filament::math::float3& size);
 
   // Returns the current transform of this renderable.
   const filament::math::mat4f& GetTransform() const;
@@ -114,12 +121,24 @@ class Renderable : public mjrRenderable {
     int elem_count = 0;
   };
 
+  // A tuple of translation, rotation, and size.
+  struct Trs {
+    filament::math::float3 translation{0.0f, 0.0f, 0.0f};
+    filament::math::mat3f rotation;
+    filament::math::float3 size{1.0f, 1.0f, 1.0f};
+    filament::math::mat4f ToTransform() const {
+      return filament::math::mat4f(rotation, translation) *
+            filament::math::mat4f::scaling(size);
+    }
+  };
+
   // When composing a multi-part renderable, each Entity will have its own
   // transform offset based on the transform of the Renderable itself.
   using GetTransformFn = std::function<filament::math::mat4f(int, const Trs&)>;
 
   void AppendMesh(const Mesh* mesh);
   void InitPartEntity(Part& part);
+  void UpdateTransform();
 
   void AssignMaterial(mjrDrawMode mode,
                       ObjectManager::MaterialType material_type);
@@ -136,6 +155,7 @@ class Renderable : public mjrRenderable {
   std::vector<Part> parts_;
   filament::math::mat4f transform_;
   GetTransformFn get_transform_fn_;
+  Trs trs_;
   bool wireframe_ = false;
 };
 

src/experimental/filament/filament_util.h

@@ -15,6 +15,7 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_UTIL_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_UTIL_H_
 
+#include <string_view>
 #include <math/mat3.h>
 #include <math/mat4.h>
 #include <math/vec2.h>
@@ -57,22 +58,6 @@ inline filament::math::mat3f ReadMat3(const T* arr, int index = 0) {
   // clang-format on
 }
 
-// A tuple of translation, rotation, and size.
-struct Trs {
-  filament::math::float3 translation{0.0f, 0.0f, 0.0f};
-  filament::math::mat3f rotation;
-  // Note: this is _slightly_ different than scale. For example, for capsules,
-  // the size determines the length of the tube and the radius of the domes,
-  // but the shape remains a capsule.
-  filament::math::float3 size{1.0f, 1.0f, 1.0f};
-
-  // Converts the TRS to a transform matrix.
-  filament::math::mat4f ToTransform() const {
-    return filament::math::mat4f(rotation, translation) *
-           filament::math::mat4f::scaling(size);
-  }
-};
-
 // Calculates a reflection matrix for a plane defined by its transform.
 filament::math::mat4 ToReflectionMatrix(const filament::math::mat4& xform);
 

src/experimental/filament/render_context_filament.cc

@@ -247,14 +247,18 @@ void mjrf_setRenderableMaterial(mjrRenderable* renderable,
 
 void mjrf_setRenderableTransform(mjrRenderable* renderable,
                                  const float position[3],
-                                 const float rotation[9], const float size[3]) {
+                                 const float rotation[9]) {
   const filament::math::float3 fposition{position[0], position[1], position[2]};
-  const filament::math::float3 fsize{size[0], size[1], size[2]};
   const filament::math::mat3f frotation{rotation[0], rotation[3], rotation[6],
                                         rotation[1], rotation[4], rotation[7],
                                         rotation[2], rotation[5], rotation[8]};
   mujoco::Renderable::downcast(renderable)
-      ->SetTransform({fposition, frotation, fsize});
+      ->SetTransform(fposition, frotation);
+}
+
+void mjrf_setRenderableSize(mjrRenderable* renderable, const float size[3]) {
+  const filament::math::float3 fsize{size[0], size[1], size[2]};
+  mujoco::Renderable::downcast(renderable)->SetSize(fsize);
 }
 
 void mjrf_setRenderableLayerMask(mjrRenderable* renderable,

src/experimental/filament/render_context_filament.h

@@ -695,12 +695,16 @@ void mjrf_setRenderableGeomMesh(mjrRenderable* renderable, mjtGeom type,
 void mjrf_setRenderableMaterial(mjrRenderable* renderable,
                                 const mjrMaterial* material);
 
-// Sets the transform (position, rotation, and size) of the renderable. Note
-// that `size` is not the same as `scale`. For example, the z-size of a capsule
-// only scales the tubular-portion of its geometry, but not the spherical caps.
+// Sets the transform position and rotation of the renderable.
 void mjrf_setRenderableTransform(mjrRenderable* renderable,
                                  const float position[3],
-                                 const float rotation[9], const float size[3]);
+                                 const float rotation[9]);
+
+// Sets the size of the renderable. Note that, for most renderables, this is
+// equivalent to setting the scale. However, for some geom-based renderables,
+// the size scale is not applied uniformly (e.g. the spherical ends of a
+// capsule are scaled such that they always remain spherical).
+void mjrf_setRenderableSize(mjrRenderable* renderable, const float size[3]);
 
 // Sets whether the renderable casts shadows or not.
 void mjrf_setRenderableCastShadows(mjrRenderable* renderable,

src/experimental/platform/ux/imgui_bridge.cc

@@ -270,10 +270,8 @@ void ImguiBridge::Update() {
       }
       mjrf_setRenderableMaterial(renderable.get(), &material);
 
-      const float position[] = {0, 0, 0};
-      const float rotation[] = {1, 0, 0, 0, 1, 0, 0, 0, 1};
       const float size[] = {scale.x, scale.y, 1.0f};
-      mjrf_setRenderableTransform(renderable.get(), position, rotation, size);
+      mjrf_setRenderableSize(renderable.get(), size);
 
       index_offset += command.ElemCount;
       ++renderable_index;