raddo 0.1.4__py3-none-any.whl

raddo/controllers/__init__.py

@@ -0,0 +1,7 @@
+from .base_controller import Controller
+from .cem_controller import CEMController
+from .feedback_controller import FeedbackController
+from .mppi_controller import MPPIController
+from .qp_controller import QPController
+from .sqp_controller import SQPController
+from .tvlqr_controller import TVLQRController

raddo/controllers/acados_controller.py

@@ -0,0 +1,356 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import casadi as cs
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+from acados_template import AcadosModel, AcadosOcp, AcadosOcpSolver
+from jax.typing import ArrayLike
+
+if TYPE_CHECKING:
+    from raddo.dynamics.base_dynamics import Dynamics
+    from raddo.utils.propagate import TrajectoryGenerator
+
+
+class SpacecraftAcadosModel(eqx.Module):
+  """
+  Constructs the AcadosModel used for defining the system dynamics.
+  """
+  propagator: TrajectoryGenerator
+  dynamics: Dynamics
+  
+  name: str
+  mass: float
+  inertia: jax.Array
+  inertia_inv: jax.Array
+  state_limits: jax.Array
+  control_limits: jax.Array
+
+  dt: float
+  
+  def __init__(
+      self,
+      propagator: TrajectoryGenerator,
+      state_limits: jax.Array,
+      control_limits: jax.Array,
+  ) -> None:
+    """
+    Constructor for the spacecraft model
+    
+    Attributes:
+      name (str): model name.
+      mass (float): robot mass [kg].
+      inertia (np.ndarray): 3x3 inertia tensor [kg*m^2].
+      omega_max (float): maximum (and defines minimum) angular rotation [rad/s].
+      moment_max (float): maximum moment applied along any axis [N*m].
+      dt (float): integration time step.
+    """
+    self.name = 'spacecraft_acados_model'
+    self.propagator = propagator
+    self.dynamics = propagator.dynamics
+    self.dt = propagator.dt
+    self.mass = self.dynamics.dynamics_params['mass']
+    self.inertia = np.array(self.dynamics.dynamics_params['inertia'])
+    self.inertia_inv = np.array(self.dynamics.dynamics_params['inertia_inv'])
+
+    self.state_limits = state_limits
+    self.control_limits = control_limits
+
+  @staticmethod
+  def S(u: cs.MX | cs.SX | cs.DM) -> cs.MX | cs.SX | cs.DM:
+    """
+    Skew symmetric map (lie algebra of SO(3), solves a x b = S(a)b)
+    """
+    return cs.vertcat(
+        cs.horzcat(0, -u[2], u[1]),
+        cs.horzcat(u[2], 0, -u[0]),
+        cs.horzcat(-u[1], u[0], 0))
+
+  def q_left(self, q: cs.MX | cs.SX | cs.DM) -> cs.MX | cs.SX | cs.DM:
+      """
+      Left quaternion product matrix.
+      Assumes q is [qs, qx, qy, qz] (scalar first).
+      """
+      qs = q[0]
+      qv = q[1:]
+      
+      # Top row: [qs, -qv.T]
+      row_top = cs.horzcat(qs, -qv.T)
+      
+      # Bottom block: [qv, qs*I + S(qv)]  (3x3)
+      row_bottom = cs.horzcat(qv, qs * cs.MX.eye(3) + self.S(qv))
+      
+      return cs.vertcat(row_top, row_bottom)
+
+  @staticmethod
+  def get_rotation(q: cs.MX | cs.SX | cs.DM) -> cs.MX | cs.SX | cs.DM:
+      """Quaternion to rotation matrix"""
+      qw, qx, qy, qz = q[0], q[1], q[2], q[3]
+      qw2 = qw * qw
+      qx2 = qx * qx
+      qy2 = qy * qy
+      qz2 = qz * qz
+      return cs.vertcat(
+          cs.horzcat(qw2 + qx2 - qy2 - qz2, 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)),
+          cs.horzcat(2 * (qx * qy + qw * qz), qw2 - qx2 + qy2 - qz2, 2 * (qy * qz - qw * qx)),
+          cs.horzcat(2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), qw2 - qx2 - qy2 + qz2))
+
+  @staticmethod
+  def q_conj(q: cs.MX | cs.SX | cs.DM) -> cs.MX | cs.SX | cs.DM:
+    """Conjugate quaternion"""
+    return cs.vertcat(q[0], -q[1], -q[2], -q[3])
+ 
+  def get_acados_model(self) -> AcadosModel:
+    """
+      Constructs the AcadosModel used for the 3-DoF attitude
+      planning problem for spacecraft. 
+      Constructed w help from https://github.com/acados/acados/blob/main/examples/acados_python/getting_started/pendulum_model.py
+
+      Returns:
+        model (AcadosModel): model for spacecraft slewing problem.
+    """
+    # Set up states
+    q = cs.MX.sym('q',4,1)
+    v = cs.MX.sym('v', 3, 1)
+    x = cs.vertcat(q,v)
+
+    # Set up control input (moment)
+    M = cs.MX.sym('M', 3, 1)
+
+    # Set up params (b-field and target_state
+    b = cs.MX.sym('b',3,1)
+    x_g = cs.MX.sym('xg',7,1)
+    p = cs.vertcat(b,x_g)
+
+    # Magnetorquer torque: rotate b-field from PCI to body frame
+    b_body = self.get_rotation(self.q_conj(q)) @ b
+    tau = cs.cross(M, b_body) * 1e-9  # convert nT to T
+
+    # Set up discrete dynamics
+    qdot = 0.5*self.q_left(q) @ cs.vertcat(0.0, v)
+    J = cs.MX(self.inertia)
+    J_inv = cs.MX(self.inertia_inv)
+
+    # Use torque directly (M) or use dipole (tau)
+    vdot = J_inv @ (tau - cs.cross(v, J @ v)) # use dipole
+    #vdot = J_inv @ (M - cs.cross(v, J @ v)) # use torque directly
+
+    f_disc = cs.vertcat((q + qdot*self.dt)/cs.norm_2(q + qdot*self.dt), v + vdot*self.dt)
+
+    model = AcadosModel()
+
+    model.disc_dyn_expr = f_disc
+    model.x = x
+    model.u = M
+    model.p = p
+    model.name = self.name
+
+    return model
+    
+class AcadosController(eqx.Module):
+  model: SpacecraftAcadosModel
+  horizon: int
+
+  Q: jax.Array
+  Qf: jax.Array
+  R: jax.Array
+
+  max_iter: int
+  ocp_solver: AcadosOcpSolver = eqx.field(static=True)
+  
+  def __init__(
+      self,
+      model: SpacecraftAcadosModel,
+      Q: jax.Array,
+      Qf: jax.Array,
+      R: jax.Array,
+      horizon: int,
+      max_iter: int,
+  ) -> None:
+    """
+    Attributes:
+      model (AcadosModel): acados model.
+      ocp_solver (AcadosOcpSolver): acados solver.
+      N (int): number of time steps in optimal control horizon.
+      x0 (np.ndarray): initial state of system.
+      Q_mat (list[float]): diagonal elements of the Q cost term.
+      R_mat (list[float]): diagonal elements of the R cost term.
+    """
+    self.model = model
+    self.horizon = horizon
+    
+    # set cost
+    self.Q = Q 
+    self.Qf = Qf 
+    self.R = R 
+    
+    self.max_iter = max_iter
+    self.ocp_solver = self.setup(self.horizon)
+
+  def setup(self, N_horizon:int) -> AcadosOcpSolver:
+    """
+    Method to set up the AcadosOcpSolver; this is where the constraints and cost terms
+    are explicitly defined for the system using the AcadosModel.
+
+    Parameters:
+      x0 (np.ndarray): initial state of system. 
+      N_horizon (int): number of time steps in optimal control horizon.
+
+    Returns:
+      ocp_solver (AcadosOcpSolver): acados solver.
+    """
+    # Create ocp object to formulate the OCP
+    ocp = AcadosOcp()
+
+    # set model
+    ocp.model = self.model.get_acados_model()
+
+    # set prediction horizon
+    ocp.solver_options.N_horizon = N_horizon
+    ocp.solver_options.tf = N_horizon*self.model.dt
+
+    # init param
+    ocp.parameter_values =  np.zeros(10) #b-field (3) and target_state (7)
+    target_state = ocp.model.p[3:]
+    q_g = target_state[:4]
+    omega_g = target_state[4:]
+    
+    # path cost
+    ocp.cost.cost_type = 'NONLINEAR_LS'
+
+    q_error = 1 - cs.dot(ocp.model.x[:4], q_g)**2
+    omega_error = omega_g - ocp.model.x[4:7]
+    
+    ocp.model.cost_y_expr = cs.vertcat(q_error, omega_error, ocp.model.u)
+    ocp.cost.yref = np.zeros(7)
+    ocp.cost.W = np.array(jax.scipy.linalg.block_diag(self.Q, self.R))
+
+    # terminal cost
+    ocp.cost.cost_type_e = 'NONLINEAR_LS'
+    ocp.model.cost_y_expr_e = cs.vertcat(q_error, omega_error)
+    ocp.cost.yref_e = np.zeros(4)
+    ocp.cost.W_e = np.array(self.Qf)
+    
+    # set constraints
+    ocp.constraints.x0 = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) # placeholder, gets set in solve
+    ocp.constraints.ubu = np.array(self.model.control_limits[:,1])
+    ocp.constraints.lbu =  np.array(self.model.control_limits[:,0])
+    ocp.constraints.idxbu = np.array([0, 1, 2])
+
+    ocp.constraints.ubx = np.array(self.model.state_limits[:,1]) 
+    ocp.constraints.lbx = np.array(self.model.state_limits[:,0])
+    ocp.constraints.idxbx = np.arange(3)
+    
+    # Set solver details
+    ocp.solver_options.integrator_type = 'DISCRETE'
+    ocp.solver_options.nlp_solver_type = 'SQP'
+    ocp.solver_options.nlp_solver_max_iter = self.max_iter
+
+    #ocp.solver_options.qp_solver = 'PARTIAL_CONDENSING_HPIPM'
+    #ocp.solver_options.hessian_approx = 'GAUSS_NEWTON'
+    #ocp.solver_options.qp_solver_iter_max = 1000
+    #ocp.solver_options.tol = 1e-4
+
+    ocp_solver = AcadosOcpSolver(ocp, verbose=False)
+
+    return ocp_solver
+
+  
+  def solve(
+      self,
+      x0: jax.Array,
+      target_state: jax.Array,
+      key: jax.Array,
+      external_params: jax.Array,
+      nominal_traj: jax.Array,
+      nominal_cntrl: jax.Array,
+  ) -> tuple[jax.Array, jax.Array]:
+    """
+    Solves the optimal control problem starting at an initial state x0.
+
+    Parameters:
+      x0 (np.ndarray: initial state.
+      verbose (bool): verbosity level of solver.
+
+    Returns:
+      x_traj (np.ndarray): state trajectory for system.
+      u_traj (np.ndarray): control trajectory for system.
+    """
+    x_traj = np.zeros((self.horizon+1, len(x0)))
+    u_traj = np.zeros((self.horizon, 3))
+
+    self.ocp_solver.set(0, 'lbx', np.array(x0))
+    self.ocp_solver.set(0, 'ubx', np.array(x0))
+
+    nominal_traj = np.array(nominal_traj)
+    nominal_cntrl = np.array(nominal_cntrl)
+
+    for i in range(self.horizon):
+      # Different parameter at each node
+      p = jnp.concatenate([external_params[i,:], target_state])
+      self.ocp_solver.set(i, 'p', np.array(p))
+
+      # initialize to init traj
+      self.ocp_solver.set(i, 'x', nominal_traj[i, :])
+      self.ocp_solver.set(i, 'u', nominal_cntrl[i, :])
+
+    # terminal condition
+    p = jnp.concatenate([external_params[-1,:], target_state])
+    self.ocp_solver.set(self.horizon, 'p', np.array(p))
+    self.ocp_solver.set(self.horizon, 'x', nominal_traj[-1, :])
+    
+    status = self.ocp_solver.solve()
+
+    verbose=False
+    
+    if verbose:
+      if status == 0:
+        print("Success")
+      elif status == 1:
+        print("Acados failure or nan")
+      elif status == 2:
+        print("Max iterations")
+      elif status == 3:
+        print("Acados minstep")
+      elif status == 4:
+        print("QP Failure")
+    
+    for i in range(self.horizon):
+      x_traj[i,:] = self.ocp_solver.get(i, 'x')
+      u_traj[i,:] = self.ocp_solver.get(i, 'u')
+    x_traj[self.horizon,:] = self.ocp_solver.get(self.horizon, 'x')
+
+    return jnp.array(x_traj), jnp.array(u_traj)
+    
+  def __call__(
+      self,
+      x0: jax.Array,
+      target_state: jax.Array,
+      key: jax.Array,
+      external_params: jax.Array,
+      nominal_traj: jax.Array,
+      nominal_cntrl: jax.Array,
+  ) -> tuple[jax.Array, jax.Array]:
+    # Ensure shapes use concrete integers (self.horizon)
+    shape1 = jax.ShapeDtypeStruct((self.horizon + 1, x0.shape[-1]), jnp.float32)
+    shape2 = jax.ShapeDtypeStruct((self.horizon, 3), jnp.float32)
+    
+    # Wrap the call to avoid hashing 'self'
+    def callback_wrapper(
+        x0: np.ndarray,
+        target: np.ndarray,
+        k: np.ndarray,
+        ext: np.ndarray,
+        nom_t: np.ndarray,
+        nom_c: np.ndarray,
+    ) -> tuple[jax.Array, jax.Array]:
+        return self.solve(x0, target, k, ext, nom_t, nom_c)
+
+    return jax.pure_callback(
+        callback_wrapper,
+        (shape1, shape2),
+        x0, target_state, key, external_params, nominal_traj, nominal_cntrl,
+        vmap_method='sequential')

raddo/controllers/base_controller.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import equinox as eqx
+
+if TYPE_CHECKING:
+    import jax
+
+
+class Controller(eqx.Module):
+  def __call__(
+      self,
+      state: jax.Array,
+      target: jax.Array,
+      key: jax.Array,
+  ) -> jax.Array | tuple[jax.Array, jax.Array]:
+    raise NotImplementedError("Must define a controller if no control sequence is given")

raddo/controllers/cem_controller.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+
+from .base_controller import Controller
+
+if TYPE_CHECKING:
+    from raddo.dynamics.base_dynamics import Dynamics
+    from raddo.utils.propagate import TrajectoryGenerator
+
+
+class CEMController(Controller):
+  propagator: TrajectoryGenerator
+  dynamics: Dynamics
+  
+  # Parameters
+  Q: jax.Array
+  Qf: jax.Array
+  R: jax.Array
+  control_limits: jax.Array
+  
+  # Hyperparameters (Static) (ints and floats default to static)
+  num_samples: int
+  num_iter: int
+  horizon: int
+  elite_percent: float
+  init_std: float
+
+  # Useful for working with dynamics
+  quat_start: int
+  nx: int
+  nu: int
+
+  def __init__(
+      self,
+      propagator: TrajectoryGenerator,
+      Q: jax.Array,
+      Qf: jax.Array,
+      R: jax.Array,
+      control_limits: jax.Array,
+      num_samples: int,
+      num_iter: int,
+      horizon: int,
+      elite_percent: float=0.1,
+      init_std: float=0.1,
+  ) -> None:
+    object.__setattr__(self, "propagator", propagator)
+    object.__setattr__(self, "dynamics", self.propagator.dynamics)
+    object.__setattr__(self, "Q", (Q + Q.T)/2) # ensure symmetric
+    object.__setattr__(self, "Qf", (Qf + Qf.T)/2)
+    object.__setattr__(self, "R", (R+R.T)/2)
+    object.__setattr__(self, "num_samples", num_samples)
+    object.__setattr__(self, "num_iter", num_iter)
+    object.__setattr__(self, "horizon", horizon)
+    object.__setattr__(self, "elite_percent", elite_percent)
+    object.__setattr__(self, "init_std", init_std)
+    object.__setattr__(self, "nu", self.dynamics.num_controls)
+    object.__setattr__(self, "control_limits", control_limits)
+
+    self.nx = self.dynamics.num_states
+    nx = self.dynamics.num_states
+    quat_start = self.dynamics.params.get("quat_start")
+
+    # Check if attitude Jacobian is square (whether using mrp's or not)
+    if not self.dynamics.E(jnp.ones(nx)).shape[0] == self.dynamics.E(jnp.ones(nx)).shape[1]:
+      nx = self.dynamics.num_states - 1
+
+    object.__setattr__(self, "nx", nx)
+    object.__setattr__(self, "quat_start", quat_start)
+
+  def compute_cost(
+      self,
+      states: jax.Array,
+      controls: jax.Array,
+      target_state: jax.Array,
+  ) -> jax.Array:
+    """Computes the total quadratic cost for a trajectory and control sequence, including terminal cost.
+        Inputs:
+        states: Sequence of states, shape (T+1, nx)
+        controls: Sequence of controls, shape (T, nu)
+        Outputs:
+        cost: Scalar total cost (float)
+    """
+    # Terminal cost
+    final_error = self.dynamics.get_error_coordinates(states[-1,:], target_state)
+    J_terminal = 0.5 * final_error.T @ self.Qf @ final_error
+
+    # Stage cost (vectorized over time) 
+    diff_stage = states[:-1, :] - target_state
+    
+    def stage_step_cost(x_k: jax.Array, u_k: jax.Array) -> jax.Array:
+      error = self.dynamics.get_error_coordinates(x_k, target_state) 
+      return (0.5 * error.T @ self.Q @ error + 0.5 * u_k.T @ self.R @ u_k)*self.propagator.dt
+        
+    J_stage = jax.vmap(stage_step_cost)(states[:-1], controls)
+    
+    return J_terminal + jnp.sum(J_stage)
+
+  @eqx.filter_jit
+  def __call__(
+      self,
+      x0: jax.Array,
+      x_goal: jax.Array,
+      key: jax.Array,
+      external_dynamics_params: jax.Array,
+      nominal_traj: jax.Array,
+      nominal_cntrl: jax.Array,
+  ) -> tuple[jax.Array, jax.Array]:
+    """
+    Cross-Entropy (CEM) algorithm.
+    Returns:
+        nominal state trajectory: (T+1, nx)
+        nominal control trajectory: (T, nu)
+    """
+    n_elites = int(self.num_samples * self.elite_percent)
+    nu = self.nu
+
+    # Initialize mean and covariance (std dev)
+    mu = jnp.zeros((self.horizon, nu)) # init to 0
+    # mu = nominal_cntrl # warm-start
+    sigma = self.init_std * jnp.ones((self.horizon, nu)) # Initial wide exploration
+    
+    # Optimization Loop
+    def scan_step(
+        carry: tuple[jax.Array, jax.Array, jax.Array],
+        _: None,
+    ) -> tuple[tuple[jax.Array, jax.Array, jax.Array], jax.Array]:
+      curr_mu, curr_sigma, key = carry
+      
+      key, sample_key, traj_key = jrandom.split(key, 3)
+      
+      # Sample controls: (K, T, nu)
+      # z ~ N(0, 1) -> u = mu + sigma * z
+      z = jrandom.normal(sample_key, shape=(self.num_samples, self.horizon, nu))
+      u_candidates = curr_mu + curr_sigma * z
+      
+      # Clip
+      u_candidates = jnp.clip(
+          u_candidates, 
+          self.control_limits[:, 0], 
+          self.control_limits[:, 1]
+      )
+
+      # Rollout batch of trajectories from different u_candidates
+      states_batch, _ = self.propagator._generate_trajectory_batch_seq(
+                              initial_states = x0,
+                              target_states = x_goal,
+                              key = traj_key,
+                              batch_size = self.num_samples,
+                              noise_std = .001, #jnp.array([0.001]*4 + [0.005]*3)
+        
+                              num_steps = self.horizon,
+                              external_dynamics_params = external_dynamics_params,
+                              control_sequence = u_candidates)
+      
+      # Compute costs (vmap over the batch dimension)
+      costs = jax.vmap(self.compute_cost, in_axes=(0, 0, None))(states_batch, u_candidates, x_goal)
+      
+      # Select Elites
+      _, idx_elites = jax.lax.top_k(-costs, n_elites)
+      
+      # Gather elite samples: (n_elites, T, nu)
+      elites = u_candidates[idx_elites]
+      elite_traj = states_batch[idx_elites]
+      
+      # Update Distribution
+      new_mu = jnp.mean(elites, axis=0)
+      new_sigma = jnp.std(elites, axis=0)
+
+      # Get mean cost of elites
+      cost = jnp.mean(costs[idx_elites])
+  
+      #return (new_mu, new_sigma,  key), carry #  [[ 1 ]] For running in closed loop
+      return (new_mu, new_sigma,  key), cost#, states_batch, elite_traj) # [[ 2 ]] Store all samples for debugging
+
+
+    # Run optimization
+    init_carry = (mu, sigma, key)
+    (final_mu, final_sigma, final_key), _ = jax.lax.scan(scan_step, init_carry, length=self.num_iter)
+
+
+    final_u_nominal = jnp.clip(final_mu, self.control_limits[:, 0], self.control_limits[:, 1])
+
+    best_states, best_controls = self.propagator._generate_trajectory_seq(
+                                initial_state = x0,
+                                target_state = x_goal,
+                                key = final_key,
+                                noise_std = 0.0,
+                                num_steps = self.horizon,
+                                external_dynamics_params = external_dynamics_params,
+                                control_sequence = final_u_nominal)
+    
+    return best_states, best_controls

raddo/controllers/feedback_controller.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import equinox as eqx
+import jax.numpy as jnp
+
+from .base_controller import Controller
+
+if TYPE_CHECKING:
+    import jax
+
+    from raddo.dynamics.base_dynamics import Dynamics
+    from raddo.utils.propagate import TrajectoryGenerator
+
+
+class FeedbackController(Controller):
+  propagator: TrajectoryGenerator
+  dynamics: Dynamics
+  cntrl_param_init: jax.Array
+  
+  # Parameters
+  Q: jax.Array
+  R: jax.Array
+  control_limits: jax.Array
+
+  nx: int = eqx.field(static=True)
+
+  def __init__(
+      self,
+      propagator: TrajectoryGenerator,
+      Q: jax.Array,
+      Qf: jax.Array,
+      R: jax.Array,
+      control_limits: jax.Array,
+  ) -> None:
+    self.propagator = propagator
+    self.dynamics = self.propagator.dynamics
+    self.Q = Q
+    self.cntrl_param_init = Qf
+    self.R = R
+    self.control_limits = control_limits
+
+    nx = self.dynamics.num_states
+
+    # Check if attitude Jacobian is square (whether using mrp's or not)
+    if not self.dynamics.E(jnp.ones(nx)).shape[0] == self.dynamics.E(jnp.ones(nx)).shape[1]:
+      nx = self.dynamics.num_states - 1
+
+    object.__setattr__(self, "nx", nx)
+
+
+  @eqx.filter_jit
+  def __call__(
+      self,
+      x0: jax.Array,
+      x_goal: jax.Array,
+      key: jax.Array,
+      external_dynamics_params: jax.Array,
+      nominal_traj: jax.Array,
+      nominal_cntrl: jax.Array,
+      cntrl_param: jax.Array,
+  ) -> tuple[jax.Array, jax.Array]:
+    """
+    Feedback controller for tracking a nominal trajectory
+
+    Returns:
+        nominal state trajectory: (T+1, nx)
+        nominal control trajectory: (T, nu)
+    """
+    dt = self.propagator.dt
+    A, B = self.dynamics.linearize_and_discretize(jnp.array([x0, x0]), nominal_cntrl[:1], external_dynamics_params[:1], dt)
+    #A, B = self.dynamics.linearize_and_discretize(nominal_traj, nominal_cntrl, external_dynamics_params, dt)
+    
+    A = A[0]
+    B = B[0]
+    Q = self.Q
+    R = self.R 
+    S = cntrl_param
+
+    dx = self.dynamics.get_error_coords(x0, nominal_traj[0])
+
+    # Ricatti equation
+    eps = 1e-8*jnp.ones(R.shape) # to avoid singular matrix
+    K = jnp.linalg.solve(R + B.T @ S @ B + eps, B.T @ S @ A)
+    # S = Q + K.T @ R @ K + (A - B@K).T @ S @ (A - B@K)
+    S = Q + A.T @ S @ A - A.T @ S @ B @ K
+
+    # Control law
+    u = nominal_cntrl[0] - K @ dx
+    
+    # Clip controls
+    u = jnp.clip(u, self.control_limits[:, 0], self.control_limits[:, 1])
+
+    return u, S