NMPC Swing Up Controller based on Crocoddyl (#86)

* add NMPC swing up controller using crocoddyl

* tune to make it work with the new urdf

* nmpc swing up : working in simulation in closed loop at 100Hz on measured data. Rm support for constraints for now while waiting for mim_solvers package.

* cast rwd

* urdf : absolute path to urdf

* logs : absolute path and create log dir in repo if it doesn't exist yet

* pid : rm config file and read json utils

* isort

---------

Co-authored-by: Armandpl <adpl33@gmail.com>

Author: pierfabre

.gitignore

@@ -7,6 +7,7 @@
 **/replay_buffer/*
 
 data/*
+logs/
 
 # ignore everything except notebooks in the notebook folder
 notebooks/*

furuta/controls/controllers.py

@@ -1,45 +1,148 @@
+import typing as tp
 from abc import ABC, abstractmethod
 
+import crocoddyl
 import numpy as np
+import pinocchio as pin
 from simple_pid import PID
 
 
 class Controller(ABC):
     @abstractmethod
-    def compute_command(self, position: float):
+    def compute_command(self, state):
         pass
 
-    @staticmethod
-    def build_controller(parameters: dict):
-        controller_type = parameters["controller_type"]
-        # match controller_type:
-        #     case "PIDController":
-        #         return PIDController(parameters)
-        #     case _:
-        #         raise ValueError(f"Invalid controller type: {controller_type}")
-        if controller_type == "PIDController":
-            return PIDController(parameters)
-        else:
-            raise ValueError(f"Invalid controller type: {controller_type}")
-
 
 class PIDController(Controller):
-    def __init__(self, parameters):
-        try:
-            sample_time = 1 / parameters["control_frequency"]
-        except KeyError:
-            sample_time = None
-
-        try:
-            self.pid = PID(
-                Kp=parameters["Kp"],
-                Ki=parameters["Ki"],
-                Kd=parameters["Kd"],
-                setpoint=np.deg2rad(parameters["setpoint"]),
-                sample_time=sample_time,
-            )
-        except KeyError:
-            raise ValueError("Invalid PID parameters")
+    def __init__(
+        self,
+        dt: float = None,
+        Kp: float = 0.0,
+        Ki: float = 0.0,
+        Kd: float = 0.0,
+        setpoint: float = 0.0,
+    ):
+
+        self.pid = PID(
+            Kp=Kp,
+            Ki=Ki,
+            Kd=Kd,
+            setpoint=setpoint,
+            sample_time=dt,
+        )
 
     def compute_command(self, position: float):
         return self.pid(position)
+
+
+class SwingUpController(Controller):
+    class FurutaActuationModel(crocoddyl.ActuationModelAbstract):
+        def __init__(self, state, nu):
+            crocoddyl.ActuationModelAbstract.__init__(self, state, nu=nu)
+
+        def calc(self, data, x, u):
+            assert len(data.tau) == 2
+            # Map the control dimensions to the joint torque
+            data.tau[0] = u
+            data.tau[1] = 0  # Underactuated joint
+
+        def calcDiff(self, data, x, u):
+            # Specify the actuation jacobian
+            data.dtau_du[0] = 1
+            data.dtau_du[1] = 0
+
+    def __init__(
+        self,
+        robot: pin.RobotWrapper,
+        x_target: np.ndarray,
+        control_freq: float = 100.0,
+        t_final: float = 0.5,
+        u_lim: float = 0.1,
+        Q: np.ndarray = np.array([10, 50, 1, 1]),
+        R: np.ndarray = np.array([0.1]),
+        S: np.ndarray = np.array([1.0]),
+        M: int = 10,
+    ):
+        self.u_lim = u_lim
+
+        # Time variables
+        dt = 1 / control_freq  # Time step
+        self.N = int(t_final * control_freq)
+
+        # Instantiate robot as a pinocchio RobotWrapper
+        self.robot = robot
+
+        # State
+        state = crocoddyl.StateMultibody(robot.model)
+
+        # Actuation model
+        nu = 1
+        actuation = self.FurutaActuationModel(state, nu)
+
+        # State Cost
+        state_residual = crocoddyl.ResidualModelState(state, xref=x_target, nu=nu)
+        state_residual_activation = crocoddyl.ActivationModelWeightedQuad(Q)
+        state_cost = crocoddyl.CostModelResidual(state, state_residual_activation, state_residual)
+
+        # Control Cost
+        control_residual = crocoddyl.ResidualModelControl(state, nu=nu)
+        control_residual_activation = crocoddyl.ActivationModelWeightedQuad(R)
+        control_cost = crocoddyl.CostModelResidual(
+            state, control_residual_activation, control_residual
+        )
+
+        # Control rate cost
+        self.control_rate_residual = crocoddyl.ResidualModelControl(state, uref=np.array([0.0]))
+        control_rate_residual_activation = crocoddyl.ActivationModelWeightedQuad(S)
+        control_rate_cost = crocoddyl.CostModelResidual(
+            state, control_rate_residual_activation, self.control_rate_residual
+        )
+
+        self.running_models = []
+        for k in range(self.N):
+            running_cost = crocoddyl.CostModelSum(state, nu=nu)
+            running_cost.addCost("state_cost", cost=state_cost, weight=1.0)
+            running_cost.addCost("control_cost", cost=control_cost, weight=1.0)
+            running_cost.addCost("control_rate_cost", cost=control_rate_cost, weight=np.exp(M - k))
+
+            running_model = crocoddyl.IntegratedActionModelEuler(
+                crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, running_cost),
+                dt,
+            )
+            self.running_models.append(running_model)
+
+        # Terminal cost
+        terminal_cost = crocoddyl.CostModelSum(state, nu=nu)
+        terminal_cost.addCost("state_cost", cost=state_cost, weight=self.N)
+
+        self.terminal_model = crocoddyl.IntegratedActionModelEuler(
+            crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, terminal_cost),
+            0.0,
+        )
+
+    def create_problem(self, state: np.ndarray):
+        self.problem = crocoddyl.ShootingProblem(state, self.running_models, self.terminal_model)
+
+    def init_solver(self):
+        self.solver = crocoddyl.SolverFDDP(self.problem)
+        callbacks = []
+        callbacks.append(crocoddyl.CallbackVerbose())
+        self.solver.setCallbacks(callbacks)
+
+    def compute_command(self, state: np.ndarray, max_iter: int = 500, x_ws=[], u_ws=[]) -> float:
+        self.create_problem(state)
+        self.init_solver()
+        self.solver.solve(x_ws, u_ws, max_iter, False, 1e-5)
+        u = np.clip(self.solver.us[0][0], -self.u_lim, self.u_lim)
+        return u
+
+    def get_trajectoy(self) -> np.ndarray:
+        return self.solver.xs.tolist()
+
+    def get_command(self) -> np.ndarray:
+        return self.solver.us.tolist()
+
+    def get_warm_start(self) -> tp.Tuple[np.ndarray, np.ndarray]:
+        x_ws = self.solver.xs.tolist()[1:] + [self.solver.xs[-1]]
+        u_ws = self.solver.us.tolist()[1:] + [self.solver.us[-1]]
+        return x_ws, u_ws

furuta/controls/utils.py

@@ -124,7 +124,7 @@ def step(self, action):
         terminated = not self.state_space.contains(self._state)
         truncated = False
 
-        return obs, rwd, terminated, truncated, {}
+        return obs, float(rwd), terminated, truncated, {}
 
     def get_obs(self):
         obs = np.float32(

furuta/rl/envs/furuta_base.py

@@ -6,15 +6,15 @@
 import pinocchio as pin
 import serial
 
+import furuta
+
 
 class RobotModel:
     def __init__(self):
-        self.robot = None
-        base_path = Path("robot/hardware/v2/")
-        if base_path.exists():
-            urdf_path = str(base_path / "robot.urdf")
-            stls_path = str(base_path / "stl")
-            self.robot = pin.RobotWrapper.BuildFromURDF(urdf_path, [stls_path])
+        base_path = Path(furuta.__path__[0]).parent / "robot" / "hardware" / "v2"
+        urdf_path = base_path / "robot.urdf"
+        stls_path = base_path / "stl"
+        self.robot = pin.RobotWrapper.BuildFromURDF(urdf_path, [stls_path])
 
 
 class Robot:

furuta/robot.py

@@ -21,17 +21,16 @@ def display(self, state: np.ndarray) -> np.ndarray:
     def close(self):
         pass
 
-    def animate(self, times: np.ndarray, states: np.ndarray):
+    def animate(self, times: np.ndarray, states: np.ndarray | list[np.ndarray]):
         # Initial state
-        q = states[:, :2]
-        self.display(q[0])
+        self.display(states[0][:2])
         time.sleep(1.0)
         tic = time.time()
         for i in range(1, len(times)):
             toc = time.time()
             time.sleep(max(0, times[i] - times[i - 1] - (toc - tic)))
             tic = time.time()
-            self.display(q[i])
+            self.display(states[i][:2])
 
 
 class Viewer3D(AbstractViewer):

furuta/viewer.py

@@ -18,7 +18,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "ipykernel>=6.26.0",
-    "matplotlib==3.7.3",
+    "matplotlib>=3.7.3",
     "pre-commit>=3.5.0",
     "stable-baselines3>=2.2.1",
     "wandb>=0.16.2",
@@ -30,8 +30,7 @@ dev = [
     "jax-metal>=0.0.5; sys_platform == 'darwin'",
     "sb3-contrib",
     "pytest>=7.4.4",
-    "pin>=2.7.0",
-    "crocoddyl>=2.0.2",
+    "crocoddyl>=3.0.0",
     "onshape-to-robot>=0.3.26",
     "meshcat",
 ]