Feature/rollout (#15)

* feat: add rollout with custom integrator

* fix: correct rk4 implementation

* chore: rename

* chore: typo

* feat: implement manifold rk4

* fix: update comparison

* feat: add controls to rollout

darli/backend/base.py

@@ -259,3 +259,12 @@ def cone(
         self, force: ArrayLike | None, mu: float, type: str, X=None, Y=None
     ) -> ConeBase:
         pass
+
+    @abstractmethod
+    def integrate_configuration(
+        self,
+        dt: float,
+        q: ArrayLike | None = None,
+        v: ArrayLike | None = None,
+    ) -> ArrayLike:
+        ...

darli/modeling/base.py

@@ -347,11 +347,21 @@ def centroidal_dynamics(
     ) -> CentroidalDynamics:
         ...
 
+    @abstractmethod
+    def update(
+        self,
+        q: ArrayLike,
+        v: ArrayLike,
+        dv: ArrayLike | None = None,
+        u: ArrayLike | None = None,
+    ) -> ArrayLike:
+        ...
+
 
 class StateSpaceBase(ABC):
     @property
     @abstractmethod
-    def model(self):
+    def model(self) -> ModelBase:
         pass
 
     @property
@@ -364,9 +374,22 @@ def force_jacobians(self):
     def state(self):
         pass
 
-    @property
     @abstractmethod
-    def state_derivative(self):
+    def time_variation(
+        self,
+        q: ArrayLike | None = None,
+        v: ArrayLike | None = None,
+        u: ArrayLike | None = None,
+    ):
+        pass
+
+    @abstractmethod
+    def derivative(
+        self,
+        q: ArrayLike | None = None,
+        v: ArrayLike | None = None,
+        u: ArrayLike | None = None,
+    ):
         pass
 
     @property

darli/modeling/functional/model.py

@@ -369,3 +369,13 @@ def centroidal_dynamics(self) -> CentroidalDynamics:
         self.__centroidal = res
 
         return res
+
+    def update(
+        self,
+        q: ArrayLike,
+        v: ArrayLike,
+        dv: ArrayLike | None = None,
+        u: ArrayLike | None = None,
+    ) -> ArrayLike:
+        # dummy implementation to satisfy base class
+        return

darli/modeling/functional/state_space.py

@@ -1,6 +1,7 @@
 from ..state_space import CasadiStateSpace
-from ..base import ModelBase, StateSpaceBase
+from ..base import ModelBase, StateSpaceBase, ArrayLike
 import casadi as cs
+from ..integrators import Integrator, ForwardEuler
 
 
 class FunctionalStateSpace(StateSpaceBase):
@@ -35,23 +36,101 @@ def state(self):
         )
 
     @property
-    def state_derivative(self):
+    def derivative(self):
         return cs.Function(
-            "state_derivative",
+            "derivative",
             [
                 self.__space.model.q,
                 self.__space.model.v,
                 self.__space.model.qfrc_u,
                 *self.__space.model.contact_forces,
             ],
-            [self.__space.state_derivative],
+            [
+                self.__space.derivative(
+                    self.__space.model.q,
+                    self.__space.model.v,
+                    self.__space.model.qfrc_u,
+                )
+            ],
             [
                 "q",
                 "v",
                 "tau",
                 *self.__space.model.contact_names,
             ],
-            ["state_derivative"],
+            ["derivative"],
+        )
+
+    @property
+    def time_variation(self):
+        return cs.Function(
+            "time_variation",
+            [
+                self.__space.model.q,
+                self.__space.model.v,
+                self.__space.model.qfrc_u,
+                *self.__space.model.contact_forces,
+            ],
+            [
+                self.__space.time_variation(
+                    self.__space.model.q,
+                    self.__space.model.v,
+                    self.__space.model.qfrc_u,
+                )
+            ],
+            [
+                "q",
+                "v",
+                "tau",
+                *self.__space.model.contact_names,
+            ],
+            ["time_variation"],
+        )
+
+    def rollout(
+        self,
+        dt: float,
+        n_steps: int,
+        control_sampling: float | None = None,
+        integrator: Integrator = ForwardEuler,
+    ) -> cs.Function:
+        if control_sampling is None:
+            control_sampling = dt
+
+        control_dim = int(n_steps * dt / control_sampling)
+        container = self.__space.model.backend.math.zeros(
+            (self.__space.model.nu, control_dim)
+        ).array
+
+        for i in range(control_dim):
+            var = self.__space.model.backend.math.array(
+                f"control_{i}", self.__space.model.nu
+            ).array
+            container[:, i] = var
+
+        return cs.Function(
+            "rollout",
+            [
+                self.__space.state,
+                container,
+                *self.__space.model.contact_forces,
+            ],
+            [
+                self.__space.rollout(
+                    self.__space.state,
+                    container,
+                    dt,
+                    n_steps,
+                    control_sampling,
+                    integrator,
+                )
+            ],
+            [
+                "state",
+                "controls",
+                *self.__space.model.contact_names,
+            ],
+            ["next_state"],
         )
 
     @property

darli/modeling/integrators/__init__.py

@@ -0,0 +1,4 @@
+"""Module for various integrators"""
+from .base import Integrator
+from .fwd_euler import ForwardEuler
+from .rk4 import RK4

darli/modeling/integrators/base.py

@@ -0,0 +1,17 @@
+"""Module for base class of integrators"""
+
+from abc import ABC, abstractstaticmethod
+from ..base import StateSpaceBase
+from ...arrays import ArrayLike
+import casadi as cs
+
+
+class Integrator(ABC):
+    @abstractstaticmethod
+    def forward(
+        state_space: StateSpaceBase,
+        x0: ArrayLike,
+        qfrc_u: ArrayLike,
+        dt: float | cs.SX,
+    ) -> ArrayLike:
+        pass

darli/modeling/integrators/fwd_euler.py

@@ -0,0 +1,31 @@
+"""Module for forward euler integrator"""
+from .base import Integrator, StateSpaceBase, ArrayLike, cs
+
+
+class ForwardEuler(Integrator):
+    @staticmethod
+    def forward(
+        state_space: StateSpaceBase,
+        x0: ArrayLike,
+        qfrc_u: ArrayLike,
+        dt: cs.SX | float,
+    ):
+        nq = state_space.model.nq
+        nv = state_space.model.nv
+
+        q, v = x0[:nq], x0[nq:]
+        vdot = state_space.model.forward_dynamics(q, v, qfrc_u)
+
+        # forward euler to integrate velocity
+        integrated_v = v + dt * vdot
+
+        # pinocchio fancy lie-group integration
+        integrated_q = state_space.model.backend.integrate_configuration(
+            dt, q, integrated_v
+        )
+
+        container = state_space.model.backend.math.zeros(nq + nv).array
+        container[:nq] = integrated_q
+        container[nq:] = integrated_v
+
+        return container

darli/modeling/integrators/rk4.py

@@ -0,0 +1,43 @@
+"""Module for RK4 integrator"""
+from .base import Integrator, StateSpaceBase, ArrayLike, cs
+
+
+class RK4(Integrator):
+    @staticmethod
+    def forward(
+        state_space: StateSpaceBase,
+        x0: ArrayLike,
+        qfrc_u: ArrayLike,
+        h: cs.SX | float,
+    ):
+        nq = state_space.model.nq
+        nv = state_space.model.nv
+
+        def tangent_int(x, tang_x):
+            """tangent integration of state by its increment"""
+
+            # simplify notation for integration on manifold
+            manifold_int = state_space.model.backend.integrate_configuration
+
+            container = state_space.model.backend.math.zeros(nq + nv)
+
+            # configuration is integrated on manifold using pinocchio implementation
+            container[:nq] = manifold_int(h, x[:nq], tang_x[nv:])
+
+            # velocity and acceleration are in the same space and do not require specific treatment
+            container[nq:] = x[nq:] + h * tang_x[nv:]
+
+            return container.array
+
+        k1 = state_space.derivative(x0[:nq], x0[nq:], qfrc_u)
+
+        k2_ = tangent_int(x0, 0.5 * h * k1)
+        k2 = state_space.derivative(k2_[:nq], k2_[nq:], qfrc_u)
+
+        k3_ = tangent_int(x0, 0.5 * h * k2)
+        k3 = state_space.derivative(k3_[:nq], k3_[nq:], qfrc_u)
+
+        k4_ = tangent_int(x0, h * k3)
+        k4 = state_space.derivative(k4_[:nq], k4_[nq:], qfrc_u)
+
+        return tangent_int(x0, (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4))

darli/modeling/robot.py

@@ -121,7 +121,7 @@ def update(
         # update current dv and qfrc_u
         self._dv = self._backend._dv
         self._qfrc_u = self._backend._tau
-        if u:
+        if u is not None:
             self._u = u
 
         # update bodies

darli/modeling/state_space/casadi.py

@@ -16,21 +16,21 @@ def __init__(self, model: ModelBase) -> None:
     def state_jacobian(self):
         # if not free flyer
         if self.model.nq == self.model.nv:
-            return cs.jacobian(self.state_derivative, self.state)
+            return cs.jacobian(self.time_variation(), self.state)
 
         map = state_tangent_map(self.state, CasadiLikeFactory)
 
-        return map.T @ cs.jacobian(self.state_derivative, self.state) @ map
+        return map.T @ cs.jacobian(self.time_variation(), self.state) @ map
 
     @property
     def input_jacobian(self):
         # if not free flyer
         if self.model.nq == self.model.nv:
-            return cs.jacobian(self.state_derivative, self.model.qfrc_u)
+            return cs.jacobian(self.time_variation(), self.model.qfrc_u)
 
         map = state_tangent_map(self.state, CasadiLikeFactory)
 
-        return map.T @ cs.jacobian(self.state_derivative, self.model.qfrc_u)
+        return map.T @ cs.jacobian(self.time_variation(), self.model.qfrc_u)
 
     def force_jacobian(self, body_name: str) -> cs.Function:
         # early quit if we have already computed the jacobian
@@ -44,7 +44,7 @@ def force_jacobian(self, body_name: str) -> cs.Function:
         if self.model.bodies[body_name].contact is None:
             raise KeyError(f"Body {body_name} has no contact")
 
-        xdot = self.state_derivative
+        xdot = self.time_variation()
         body = self.model.body(body_name)
 
         map = state_tangent_map(self.state, CasadiLikeFactory)