From 1f8f3a500439668f7e4b5fb55cd6e7bccaa248c8 Mon Sep 17 00:00:00 2001
From: Robin Alazard <robinalazard@google.com>
Date: Thu, 10 Jul 2025 08:31:15 -0700
Subject: [PATCH] Add capsule axis rotation for USD->MJ.

PiperOrigin-RevId: 781546484
Change-Id: Ie6613de8faac406664ca22cb246fe7db2a728317
---
 src/experimental/usd/usd_to_mjspec.cc | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/experimental/usd/usd_to_mjspec.cc b/src/experimental/usd/usd_to_mjspec.cc
index f9134c98..5bf7855b 100644
--- a/src/experimental/usd/usd_to_mjspec.cc
+++ b/src/experimental/usd/usd_to_mjspec.cc
@@ -35,6 +35,7 @@
 #include "experimental/usd/kinematic_tree.h"
 #include <pxr/base/gf/declare.h>
 #include <pxr/base/gf/matrix4d.h>
+#include <pxr/base/gf/rotation.h>
 #include <pxr/base/gf/vec3d.h>
 #include <pxr/base/tf/token.h>
 #include <pxr/base/vt/types.h>
@@ -193,6 +194,25 @@ bool MaybeParseGeomPrimitive(const pxr::UsdPrim& prim, T* element,
     element->size[0] = scale[0] * radius;
     element->size[1] = scale[1] * height / 2.0f;
     element->size[2] = 0;
+
+    TfToken axis;
+    capsule.GetAxisAttr().Get(&axis);
+
+    // Mujoco (and USD) capsules are aligned with Z by default.
+    // When USD axis is X or Y, we apply a rotation to align with the Z axis.
+    pxr::GfQuatd axis_rot(1.0);
+    if (axis == pxr::UsdGeomTokens->x) {
+      axis_rot = pxr::GfRotation(pxr::GfVec3d::XAxis(), pxr::GfVec3d::ZAxis())
+                     .GetQuat();
+    } else if (axis == pxr::UsdGeomTokens->y) {
+      axis_rot = pxr::GfRotation(pxr::GfVec3d::YAxis(), pxr::GfVec3d::ZAxis())
+                     .GetQuat();
+    }
+
+    pxr::GfQuatd current_rot(element->quat[0], element->quat[1],
+                             element->quat[2], element->quat[3]);
+    pxr::GfQuatd new_rot = current_rot * axis_rot;
+    SetDoubleArrFromGfQuatd(element->quat, new_rot);
   } else if (prim.IsA<pxr::UsdGeomCube>()) {
     element->type = mjGEOM_BOX;
     auto cube = pxr::UsdGeomCube(prim);
@@ -1452,4 +1472,3 @@ mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) {
 
   return spec;
 }
-