Add MjcPhysicsSiteAPI support and testing.

PiperOrigin-RevId: 759124593
Change-Id: Ic2fa2a00f8384f432ee2eef95c80f5c180aa4ca9

src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc

@@ -21,6 +21,7 @@
 #include <vector>
 
 #include <mujoco/mujoco.h>
+#include "third_party/mujoco/src/experimental/usd/mjcPhysics/tokens.h"
 #include "mjcf/utils.h"
 #include <pxr/base/arch/attributes.h>
 #include <pxr/base/gf/matrix4d.h>
@@ -834,6 +835,8 @@ class ModelWriter {
     pxr::SdfPath site_path = WriteSiteGeom(site, body_path);
     SetPrimPurpose(data_, site_path, pxr::UsdGeomTokens->guide);
 
+    ApplyApiSchema(data_, site_path, pxr::MjcPhysicsTokens->SiteAPI);
+
     int site_id = mjs_getId(site->element);
     auto transform = MujocoPosQuatToTransform(&model_->site_pos[3 * site_id],
                                               &model_->site_quat[4 * site_id]);

test/experimental/usd/mjcPhysics/mjc_site_api_test.cc

@@ -0,0 +1,96 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <mujoco/experimental/usd.h>
+#include <mujoco/mujoco.h>
+#include "src/experimental/usd/mjcPhysics/siteAPI.h"
+#include "test/fixture.h"
+#include <pxr/base/tf/token.h>
+#include <pxr/usd/sdf/path.h>
+#include <pxr/usd/usd/common.h>
+#include <pxr/usd/usd/stage.h>
+#include <pxr/usd/usdGeom/capsule.h>
+#include <pxr/usd/usdGeom/cube.h>
+#include <pxr/usd/usdGeom/cylinder.h>
+#include <pxr/usd/usdGeom/sphere.h>
+#include <pxr/usd/usdGeom/xform.h>
+#include <pxr/usd/usdPhysics/collisionAPI.h>
+#include <pxr/usd/usdPhysics/rigidBodyAPI.h>
+
+#define EXPECT_SITE_TYPE(spec, site_path, expected_type)                  \
+  {                                                                       \
+    mjsElement* site_element =                                            \
+        mjs_findElement(spec, mjOBJ_SITE, site_path.GetString().c_str()); \
+    EXPECT_THAT(site_element, NotNull());                                 \
+    mjsSite* site = mjs_asSite(site_element);                             \
+    EXPECT_EQ(site->type, expected_type);                                 \
+  }
+
+namespace mujoco {
+namespace {
+
+using pxr::MjcPhysicsSiteAPI;
+using pxr::SdfPath;
+using MjcSiteApiTest = MujocoTest;
+using testing::NotNull;
+
+TEST_F(MjcSiteApiTest, TestApply) {
+  auto stage = pxr::UsdStage::CreateInMemory();
+
+  auto test_body_path = SdfPath("/World/TestBody");
+  auto body = pxr::UsdGeomXform::Define(stage, test_body_path);
+  pxr::UsdPhysicsRigidBodyAPI::Apply(body.GetPrim());
+
+  auto test_collider_path =
+      test_body_path.AppendChild(pxr::TfToken("Collider"));
+  auto collider = pxr::UsdGeomSphere::Define(stage, test_collider_path);
+  pxr::UsdPhysicsCollisionAPI::Apply(collider.GetPrim());
+
+  auto test_sphere_site_path =
+      test_body_path.AppendChild(pxr::TfToken("SphereSite"));
+  auto test_cylinder_site_path =
+      test_body_path.AppendChild(pxr::TfToken("CylinderSite"));
+  auto test_capsule_site_path =
+      test_body_path.AppendChild(pxr::TfToken("CapsuleSite"));
+  auto test_box_site_path = test_body_path.AppendChild(pxr::TfToken("BoxSite"));
+
+  auto sphere = pxr::UsdGeomSphere::Define(stage, test_sphere_site_path);
+  MjcPhysicsSiteAPI::Apply(sphere.GetPrim());
+
+  auto cylinder = pxr::UsdGeomCylinder::Define(stage, test_cylinder_site_path);
+  MjcPhysicsSiteAPI::Apply(cylinder.GetPrim());
+
+  auto capsule = pxr::UsdGeomCapsule::Define(stage, test_capsule_site_path);
+  MjcPhysicsSiteAPI::Apply(capsule.GetPrim());
+
+  auto box = pxr::UsdGeomCube::Define(stage, test_box_site_path);
+  MjcPhysicsSiteAPI::Apply(box.GetPrim());
+
+  mjSpec* spec = mj_parseUSDStage(stage);
+  mjModel* default_model = mj_compile(spec, nullptr);
+  EXPECT_THAT(default_model, NotNull()) << mjs_getError(spec);
+
+  EXPECT_SITE_TYPE(spec, test_sphere_site_path, mjGEOM_SPHERE);
+  EXPECT_SITE_TYPE(spec, test_cylinder_site_path, mjGEOM_CYLINDER);
+  EXPECT_SITE_TYPE(spec, test_capsule_site_path, mjGEOM_CAPSULE);
+  EXPECT_SITE_TYPE(spec, test_box_site_path, mjGEOM_BOX);
+
+  mj_deleteModel(default_model);
+  mj_deleteSpec(spec);
+}
+
+}  // namespace
+}  // namespace mujoco

test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc

@@ -17,6 +17,8 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "src/experimental/usd/mjcPhysics/sceneAPI.h"
+#include "src/experimental/usd/mjcPhysics/siteAPI.h"
 #include "test/experimental/usd/test_utils.h"
 #include "test/fixture.h"
 #include <pxr/base/gf/vec2f.h>
@@ -518,15 +520,28 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsAuthored) {
   auto stage = pxr::UsdStage::Open(layer);
   EXPECT_PRIM_VALID(stage, "/test/box_site");
   EXPECT_PRIM_IS_A(stage, "/test/box_site", pxr::UsdGeomCube);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/box_site", pxr::MjcPhysicsSiteAPI);
+
   EXPECT_PRIM_VALID(stage, "/test/ball/ball/sphere_site");
   EXPECT_PRIM_IS_A(stage, "/test/ball/ball/sphere_site", pxr::UsdGeomSphere);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/ball/ball/sphere_site",
+                          pxr::MjcPhysicsSiteAPI);
+
   EXPECT_PRIM_VALID(stage, "/test/ball/ball/capsule_site");
   EXPECT_PRIM_IS_A(stage, "/test/ball/ball/capsule_site", pxr::UsdGeomCapsule);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/ball/ball/capsule_site",
+                          pxr::MjcPhysicsSiteAPI);
+
   EXPECT_PRIM_VALID(stage, "/test/ball/ball/cylinder_site");
   EXPECT_PRIM_IS_A(stage, "/test/ball/ball/cylinder_site",
                    pxr::UsdGeomCylinder);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/ball/ball/cylinder_site",
+                          pxr::MjcPhysicsSiteAPI);
+
   EXPECT_PRIM_VALID(stage, "/test/ball/ball/ellipsoid_site");
   EXPECT_PRIM_IS_A(stage, "/test/ball/ball/ellipsoid_site", pxr::UsdGeomSphere);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/ball/ball/ellipsoid_site",
+                          pxr::MjcPhysicsSiteAPI);
 }
 
 TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsPurpose) {