Control Utilities¶

This page documents utility classes and components used by controllers in SoRoMoX.

PID Control¶

The PIDControl class provides a flexible PID control implementation that can be used as the error-based feedback term in model-based controllers.

Features¶

Flexible gain specification: Scalar, diagonal vector, or full matrix gains
Anti-windup: Optional saturation functions for integral error
JAX-compatible: Fully jittable for high-performance execution

Control Law¶

The PID control law is:

\[ u = K_p e + K_i \int_0^t e(\tau) d\tau + K_d \dot{e} \]

The integral error dynamics include optional saturation for anti-windup:

\[ \frac{d}{dt}\left(\int e \, dt\right) = \mathrm{sat}(e) \]

For the built-in hyperbolic-tangent saturation, gamma is an inverse error scale. The componentwise scalar/vector form is

\[ \mathrm{sat}_{\gamma}(e) = \frac{\tanh(\gamma e)}{\gamma} = e_{\mathrm{sat}}\tanh\!\left(\frac{e}{e_{\mathrm{sat}}}\right), \qquad e_{\mathrm{sat}}=\frac{1}{\gamma}. \]

This keeps the saturation argument dimensionless, preserves the units of the integrated error, and has unit slope at the origin. A vector gamma can assign different physical error scales to heterogeneous coordinates. Scalar and vector values of gamma must be strictly positive; a matrix gamma must be symmetric positive definite.

PIDControl¶

soromox.control.PIDControl ¶

PIDControl(Kp: float | Array, Ki: float | Array, Kd: float | Array, saturation_fn: SaturationFnName | Callable[[Array], Array] | None = None, gamma: float | Array = 1.0)

Bases: Module


              flowchart TD
              soromox.control.PIDControl[PIDControl]

              

              click soromox.control.PIDControl href "" "soromox.control.PIDControl"

PID controller with configurable gains and optional integral error saturation.

This class implements a standard PID control law

u = Kp @ e + Ki @ integral_error + Kd @ ed

where

e is the tracking error
ed is the error derivative
integral_error is the accumulated integral of the (possibly saturated) error

The integral error dynamics are

integral_error_dot = sat(e)

where sat() is a saturation function that can be used for anti-windup.

Attributes:

Name	Type	Description
`Kp`	`Array`	Proportional gain (scalar, diagonal vector, or matrix).
`Ki`	`Array`	Integral gain (scalar, diagonal vector, or matrix).
`Kd`	`Array`	Derivative gain (scalar, diagonal vector, or matrix).
`gamma`	`Array`	Inverse error scale for built-in saturation functions (e.g., `"tanh"`). Its reciprocal has the same units as the corresponding error and sets the asymptotic magnitude of the saturated error.

References

Pustina, P., Borja, P., Della Santina, C., & De Luca, A. (2022). P-satI-D shape regulation of soft robots. IEEE Robotics and Automation Letters, 8(1), 1-8.

Initialize the PID controller.

Parameters:

Name	Type	Description	Default
`Kp`	`float \| Array`	Proportional gain. Can be: - A scalar (float or 0-d array): applied uniformly to all error components. - A 1-d array (diagonal): element-wise multiplication with error. - A 2-d array (matrix): full matrix multiplication with error.	required
`Ki`	`float \| Array`	Integral gain. Same format options as Kp.	required
`Kd`	`float \| Array`	Derivative gain. Same format options as Kp.	required
`saturation_fn`	`SaturationFnName \| Callable[[Array], Array] \| None`	Optional saturation function for anti-windup. Can be: - None or "identity": No saturation (default). - "tanh": Uses `solve(gamma, tanh(gamma @ e))` for a matrix `gamma` and `tanh(gamma * e) / gamma` otherwise. This preserves the units and small-error slope of `e`. - A callable `f(e) -> saturated_e`: Custom saturation function.	`None`
`gamma`	`float \| Array`	Inverse error scale for built-in saturation functions. Can be a positive scalar, positive diagonal vector, or symmetric positive-definite matrix. For a scalar or vector, `1 / gamma` sets the componentwise saturation magnitude. Only used when `saturation_fn="tanh"`. Defaults to 1.0.	`1.0`

call ¶

__call__(e: Array, ed: Array, integral_error: Array) -> tuple[Array, Array]

Compute the PID control output and integral error derivative.

Parameters:

Name	Type	Description	Default
`e`	`Array`	Tracking error vector, shape (n,).	required
`ed`	`Array`	Error derivative vector, shape (n,).	required
`integral_error`	`Array`	Current integral of error, shape (n,).	required

Returns:

Name	Type	Description
`u`	`Array`	Control output, shape (n,) or (m,) depending on gain shapes.
`integral_error_dot`	`Array`	Time derivative of integral error, shape (n,). This is the (possibly saturated) error that should be integrated.

update_gains ¶

update_gains(gains: dict[str, Array]) -> PIDControl

This function updates the gains of the PIDControl object.

Parameters:

Name	Type	Description	Default
`gains`	`dict[str, Array]`	proportional, integral, and derivative gains	required

Returns:

Name	Type	Description
`updated_self`	`PIDControl`	self object with updated parameters

PIDControllerState¶

Container for PID controller state (integral error).

soromox.control.PIDControllerState `dataclass` ¶

PIDControllerState(integral_error: Array)

State container for the PID controller.

This dataclass holds all internal state variables that evolve over time during closed-loop control. It is registered as a JAX pytree to enable use with JAX transformations and ODE solvers.

Attributes:

Name	Type	Description
`integral_error`	`Array`	Accumulated integral of the (possibly saturated) error. Shape: (num_dofs,)

tree_flatten ¶

tree_flatten()

Flatten the state for JAX pytree operations.

tree_unflatten `classmethod` ¶

tree_unflatten(aux_data, children)

Unflatten the state from JAX pytree operations.

zero `classmethod` ¶

zero(num_dofs: int) -> PIDControllerState

Create a zero-initialized controller state.

Parameters:

Name	Type	Description	Default
`num_dofs`	`int`	Number of degrees of freedom.	required

Returns:

Type	Description
`PIDControllerState`	PIDControllerState with zero integral error.

Usage Example¶

import jax.numpy as jnp
from soromox.control import PIDControl, PIDControllerState

# Create PID controller with diagonal gains
pid = PIDControl(
    Kp=jnp.array([100.0, 100.0]),  # Proportional gains
    Ki=jnp.array([10.0, 10.0]),    # Integral gains
    Kd=jnp.array([20.0, 20.0]),    # Derivative gains
    saturation_fn="tanh",          # Anti-windup saturation
    gamma=1.0,                     # Inverse saturation-error scale
)

# Initialize controller state
control_state = PIDControllerState.zero(num_dofs=2)

# Compute control output
e = q_des - q          # Position error
ed = qd_des - qd       # Velocity error
u, integral_error_dot = pid(e, ed, control_state.integral_error)

Reference Trajectory¶

The ReferenceTrajectory class provides a flexible way to specify desired trajectories for controllers. It supports both discrete time-series data and continuous functions, automatically deriving velocities and accelerations when not provided.

Features¶

Dual representation: Discrete time series and continuous functions
Auto-derivation: Automatically computes velocity and acceleration using JAX autodiff
Linear interpolation: Creates smooth continuous trajectories from discrete data
JAX PyTree: Compatible with JAX transformations and ODE solvers

ReferenceTrajectory¶

soromox.control.ReferenceTrajectory `dataclass` ¶

ReferenceTrajectory(ts: Array, x_des_ts: Array | None = None, x_des_fn: Callable[[Array], Array] | None = None, xd_des_ts: Array | None = None, xd_des_fn: Callable[[Array], Array] | None = None, xdd_des_ts: Array | None = None, xdd_des_fn: Callable[[Array], Array] | None = None, rotation_representation: RotationRepresentation | None = None, n_points: int = 1, is_planar: bool = False)

Container for reference trajectory with both discrete and continuous representations.

This class provides flexible specification of reference trajectories for controllers. Users can provide either discrete time-series data or continuous functions, and the class will automatically generate the other representation.

When working with poses that include orientation (for 3D robots), you can specify a rotation_representation to enable proper velocity derivation. This ensures that angular velocities are computed correctly from orientation changes (not just naive differentiation).

Parameters¶

ts : Array Time steps array, shape (T,). x_des_ts : Optional[Array] Desired poses at discrete time steps, shape (T, dim). If not provided, will be computed from x_des_fn. x_des_fn : Optional[Callable[[Array], Array]] Function that returns desired pose at any time t. If not provided, will be created via linear interpolation of x_des_ts. xd_des_ts : Optional[Array] Desired velocities (twists) at discrete time steps, shape (T, dim). If not provided, will be computed from xd_des_fn or derived from x_des_fn. xd_des_fn : Optional[Callable[[Array], Array]] Function that returns desired velocity at any time t. If not provided, will be created via interpolation or derived from x_des_fn. xdd_des_ts : Optional[Array] Desired accelerations at discrete time steps, shape (T, dim). If not provided, will be computed from xdd_des_fn or derived from xd_des_fn. xdd_des_fn : Optional[Callable[[Array], Array]] Function that returns desired acceleration at any time t. If not provided, will be created via interpolation or derived from xd_des_fn. rotation_representation : Optional[RotationRepresentation] The rotation representation used in the pose. If provided, velocity derivation will properly convert orientation derivatives to angular velocities. - ROTATION_VECTOR: 3D rotation vector (axis-angle) - QUATERNION: 4D unit quaternion - ROTATION_MATRIX_6D: 6D continuous representation If None, simple differentiation is used (appropriate for position-only trajectories or planar robots). n_points : int Number of points in the pose (default 1). Used with rotation_representation for proper velocity derivation. is_planar : bool Whether the trajectory is for a planar robot. If True, velocity derivation uses simple differentiation even if rotation_representation is set.

Examples¶

From discrete time series (position-only):

ts = jnp.linspace(0, 1, 100)
x_des_ts = jnp.sin(ts)[:, None]  # shape (100, 1)
reference = ReferenceTrajectory(ts=ts, x_des_ts=x_des_ts)
reference.x_des_fn(0.5)

From a continuous function with proper twist derivation:

def x_fn(t):
    return jnp.array([0.0, 0.0, t * 0.1, 0.0, 0.0, 0.2 + 0.01 * t])

reference = ReferenceTrajectory(
    ts=ts,
    x_des_fn=x_fn,
    rotation_representation=RotationRepresentation.ROTATION_VECTOR,
)
reference.xd_des_fn(0.5)

tree_flatten ¶

tree_flatten()

Flatten the ReferenceTrajectory for JAX PyTree operations.

tree_unflatten `classmethod` ¶

tree_unflatten(aux, children)

Unflatten a ReferenceTrajectory from JAX PyTree operations.

Usage Examples¶

From Discrete Time Series¶

import jax.numpy as jnp
from soromox.control import ReferenceTrajectory

# Time array
ts = jnp.linspace(0, 10, 1000)

# Desired positions at each time step
x_des_ts = jnp.column_stack([
    jnp.sin(ts),           # First DOF
    jnp.cos(ts),           # Second DOF
])

# Create trajectory - velocities and accelerations auto-derived
ref_traj = ReferenceTrajectory(ts=ts, x_des_ts=x_des_ts)

# Access continuous function (interpolated)
x_at_5 = ref_traj.x_des_fn(5.0)
xd_at_5 = ref_traj.xd_des_fn(5.0)
xdd_at_5 = ref_traj.xdd_des_fn(5.0)

From Continuous Function¶

import jax.numpy as jnp
from soromox.control import ReferenceTrajectory

# Define trajectory as a function
def x_des_fn(t):
    return jnp.array([jnp.sin(t), jnp.cos(t)])

# Time array for discrete evaluation
ts = jnp.linspace(0, 10, 1000)

# Create trajectory - discrete arrays and derivatives auto-computed
ref_traj = ReferenceTrajectory(ts=ts, x_des_fn=x_des_fn)

# Access discrete arrays
x_des_ts = ref_traj.x_des_ts    # Shape: (1000, 2)
xd_des_ts = ref_traj.xd_des_ts  # Auto-derived via JAX autodiff

Constant Setpoint (Regulation)¶

import jax.numpy as jnp
from soromox.control import ReferenceTrajectory

# For regulation, create a constant trajectory
q_setpoint = jnp.array([0.5, 0.3])

ts = jnp.linspace(0, 10, 100)
x_des_ts = jnp.tile(q_setpoint, (len(ts), 1))

ref_traj = ReferenceTrajectory(ts=ts, x_des_ts=x_des_ts)
# xd_des and xdd_des will be (approximately) zero

Control Utilities¶

PID Control¶

Features¶

Control Law¶

PIDControl¶

soromox.control.PIDControl ¶

__call__ ¶

update_gains ¶

PIDControllerState¶

soromox.control.PIDControllerState dataclass ¶

tree_flatten ¶

tree_unflatten classmethod ¶

zero classmethod ¶

Usage Example¶

Reference Trajectory¶

Features¶

ReferenceTrajectory¶

soromox.control.ReferenceTrajectory dataclass ¶

Parameters¶

Examples¶

tree_flatten ¶

tree_unflatten classmethod ¶

Usage Examples¶

From Discrete Time Series¶

From Continuous Function¶

Constant Setpoint (Regulation)¶

call ¶

soromox.control.PIDControllerState `dataclass` ¶

tree_unflatten `classmethod` ¶

zero `classmethod` ¶

soromox.control.ReferenceTrajectory `dataclass` ¶

tree_unflatten `classmethod` ¶