diff --git a/python/quadruped_reactive_walking/WB_MPC/ShootingNode.py b/python/quadruped_reactive_walking/WB_MPC/ShootingNode.py
index 4bff371a932e07a9b38c41dc1da8e7d1a55dc2e8..63f667862f171586b0ef03e1ac5393b63e1e286b 100644
--- a/python/quadruped_reactive_walking/WB_MPC/ShootingNode.py
+++ b/python/quadruped_reactive_walking/WB_MPC/ShootingNode.py
@@ -9,36 +9,76 @@ from .ProblemData import ProblemData
class ShootingNode:
+ """Shooting node (either running or terminal) of the optimal control problem
+
+ Args:
+ pd (ProblemData): data related to the OCP
+ contactIds (tuple): IDs of the feet frames in active contacts
+
+ Attributes:
+ pd (ProblemData): data related to the OCP
+ robotweight (float): total weight of the robot (mass * gravity)
+ contactIds (tuple): IDs of feet frames in active contacts
+ freeIds (list): : IDs of feet frames not in active contacts
+ baseID (int): ID of the base frame
+ forces (list): list of spatial forces for each joint
+ acc (function): compute acceleration with ABA
+ integrate (function): integration of a state vector
+ integrate_q (function): integration of a configuration vector
+ difference (function): difference between two state vectors
+ com (function): position of the center of mass
+ baseTranslation (function): position of the base frame
+ baseRotation (function): orientation of the base frame
+ baseVelocityLin (function): linear velocity of the base frame
+ baseVelocityAng (function): angular velocity of the base frame
+ feet (function): position of a foot frame
+ Rfeet (function): orientation of a foot frame
+ vfeet (function): linear velocity of a foot frame
+ afeet (function): linear acceleration of a foot frame
+ log3 (function): inverse of the exponential map
+ data (NodeData): data stored by each shooting node of the OCP
+ x (vector): State vector (q + dq)
+ a (vector): Acceleration slack variables (ddq)
+ u (vector): Control vector (tau)
+ tau (vector): Effort vector
+ fs (vector): Contact forces
+ isTerminal (bool): True if last node of shooting problem
+
+ """
+
def __init__(self, pd: ProblemData, contactIds):
self.pd = pd
- self.contactIds = contactIds
- self.freeIds = []
-
cmodel = cpin.Model(pd.model)
cdata = cmodel.createData()
- self.baseID = pd.model.getFrameId("base_link")
+ # Compute total weight of the robot
pin.framesForwardKinematics(pd.model, pd.rdata, pd.q0)
self.robotweight = (
-sum([Y.mass for Y in pd.model.inertias]) * pd.model.gravity.linear[2]
)
+
+ # Retrieve IDs of base and feet frames
+ self.baseID = pd.model.getFrameId("base_link")
+ self.contactIds = contactIds
+ self.freeIds = []
[self.freeIds.append(idf) for idf in pd.allContactIds if idf not in contactIds]
- cx = casadi.SX.sym("x", pd.nx, 1)
- cq = casadi.SX.sym("q", pd.nq, 1)
- cv = casadi.SX.sym("v", pd.nv, 1)
- cx2 = casadi.SX.sym("x2", pd.nx, 1)
- cu = casadi.SX.sym("u", pd.nu, 1)
- ca = casadi.SX.sym("a", pd.nv, 1)
- ctau = casadi.SX.sym("tau", pd.ntau, 1)
- cdx = casadi.SX.sym("dx", pd.ndx, 1)
+ # Declare symbolic variables
+ cx = casadi.SX.sym("x", pd.nx, 1) # State vector
+ cq = casadi.SX.sym("q", pd.nq, 1) # Configuration vector
+ cv = casadi.SX.sym("v", pd.nv, 1) # Velocity vector
+ cx2 = casadi.SX.sym("x2", pd.nx, 1) # Configuration vector bis (for difference)
+ cu = casadi.SX.sym("u", pd.nu, 1) # Control vector
+ ca = casadi.SX.sym("a", pd.nv, 1) # Acceleration vector
+ ctau = casadi.SX.sym("tau", pd.ntau, 1) # Effort vector
+ cdx = casadi.SX.sym("dx", pd.ndx, 1) # State derivative vector
cfs = [
casadi.SX.sym("f" + cmodel.frames[idf].name, 3, 1)
for idf in self.contactIds
- ]
- R = casadi.SX.sym("R", 3, 3)
- R_ref = casadi.SX.sym("R_ref", 3, 3)
+ ] # Contact forces vector
+ R = casadi.SX.sym("R", 3, 3) # Rotation matrix
+ R_ref = casadi.SX.sym("R_ref", 3, 3) # Reference rotation matrix
# Build force list for ABA
forces = [cpin.Force.Zero() for _ in cmodel.joints]
@@ -55,13 +95,14 @@ class ShootingNode:
for f in forces:
self.forces.append(f)
- acc = cpin.aba(cmodel, cdata, cx[: pd.nq], cx[pd.nq :], ctau, self.forces)
+ # Defining various casadi functions to wrap pinocchio methods
+ # (needed cause pinocchio.casadi methods cannot handle MX variables used in potimization)
- # Casadi Functions to wrap pinocchio methods (needed cause pinocchio.casadi methods can't handle MX variables used in potimization)
- # acceleration(x,u,f) = ABA(q,v,tau,f) with x=q,v, tau=u, and f built using StdVec_Force syntaxt
+ # acceleration(x,u,f) = ABA(q,v,tau,f) with x=q,v, tau=u, and f built using StdVec_Force syntax
+ acc = cpin.aba(cmodel, cdata, cx[: pd.nq], cx[pd.nq :], ctau, self.forces)
self.acc = casadi.Function("acc", [cx, ctau] + cfs, [acc])
- # integrate(x,dx) = [q+dq,v+dv], with the q+dq function implemented with pin.integrate.
+ # integrate(x,dx) = [q+dq,v+dv], with the q+dq function implemented with pin.integrate.
self.integrate = casadi.Function(
"xplus",
[cx, cdx],
@@ -90,7 +131,7 @@ class ShootingNode:
],
)
- # Reference frame in which velocities of base will be expressed (for locomotion)
+ # Reference frame in which base velocities will be expressed (for locomotion)
vel_reference_frame = pin.LOCAL
cpin.forwardKinematics(cmodel, cdata, cx[: pd.nq], cx[pd.nq :], ca)
@@ -173,6 +214,17 @@ class ShootingNode:
self.log3 = casadi.Function("log", [R, R_ref], [cpin.log3(R.T @ R_ref)])
def init(self, data, x, a=None, u=None, fs=None, isTerminal=False):
+ """Initialize the data of the shooting node
+
+ Args:
+ data (NodeData): Data stored by each shooting node of the OCP
+ x (vector): state vector (q + dq)
+ a (vector): acceleration slack variables (ddq)
+ u (vector): control vector (tau)
+ fs (vector): contact forces
+ isTerminal (bool): True if last node of shooting problem
+ """
+
self.data = data
self.data.x = x
self.data.a = a
@@ -188,8 +240,16 @@ class ShootingNode:
self.fs = fs
def calc(self, x_ref, u_ref=None, target={}):
- """
- This function return xnext,cost
+ """Compute the evolution of the state (shooting) and the associated cost
+
+ Args:
+ x_ref (vector): reference state vector
+ u_ref (vector): reference control vector
+ target (list of arrays): targets of the whole-body tasks (lin and ang velocities)
+
+ Returns:
+ xnext (vector): evolution of the state vector
+ cost (casadi.MX): cost of the node associated with the shooting
"""
dt = self.pd.dt
@@ -220,6 +280,12 @@ class ShootingNode:
return xnext, self.cost
def constraint_landing_feet_cost(self, x_ref):
+ """Formulate the costs associated with the landing constraints
+ that are related to touchdown altitude and touchdown velocity
+
+ Args:
+ xref (vector): reference state vector
+ """
eq = []
for stFoot in self.contactIds:
self.cost += (
@@ -238,6 +304,7 @@ class ShootingNode:
)
def constraint_standing_feet_eq(self):
+ """Formulate the acceleration equality constraint for feet in active contact"""
eq = []
for stFoot in self.contactIds:
eq.append(self.afeet[stFoot](self.x, self.a)) # stiff contact
@@ -245,6 +312,7 @@ class ShootingNode:
return casadi.vertcat(*eq)
def constraint_standing_feet_ineq(self):
+ """Formulate the friction cone inequality constraint for feet in active contact"""
# Friction cone
ineq = []
for i, stFoot in enumerate(self.contactIds):
@@ -260,6 +328,7 @@ class ShootingNode:
return casadi.vertcat(*ineq)
def constraint_standing_feet_cost(self):
+ """Formulate the costs associated with friction cone constraints"""
# Friction cone
for i, stFoot in enumerate(self.contactIds):
R = self.Rfeet[stFoot](self.x)
@@ -292,6 +361,7 @@ class ShootingNode:
)
def constraint_control_cost(self):
+ """Formulate the costs associated with control bound constraints"""
self.cost += (
self.pd.control_bound_w
* casadi.sumsqr(casadi.if_else(self.u <= -self.pd.effort_limit, self.u, 0))
@@ -306,11 +376,17 @@ class ShootingNode:
)
def constraint_dynamics_eq(self):
+ """Formulate the equality constraint of the equation of the dynamics"""
eq = []
eq.append(self.acc(self.x, self.tau, *self.fs) - self.a)
return casadi.vertcat(*eq)
def constraint_swing_feet_cost(self, x_ref):
+ """Formulate the cost to avoid collision of feet in swing phase
+
+ Args:
+ xref (vector): reference state vector
+ """
ineq = []
for sw_foot in self.freeIds:
r = self.feet[sw_foot](self.x)[2] - self.feet[sw_foot](x_ref)[2]
@@ -322,6 +398,7 @@ class ShootingNode:
# ineq.append(self.vfeet[sw_foot](self.x)[1] - k[1]* self.feet[sw_foot](self.x)[2])
def force_reg_cost(self):
+ """Formulate the force regularization cost"""
for i, stFoot in enumerate(self.contactIds):
R = self.Rfeet[stFoot](self.x)
f_ = self.fs[i]
@@ -333,11 +410,13 @@ class ShootingNode:
)
def control_cost(self, u_ref):
+ """Formulate the control regularization cost"""
self.cost += (
1 / 2 * self.pd.control_reg_w * casadi.sumsqr(self.u - u_ref) * self.pd.dt
)
def body_reg_cost(self, x_ref):
+ """Formulate the state regularization cost"""
if self.isTerminal:
self.cost += (
1
@@ -359,30 +438,45 @@ class ShootingNode:
* self.pd.dt
)
- def target_cost(self, task):
+ def target_cost(self, targets):
+ """Formulate the base velocity tracking cost
+
+ Args:
+ targets (list of arrays): targets of the whole-body tasks (lin and ang velocities)
+ """
self.cost += (
casadi.sumsqr(
self.pd.base_velocity_tracking_w
- * (self.baseVelocityLin(self.x) - task[0])
+ * (self.baseVelocityLin(self.x) - targets[0])
)
* self.pd.dt
)
self.cost += (
casadi.sumsqr(
self.pd.base_velocity_tracking_w
- * (self.baseVelocityAng(self.x) - task[1])
+ * (self.baseVelocityAng(self.x) - targets[1])
)
* self.pd.dt
)
- def compute_cost(self, x_ref, u_ref, target):
+ def compute_cost(self, x_ref, u_ref, targets):
+ """Formulate the total cost for a given node of the multiple shooting problem
+
+ Args:
+ x_ref (vector): reference state vector
+ u_ref (vector): reference control vector
+ targets (list of arrays): targets of the whole-body tasks (lin and ang velocities)
+
+ Returns:
+ cost (casadi.MX): cost of this node
+ """
self.cost = 0
if not self.isTerminal:
self.control_cost(u_ref)
self.constraint_control_cost()
self.constraint_standing_feet_cost()
self.constraint_swing_feet_cost(x_ref)
- self.target_cost(target)
+ self.target_cost(targets)
self.body_reg_cost(x_ref=x_ref)