Merge branch 'monica-branch' into rory-branch

Author: rrwwee

control.py

@@ -8,51 +8,125 @@
 
 import numpy as np
 
+from zmq.constants import THREAD_AFFINITY_CPU_REMOVE
 from bezier import Bezier
+from typing import List
+import pinocchio as pin
+from cube_trajectory_wrapper import GlobalMinJerkLinearJTraj
+from tools import setcubeplacement
+from config import LEFT_HAND, RIGHT_HAND
+
+import pybullet as pyb
 
 # in my solution these gains were good enough for all joints but you might want to tune this.
-Kp = 300.               # proportional gain (P of PD)
-Kv = 2 * np.sqrt(Kp)   # derivative gain (D of PD)
+Kp = 1000               # proportional gain (P of PD)
+Kv = 30 * np.sqrt(Kp)   # derivative gain (D of PD)
+
+Kx_lin = 2000  # Very stiff position control
+Dx_lin = 5 * np.sqrt(Kx_lin)
+
+Kx_rot = 1000
+Dx_rot = 10           
 
-def controllaw(sim, robot, trajs, tcurrent, cube):
+def controllaw(sim, robot, trajs, tcurrent):
+    """Joint trajectory tracking with end-effector force correction"""
     q, vq = sim.getpybulletstate()
 
     q_target = trajs[0](tcurrent)
     vq_target = trajs[1](tcurrent)
+    vvq_target = trajs[2](tcurrent)
+    
+    vvq_des = vvq_target + Kp * (q_target - q) + Kv * (vq_target - vq)
+    M = robot.mass(q)
+    h = robot.nle(q, vq)
+    tau_base = M @ vvq_des + h
 
+    cube_pos, cube_quat = pyb.getBasePositionAndOrientation(sim.cubeId)
+    R = np.array(pyb.getMatrixFromQuaternion(cube_quat)).reshape(3, 3)
+    cube_current_SE3 = pin.SE3(R, np.array(cube_pos))
 
-    error = q_target - q
-    vvq_des = Kp * error + Kv * vq_target
+    # cube_target_SE3 = pin.SE3(rotate('z', 0.), cube_trajs(tcurrent))
+    setcubeplacement(robot, cube, cube_current_SE3)
 
-    # data = robot.data
+    # Forward kinematics
+    model, data = robot.model, robot.data
+    left_id = model.getFrameId(LEFT_HAND)
+    right_id = model.getFrameId(RIGHT_HAND)
 
-    # Mass matrix
-    M = robot.mass(q)            
-    h = robot.nle(q, vq)
+    pin.forwardKinematics(model, data, trajs[0](0))
+    pin.updateFramePlacements(model, data)
 
-    torques = M @ vvq_des + h
+    oMLhand_desired = data.oMf[left_id]   # From joint trajectory
+    oMRhand_desired = data.oMf[right_id]
 
-    # print('Order: ', sim.bulletCtrlJointsInPinOrder)
-     
-    # torques = [torques[i] for i in sim.bulletCtrlJointsInPinOrder]
-    sim.step(torques)
+    
+    pin.forwardKinematics(model, data, q, vq)
+    pin.updateFramePlacements(model, data)
+    pin.computeJointJacobians(model, data, q)
+    
+    # Jacobians
+    J_l = pin.computeFrameJacobian(model, data, q, left_id, 
+                                    pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
+    J_r = pin.computeFrameJacobian(model, data, q, right_id,
+                                    pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
+                                    
+    oMLhand_actual = data.oMf[left_id]
+    oMRhand_actual = data.oMf[right_id]
 
-if __name__ == "__main__":
-        
-    from tools import setupwithpybullet, setupwithpybulletandmeshcat, rununtil
-    from config import DT
+    error_pos_hands = oMLhand_actual.translation - oMRhand_actual.translation
 
-    # `setupwithpybullet` currently returns (robot, sim, cube).
-    # Unpack and avoid depending on any `table` object.
-    robot, sim, cube = setupwithpybullet()
+    # Orientation errors
+    R_err_l = oMLhand_desired.rotation @ oMLhand_actual.rotation.T
+    err_rot_l = pin.log3(R_err_l)
+
+    R_err_r = oMRhand_desired.rotation @ oMRhand_actual.rotation.T
+    err_rot_r = pin.log3(R_err_r)
+
+    # Velocities
+    v_l = pin.getFrameVelocity(model, data, left_id, 
+                               pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
+    vel_l = np.hstack([v_l.linear, v_l.angular])
+    v_r = pin.getFrameVelocity(model, data, right_id,
+                               pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
+    vel_r = np.hstack([v_r.linear, v_r.angular])
 
+    # Compute correction forces
+    F_l = np.hstack([
+        -Kx_lin * error_pos_hands - Dx_lin * vel_l[:3],
+        Kx_rot * err_rot_l - Dx_rot * vel_l[3:]
+    ])
+    F_r = np.hstack([
+        Kx_lin * error_pos_hands - Dx_lin * vel_r[:3],
+        Kx_rot * err_rot_r - Dx_rot * vel_r[3:]
+    ])
 
-    from config import CUBE_PLACEMENT, CUBE_PLACEMENT_TARGET    
+    # Map to joint torques
+    tau_l = J_l[:, :].T @ F_l
+    tau_r = J_r[:, :].T @ F_r
+
+    tau_l = np.clip(tau_l, -200.0, 200.0)
+    tau_r = np.clip(tau_r, -200.0, 200.0)
+
+    # Combine: base tracking + end-effector correction
+    tau_total = tau_base + tau_r + tau_l
+    sim.step(tau_total)
+ 
+if __name__ == "__main__":
+        
+    from tools import setupwithpybullet, setupwithpybulletandmeshcat, rununtil
+    from config import CUBE_PLACEMENT, CUBE_PLACEMENT_TARGET, DT
     from inverse_geometry import computeqgrasppose
     from path import computepath
 
-    q0,successinit = computeqgrasppose(robot, robot.q0, cube, CUBE_PLACEMENT, None)
-    qe,successend = computeqgrasppose(robot, robot.q0, cube, CUBE_PLACEMENT_TARGET,  None)
+    robot, sim, table, cube = setupwithpybullet()
+    
+    q0, successinit = computeqgrasppose(robot, robot.q0, cube, CUBE_PLACEMENT, None)
+    qe, successend = computeqgrasppose(robot, robot.q0, cube, CUBE_PLACEMENT_TARGET,  None)
+
+    if not successinit or not successend:
+        print('Failed to compute grasp pose')
+        exit()
+
     path, cube_placements = computepath(
         q0,
         qe,
@@ -62,44 +136,52 @@ def controllaw(sim, robot, trajs, tcurrent, cube):
         cube=cube,
         computeqgrasppose=computeqgrasppose)
 
-
-    print('Success Init: ', successinit)
-    print('Success End: ', successend)
-
-    
-    #setting initial configuration
     sim.setqsim(q0)
-    
-    
-    #TODO this is just an example, you are free to do as you please.
-    #In any case this trajectory does not follow the path 
-    #0 init and end velocities
-    def maketraj(q0,q1,path, T): #TODO compute a real trajectory !
-        poimtlist = [q0, q0] + path + [q1, q1]
-        # point = [q0, q0, q1, q1]
-        q_of_t = Bezier(poimtlist,t_max=T)
+
+    def maketraj(q0, q1, path, T): #TODO compute a real trajectory !
+        points = [q0, q0] + path + [q1, q1]
+        q_of_t = Bezier(points,t_max=T)
         vq_of_t = q_of_t.derivative(1)
         vvq_of_t = vq_of_t.derivative(1)
         return q_of_t, vq_of_t, vvq_of_t
+
+
+    def maketraj_with_joint_space_linear(
+            waypoints: List[np.array],
+            T: int):
+        joint_space_trajectory = GlobalMinJerkLinearJTraj(
+            waypoints=waypoints,
+            T_max=T,
+        )
+
+        return (
+            joint_space_trajectory,
+            joint_space_trajectory.derivative(1),
+            joint_space_trajectory.derivative(2)
+        )
 
-    
-    #TODO this is just a random trajectory, you need to do this yourself
     total_time=10.
-    trajs = maketraj(q0, qe, path, total_time)   
-    
+
+    trajs = maketraj_with_joint_space_linear(
+        waypoints=path,
+        T=total_time
+    )
+
+    cube_trajs = maketraj_with_joint_space_linear(
+        waypoints=[p.translation for p in cube_placements],
+        T=total_time
+    )[0]
+
     tcur = 0.
 
-    
     while tcur < total_time:
         rununtil(
             controllaw, 
             DT, 
             sim, 
             robot, 
             trajs, 
-            tcur, 
-            cube)
+            tcur)
         tcur += DT
-    
-    
+