From b788f8ad2aea655ade37e957aa6a2978596297ad Mon Sep 17 00:00:00 2001
From: Rooholla-Khorrambakht <learn.kho@gmail.com>
Date: Wed, 3 Apr 2024 02:24:17 -0400
Subject: [PATCH] Ground force reaction estimator

---
 Go2Py/robot/model.py            |  11 +--
 Go2Py/sim/mujoco.py             |   5 +-
 examples/contact_detector.ipynb | 126 ++++++++++++++++++++++++++++++++
 3 files changed, 136 insertions(+), 6 deletions(-)
 create mode 100644 examples/contact_detector.ipynb

diff --git a/Go2Py/robot/model.py b/Go2Py/robot/model.py
index 196f2ae..d0c328f 100644
--- a/Go2Py/robot/model.py
+++ b/Go2Py/robot/model.py
@@ -56,11 +56,11 @@ class Go2Model:
 
         print(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_RR_HAA].translation)
 
-        print(self.h)
-        print(self.b)
-        print(self.l1)
-        print(self.l2)
-        print(self.l3)
+        # print(self.h)
+        # print(self.b)
+        # print(self.l1)
+        # print(self.l2)
+        # print(self.l3)
     
     def inverseKinematics(self, T, feet_pos):
         """
@@ -156,6 +156,7 @@ class Go2Model:
         self.Minv_ = pin.computeMinverse(self.robot.model, self.robot.data, q_)[self.dq_reordering_idx,:]
         self.Minv_ = self.Minv_[:,self.dq_reordering_idx]
         for ef_frame in self.ef_frames:
+            # J = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
             J = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL)
             self.ef_J_[ef_frame]=J[:, self.dq_reordering_idx]
 
diff --git a/Go2Py/sim/mujoco.py b/Go2Py/sim/mujoco.py
index dbca0ae..478ecd7 100644
--- a/Go2Py/sim/mujoco.py
+++ b/Go2Py/sim/mujoco.py
@@ -69,6 +69,7 @@ class Go2Sim:
         self.kp = np.zeros(12)
         self.kv = np.zeros(12)
         self.latest_command_stamp = time.time()
+        self.actuator_tau = np.zeros(12)
 
     def reset(self):
         self.q_nominal = np.hstack(
@@ -95,7 +96,8 @@ class Go2Sim:
 
     def getJointStates(self):
         return {"q":self.data.qpos[7:], 
-               "dq":self.data.qvel[6:]}
+               "dq":self.data.qvel[6:],
+               'tau_est':self.actuator_tau}
 
     def getPose(self):
         return self.data.qpos[:3], self.data.qpos[3:7]
@@ -122,6 +124,7 @@ class Go2Sim:
         q, dq = state['q'], state['dq']
         tau = np.diag(self.kp)@(self.q_des-q).reshape(12,1)+ \
               np.diag(self.kv)@(self.dq_des-dq).reshape(12,1)+self.tau_ff.reshape(12,1)
+        self.actuator_tau = tau
         self.data.ctrl[:] = tau.squeeze()
 
         self.step_counter += 1
diff --git a/examples/contact_detector.ipynb b/examples/contact_detector.ipynb
new file mode 100644
index 0000000..cb0de85
--- /dev/null
+++ b/examples/contact_detector.ipynb
@@ -0,0 +1,126 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from Go2Py.sim.mujoco import Go2Sim\n",
+    "from Go2Py.robot.model import Go2Model\n",
+    "import pinocchio as pin \n",
+    "import numpy as np\n",
+    "import time"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = Go2Model()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "robot = Go2Sim()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "robot.step()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "robot.standUpReset()\n",
+    "while True:\n",
+    "    time.sleep(0.001)\n",
+    "    robot.setCommands(robot.standing_q, np.zeros(12), np.ones(12)*100, np.ones(12)*0.1, np.zeros(12))\n",
+    "    robot.step()    \n",
+    "    trans, quat = robot.getPose()\n",
+    "    state = robot.getJointStates()\n",
+    "    quat = pin.Quaternion(np.hstack([quat[1:],quat[0]]))\n",
+    "    R = quat.matrix()\n",
+    "    T = np.eye(4)\n",
+    "    T[0:3,0:3] = R\n",
+    "    vel = np.zeros(6)\n",
+    "    model.update(state['q'], state['dq'],T,vel)\n",
+    "\n",
+    "    J = model.getInfo()['J']['FR_foot'][:,6:]\n",
+    "    nle = model.getInfo()['nle'][6:]\n",
+    "    # tau = (robot.data.qfrc_smooth+robot.data.qfrc_constraint)[6:]\n",
+    "    tau = state['tau_est'].squeeze()\n",
+    "    print(np.linalg.pinv(J.T)[0:3]@(tau - nle))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "question: What exactly is the difference between the local world aligned and local jacobians in the pinocchio? Why it fits my expectations for the FR3 and not the Go2?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "robot.getJointStates()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "J[3:,:].T"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}