Predictive Control API

ambersim/control/base.py

@@ -0,0 +1,32 @@
+import jax
+from flax import struct
+
+
+@struct.dataclass
+class ControllerParams:
+    """The parameters for generic controllers.
+
+    This is left completely empty for maximum flexibility in the API. Some examples:
+        - "Regular" inputs into feedback controllers (e.g., the state) belong here.
+        - Non-Markovian controllers can pass histories in this params object.
+        - Parameters of the controller that you may randomize/optimize go here.
+    """
+
+
+@struct.dataclass
+class Controller:
+    """The API for a generic controller.
+
+    See the notes in TrajectoryOptimizer on the generality of this class - much of the same applies.
+    """
+
+    def compute(self, ctrl_params: ControllerParams) -> jax.Array:
+        """Computes a control input.
+
+        Args:
+            ctrl_params: ControllerParams
+
+        Returns:
+            u (shape=(nu,)): The control input.
+        """
+        raise NotImplementedError

ambersim/control/predictive_control.py

@@ -0,0 +1,149 @@
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+from flax import struct
+from mujoco import mjx
+
+from ambersim.control.base import Controller, ControllerParams
+from ambersim.trajopt.base import TrajectoryOptimizer
+from ambersim.trajopt.shooting import PDPredictiveSamplerParams, PredictiveSampler, VanillaPredictiveSamplerParams
+
+# ########### #
+# GENERIC API #
+# ########### #
+
+
+@struct.dataclass
+class PredictiveControllerParams(ControllerParams):
+    """The generic API for predictive controller params."""
+
+
+@struct.dataclass
+class PredictiveController(Controller):
+    """The generic API for a predictive controller."""
+
+    trajectory_optimizer: TrajectoryOptimizer
+    model: mjx.Model
+
+    def compute(self, ctrl_params: PredictiveControllerParams) -> jax.Array:
+        """Computes a control input using forward prediction."""
+        raise NotImplementedError
+
+
+# ################### #
+# PREDICTIVE SAMPLING #
+# ################### #
+
+
+@struct.dataclass
+class VanillaPredictiveSamplingControllerParams(PredictiveControllerParams):
+    """Vanilla predictive sampling controller params."""
+
+    key: jax.Array  # random key for sampling
+    x: jax.Array  # shape=(nq + nv,) current state
+    guess: jax.Array  # shape=(N, nu) current guess
+
+
+@struct.dataclass
+class VanillaPredictiveSamplingController(PredictiveController):
+    """Vanilla predictive sampling controller."""
+
+    def __post_init__(self) -> None:
+        """Post-initialization check."""
+        assert isinstance(
+            self.trajectory_optimizer, PredictiveSampler
+        ), "trajectory_optimizer must be a PredictiveSampler!"
+
+    def compute(self, ctrl_params: VanillaPredictiveSamplingControllerParams) -> jax.Array:
+        """Computes a control input using forward prediction.
+
+        Args:
+            ctrl_params: Inputs into the controller.
+
+        Returns:
+            u (shape=(nu,)): The control input.
+        """
+        return self.compute_with_us_star(ctrl_params)[0]
+
+    def compute_with_us_star(
+        self, ctrl_params: VanillaPredictiveSamplingControllerParams
+    ) -> Tuple[jax.Array, jax.Array]:
+        """Computes a control input using forward prediction + the optimal sequence of guesses.
+
+        This is needed in practice because the current optimal sequence is used to warm start the sampling distribution
+        for the next call of the controller.
+
+        Args:
+            ctrl_params: Inputs into the controller.
+
+        Returns:
+            u (shape=(nu,)): The control input.
+            us_star (shape=(N, nu)): The optimal control sequence.
+        """
+        to_params = VanillaPredictiveSamplerParams(
+            key=ctrl_params.key,
+            x0=ctrl_params.x,
+            guess=ctrl_params.guess,
+        )
+        xs_star, us_star = self.trajectory_optimizer.optimize(to_params)
+        u = us_star[0, :]
+        return u, us_star
+
+
+@struct.dataclass
+class PDPredictiveSamplingControllerParams(PredictiveControllerParams):
+    """PD predictive sampling controller params."""
+
+    key: jax.Array  # random key for sampling
+    x: jax.Array  # shape=(nq + nv,) current state
+    guess: jax.Array  # shape=(N, nq) current guess
+    kp: float  # proportional gain
+    kd: float  # derivative gain
+
+
+@struct.dataclass
+class PDPredictiveSamplingController(PredictiveController):
+    """PD predictive sampling controller."""
+
+    def __post_init__(self) -> None:
+        """Post-initialization check."""
+        assert isinstance(
+            self.trajectory_optimizer, PredictiveSampler
+        ), "trajectory_optimizer must be a PredictiveSampler!"
+
+    def compute(self, ctrl_params: PDPredictiveSamplingControllerParams) -> jax.Array:
+        """Computes a control input using forward prediction.
+
+        Args:
+            ctrl_params: Inputs into the controller.
+
+        Returns:
+            u (shape=(nu,)): The control input.
+        """
+        return self.compute_with_qs_star(ctrl_params)[0]
+
+    def compute_with_qs_star(self, ctrl_params: PDPredictiveSamplingControllerParams) -> Tuple[jax.Array, jax.Array]:
+        """Computes a control input using forward prediction + the optimal sequence of guesses.
+
+        This is needed in practice because the current optimal sequence is used to warm start the sampling distribution
+        for the next call of the controller.
+
+        Args:
+            ctrl_params: Inputs into the controller.
+
+        Returns:
+            u (shape=(nu,)): The control input.
+            us_star (shape=(N, nu)): The optimal control sequence.
+        """
+        to_params = PDPredictiveSamplerParams(
+            key=ctrl_params.key,
+            x0=ctrl_params.x,
+            guess=ctrl_params.guess,
+            kp=ctrl_params.kp,
+            kd=ctrl_params.kd,
+        )
+        xs_star, us_star = self.trajectory_optimizer.optimize(to_params)
+        # u = us_star[0, :]
+        qs_star = xs_star[:, : self.model.nq]  # the 0th index is the current state, so return the 1st index
+        return qs_star[1, :], qs_star

ambersim/models/allegro_hand/cube.xml

@@ -0,0 +1,29 @@
+<mujoco>
+  <asset>
+    <texture name="cube" type="cube"
+             fileup="assets/fileup.png" fileback="assets/fileback.png"
+             filedown="assets/filedown.png" filefront="assets/filefront.png"
+             fileleft="assets/fileleft.png" fileright="assets/fileright.png"/>
+    <material name="cube" texture="cube"/>
+    <texture name="graycube" type="cube" fileup="assets/grayup.png"
+             fileback="assets/grayback.png" filedown="assets/graydown.png"
+             filefront="assets/grayfront.png" fileleft="assets/grayleft.png"
+             fileright="assets/grayright.png"/>
+    <material name="graycube" texture="graycube"/>
+  </asset>
+  <worldbody>
+    <light pos="0 0 1"/>
+    <body name="cube" pos="0.0 0.0 0.035" quat="1 0 0 0">
+      <freejoint/>
+      <geom name="cube" type="box" size=".022 .022 .022" mass=".126" material="cube"/>
+    </body>
+  </worldbody>
+
+   <sensor>
+        <framepos name="trace0" objtype="body" objname="cube"/>
+        <framepos name="cube_position" objtype="body" objname="cube"/>
+        <framequat name="cube_orientation" objtype="body" objname="cube"/>
+        <framelinvel name="cube_linear_velocity" objtype="body" objname="cube"/>
+        <frameangvel name="cube_angular_velocity" objtype="body" objname="cube"/>
+    </sensor>
+</mujoco>