expose link incoming joint force

haosulab · May 17, 2024 · 3037022 · 3037022
1 parent 55658e9
commit 3037022
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 0 deletions.
diff --git a/include/sapien/physx/articulation.h b/include/sapien/physx/articulation.h
@@ -39,6 +39,9 @@ class PhysxArticulation {
 
   Eigen::Matrix<float, Eigen::Dynamic, 2, Eigen::RowMajor> getQLimit();
 
+  Eigen::Matrix<float, Eigen::Dynamic, 6, Eigen::RowMajor> getLinkIncomingJointForces();
+
+
   Eigen::VectorXf computePassiveForce(bool gravity, bool coriolisAndCentrifugal);
 
   Pose getRootPose();

diff --git a/include/sapien/physx/physx_system.h b/include/sapien/physx/physx_system.h
@@ -189,6 +189,10 @@ class PhysxSystemGpu : public PhysxSystem {
   CudaArrayHandle gpuGetArticulationQTargetPosCudaHandle() const { return mCudaQTargetPosHandle; }
   CudaArrayHandle gpuGetArticulationQTargetVelCudaHandle() const { return mCudaQTargetVelHandle; }
 
+  CudaArrayHandle gpuGetArticulationLinkIncomingJointForceHandle() const {
+    return mCudaArticulationLinkIncomingJointForceBuffer.handle();
+  }
+
   void gpuFetchRigidDynamicData();
   void gpuFetchArticulationLinkPose();
   void gpuFetchArticulationLinkVel();
@@ -197,6 +201,7 @@ class PhysxSystemGpu : public PhysxSystem {
   void gpuFetchArticulationQacc();
   void gpuFetchArticulationQTargetPos();
   void gpuFetchArticulationQTargetVel();
+  void gpuFetchArticulationLinkIncomingJointForce();
 
   void gpuApplyRigidDynamicData(CudaArrayHandle const &indices);
   void gpuApplyArticulationRootPose(CudaArrayHandle const &indices);
@@ -300,6 +305,8 @@ class PhysxSystemGpu : public PhysxSystem {
   CudaArrayHandle mCudaQTargetPosHandle;
   CudaArrayHandle mCudaQTargetVelHandle;
 
+  CudaArray mCudaArticulationLinkIncomingJointForceBuffer;
+
   CudaArray mCudaContactBuffer;
   CudaArray mCudaContactCount;
 

diff --git a/python/py_package/pysapien/physx.pyi b/python/py_package/pysapien/physx.pyi
@@ -40,6 +40,8 @@ class PhysxArticulation:
         ...
     def get_joints(self) -> list[PhysxArticulationJoint]:
         ...
+    def get_link_incoming_joint_forces(self) -> numpy.ndarray[tuple[M, typing.Literal[6]], numpy.dtype[numpy.float32]]:
+        ...
     def get_links(self) -> list[PhysxArticulationLinkComponent]:
         ...
     def get_name(self) -> str:
@@ -111,6 +113,9 @@ class PhysxArticulation:
     def joints(self) -> list[PhysxArticulationJoint]:
         ...
     @property
+    def link_incoming_joint_forces(self) -> numpy.ndarray[tuple[M, typing.Literal[6]], numpy.dtype[numpy.float32]]:
+        ...
+    @property
     def links(self) -> list[PhysxArticulationLinkComponent]:
         ...
     @property
@@ -638,6 +643,8 @@ class PhysxGpuSystem(PhysxSystem):
         ...
     def gpu_create_contact_pair_impulse_query(self, body_pairs: list[tuple[PhysxRigidBaseComponent, PhysxRigidBaseComponent]]) -> PhysxGpuContactPairImpulseQuery:
         ...
+    def gpu_fetch_articulation_link_incoming_joint_forces(self) -> None:
+        ...
     def gpu_fetch_articulation_link_pose(self) -> None:
         ...
     def gpu_fetch_articulation_link_velocity(self) -> None:
@@ -705,6 +712,9 @@ class PhysxGpuSystem(PhysxSystem):
     def cuda_articulation_link_data(self) -> sapien.pysapien.CudaArray:
         ...
     @property
+    def cuda_articulation_link_incoming_joint_forces(self) -> sapien.pysapien.CudaArray:
+        ...
+    @property
     def cuda_articulation_qacc(self) -> sapien.pysapien.CudaArray:
         ...
     @property

diff --git a/python/pybind/physx.cpp b/python/pybind/physx.cpp
@@ -473,6 +473,8 @@ kernels will be launched to the provided stream.
                              &PhysxSystemGpu::gpuGetArticulationQTargetPosCudaHandle)
       .def_property_readonly("cuda_articulation_target_qvel",
                              &PhysxSystemGpu::gpuGetArticulationQTargetVelCudaHandle)
+      .def_property_readonly("cuda_articulation_link_incoming_joint_forces",
+                             &PhysxSystemGpu::gpuGetArticulationLinkIncomingJointForceHandle)
 
       .def("gpu_fetch_rigid_dynamic_data", &PhysxSystemGpu::gpuFetchRigidDynamicData)
       .def("gpu_fetch_articulation_link_pose", &PhysxSystemGpu::gpuFetchArticulationLinkPose)
@@ -482,6 +484,8 @@ kernels will be launched to the provided stream.
       .def("gpu_fetch_articulation_qacc", &PhysxSystemGpu::gpuFetchArticulationQacc)
       .def("gpu_fetch_articulation_target_qpos", &PhysxSystemGpu::gpuFetchArticulationQTargetPos)
       .def("gpu_fetch_articulation_target_qvel", &PhysxSystemGpu::gpuFetchArticulationQTargetVel)
+      .def("gpu_fetch_articulation_link_incoming_joint_forces",
+           &PhysxSystemGpu::gpuFetchArticulationLinkIncomingJointForce)
 
       .def("gpu_create_contact_pair_impulse_query",
            &PhysxSystemGpu::gpuCreateContactPairImpulseQuery, py::arg("body_pairs"))
@@ -994,6 +998,10 @@ set some motion axes of the dynamic rigid body to be locked
       .def("get_qlimit", &PhysxArticulation::getQLimit, R"doc(same as get_qlimit)doc")
       .def("get_qlimits", &PhysxArticulation::getQLimit)
 
+      .def_property_readonly("link_incoming_joint_forces",
+                             &PhysxArticulation::getLinkIncomingJointForces)
+      .def("get_link_incoming_joint_forces", &PhysxArticulation::getLinkIncomingJointForces)
+
       .def_property("root_pose", &PhysxArticulation::getRootPose, &PhysxArticulation::setRootPose)
       .def("get_root_pose", &PhysxArticulation::getRootPose)
       .def("set_root_pose", &PhysxArticulation::setRootPose, py::arg("pose"))

diff --git a/src/physx/articulation.cpp b/src/physx/articulation.cpp
@@ -326,6 +326,18 @@ Eigen::Matrix<float, Eigen::Dynamic, 2, Eigen::RowMajor> PhysxArticulation::getQ
   return m;
 }
 
+Eigen::Matrix<float, Eigen::Dynamic, 6, Eigen::RowMajor>
+PhysxArticulation::getLinkIncomingJointForces() {
+  if (getRoot()->isUsingDirectGPUAPI()) {
+    throw std::runtime_error("getting qf is not supported in GPU simulation.");
+  }
+  mPxArticulation->copyInternalStateToCache(*mCache,
+                                            PxArticulationCacheFlag::eLINK_INCOMING_JOINT_FORCE);
+  auto mat8 = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, 8, Eigen::RowMajor>>(
+      reinterpret_cast<float *>(mCache->linkIncomingJointForce), mPxArticulation->getNbLinks(), 8);
+  return mat8(Eigen::all, {0, 1, 2, 4, 5, 6});
+}
+
 Pose PhysxArticulation::getRootPose() {
   // if (getRoot()->isUsingDirectGPUAPI()) {
   //   throw std::runtime_error("getting root pose is not supported in GPU simulation.");

diff --git a/src/physx/physx_system.cpp b/src/physx/physx_system.cpp
@@ -642,6 +642,22 @@ void PhysxSystemGpu::gpuFetchArticulationQTargetVel() {
   mCudaEventWait.wait(mCudaStream);
 }
 
+void PhysxSystemGpu::gpuFetchArticulationLinkIncomingJointForce() {
+  checkGpuInitialized();
+
+  if (mGpuArticulationCount == 0) {
+    return;
+  }
+
+  ensureCudaDevice();
+
+  mPxScene->copyArticulationData(mCudaArticulationLinkIncomingJointForceBuffer.ptr,
+                                 mCudaArticulationIndexBuffer.ptr,
+                                 PxArticulationGpuDataType::eLINK_INCOMING_JOINT_FORCE,
+                                 mCudaArticulationIndexBuffer.shape.at(0), mCudaEventWait.event);
+  mCudaEventWait.wait(mCudaStream);
+}
+
 void PhysxSystemGpu::gpuFetchArticulationQacc() {
   checkGpuInitialized();
 
@@ -1000,6 +1016,9 @@ void PhysxSystemGpu::allocateCudaBuffers() {
                       .cudaId = mCudaRigidBodyBuffer.cudaId,
                       .ptr = (float *)mCudaRigidBodyBuffer.ptr + 13 * rigidDynamicCount};
 
+  mCudaArticulationLinkIncomingJointForceBuffer =
+      CudaArray({mGpuArticulationCount, mGpuArticulationMaxLinkCount, 6}, "f4");
+
   // rigid body force torque buffer
   mCudaRigidBodyForceBuffer = CudaArray({rigidBodyCount, 4}, "f4");
   mCudaRigidDynamicForceHandle = CudaArrayHandle{.shape = {rigidDynamicCount, 4},