From e76acc01a31f5cad58f5dc37f31e7c646f7d0a23 Mon Sep 17 00:00:00 2001
From: Rooholla-KhorramBakht <learn.kho@gmail.com>
Date: Sat, 16 Mar 2024 01:25:53 -0400
Subject: [PATCH] kinematic and dynamic param APIs added to the mujocosim

---
 Go2Py/assets/mujoco/go2.xml | 16 +++----
 Go2Py/sim/mujoco.py         | 18 ++++++++
 examples/mujoco_test.ipynb  | 90 +++++++++++++++++++++++++++++++++++--
 3 files changed, 113 insertions(+), 11 deletions(-)

diff --git a/Go2Py/assets/mujoco/go2.xml b/Go2Py/assets/mujoco/go2.xml
index f155e12..5522c1c 100644
--- a/Go2Py/assets/mujoco/go2.xml
+++ b/Go2Py/assets/mujoco/go2.xml
@@ -103,7 +103,7 @@
             <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
             <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
             <geom name="FR" class="foot"/>
-            <site name="FR_foot_contact" pos="0 0 -0.213" size="0.02"/>
+            <site name="FR_foot" pos="0 0 -0.213" size="0.02"/>
           </body>
         </body>
       </body>
@@ -132,7 +132,7 @@
             <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
             <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
             <geom name="FL" class="foot"/>
-            <site name="FL_foot_contact" pos="0 0 -0.213" size="0.02"/>
+            <site name="FL_foot" pos="0 0 -0.213" size="0.02"/>
 
           </body>
         </body>
@@ -162,7 +162,7 @@
             <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
             <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
             <geom name="RR" class="foot"/>
-            <site name="RR_foot_contact" pos="0 0 -0.213" size="0.02"/>
+            <site name="RR_foot" pos="0 0 -0.213" size="0.02"/>
           </body>
         </body>
       </body>
@@ -191,7 +191,7 @@
             <geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder" class="collision"/>
             <geom pos="0 0 -0.213" mesh="foot" class="visual" material="black"/>
             <geom name="RL" class="foot"/>
-            <site name="RL_foot_contact" pos="0 0 -0.213" size="0.02"/>
+            <site name="RL_foot" pos="0 0 -0.213" size="0.02"/>
 
           </body>
         </body>
@@ -222,10 +222,10 @@
     <accelerometer name="imu_accel" site="imu"/>
     <gyro name="imu_gyro" site="imu"/>
     <!-- Contact sensors for the feet -->
-    <touch name="FR_foot_contact" site="FR_foot_contact"/>
-    <touch name="FL_foot_contact" site="RL_foot_contact"/>
-    <touch name="RR_foot_contact" site="RR_foot_contact"/>
-    <touch name="RL_foot_contact" site="RL_foot_contact"/>
+    <touch name="FR_foot_contact" site="FR_foot"/>
+    <touch name="FL_foot_contact" site="RL_foot"/>
+    <touch name="RR_foot_contact" site="RR_foot"/>
+    <touch name="RL_foot_contact" site="RL_foot"/>
   </sensor>
 
   <keyframe>
diff --git a/Go2Py/sim/mujoco.py b/Go2Py/sim/mujoco.py
index 8671601..fa3296a 100644
--- a/Go2Py/sim/mujoco.py
+++ b/Go2Py/sim/mujoco.py
@@ -45,6 +45,10 @@ class Go2Sim:
         self.reset()
         mujoco.mj_step(self.model, self.data)
         self.viewer.sync()
+        self.nv = self.model.nv
+        self.jacp = np.zeros((3, self.nv))
+        self.jacr = np.zeros((3, self.nv))
+        self.M = np.zeros((self.nv, self.nv))
 
     def reset(self):
         self.q_nominal = np.hstack(
@@ -103,5 +107,19 @@ class Go2Sim:
         if self.render and (self.step_counter%self.render_ds_ratio)==0:
             self.viewer.sync()
 
+    def getSiteJacobian(self, site_name):
+        id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE,site_name)
+        assert id>0, 'The requested site could not be found'
+        mujoco.mj_jacSite(self.model, self.data, self.jacp, self.jacr, id)
+        return self.jacp, self.jacr
+
+    def getDynamicsParams(self):
+        mujoco.mj_fullM(self.model, self.M, self.data.qM)
+        nle = self.data.qfrc_bias.reshape(self.nv,1)
+        return {
+            'M':self.M,
+            'nle':nle
+        }
+
     def close(self):
         self.viewer.close()
diff --git a/examples/mujoco_test.ipynb b/examples/mujoco_test.ipynb
index d9befa3..4e8857b 100644
--- a/examples/mujoco_test.ipynb
+++ b/examples/mujoco_test.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -43,8 +43,92 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "robot.getFootContact(), robot.getIMU()"
+    "n = robot.model.nv\n",
+    "full_M = np.zeros((n, n))"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import mujoco\n",
+    "nv = robot.model.nv\n",
+    "M = np.zeros((nv, nv))\n",
+    "mujoco.mj_fullM(robot.model, M, robot.data.qM)\n",
+    "nle = robot.data.qfrc_bias.reshape(nv,1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "id = mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_SITE,'FR_foot')\n",
+    "nv = robot.model.nv\n",
+    "jacp = np.zeros((3, nv))\n",
+    "jacr = np.zeros((3, nv))\n",
+    "mujoco.mj_jacSite(robot.model, robot.data, jacp, jacr, id)\n",
+    "jacp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(array([[ 1.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
+       "          0.00000000e+00, -3.03599143e-01,  1.39266782e-01,\n",
+       "          0.00000000e+00, -3.02750500e-01, -1.51075728e-01,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
+       "        [ 0.00000000e+00,  1.00000000e+00,  0.00000000e+00,\n",
+       "          3.03599143e-01,  0.00000000e+00,  1.94005151e-01,\n",
+       "          3.03599143e-01,  5.45551624e-06,  1.35362601e-03,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
+       "        [ 0.00000000e+00,  0.00000000e+00,  1.00000000e+00,\n",
+       "         -1.39266782e-01, -1.94005151e-01,  0.00000000e+00,\n",
+       "         -9.27667816e-02, -6.05126024e-04, -1.50144238e-01,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00,\n",
+       "          0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]),\n",
+       " array([[0.        , 0.        , 0.        , 1.        , 0.        ,\n",
+       "         0.        , 1.        , 0.        , 0.        , 0.        ,\n",
+       "         0.        , 0.        , 0.        , 0.        , 0.        ,\n",
+       "         0.        , 0.        , 0.        ],\n",
+       "        [0.        , 0.        , 0.        , 0.        , 1.        ,\n",
+       "         0.        , 0.        , 0.99995936, 0.99995936, 0.        ,\n",
+       "         0.        , 0.        , 0.        , 0.        , 0.        ,\n",
+       "         0.        , 0.        , 0.        ],\n",
+       "        [0.        , 0.        , 0.        , 0.        , 0.        ,\n",
+       "         1.        , 0.        , 0.00901514, 0.00901514, 0.        ,\n",
+       "         0.        , 0.        , 0.        , 0.        , 0.        ,\n",
+       "         0.        , 0.        , 0.        ]]))"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "robot.getSiteJacobian('FR_foot')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {