Simulation Integrators¶

Soromox provides Diffrax-compatible solvers for state layouts and dynamics that need integration behavior beyond Diffrax's generic solver collection.

Diffrax solvers¶

Soromox rollouts accept any compatible diffrax.AbstractSolver through the solver argument. This includes standard Diffrax methods such as Tsit5, Dopri5, Bosh3, Heun, and Euler; Tsit5 remains the default when no solver is supplied. See Diffrax's ODE solver catalog for the complete collection. Use a compatible step-size controller when adaptive stepping is desired.

Standard Diffrax solvers operate on the complete first-order Soromox state, so they support [q, qd, auxiliary] layouts, including unequal spatial floating-base configuration and velocity dimensions. Before every dynamics evaluation and on saved output, Soromox projects the state through the system's project_state hook. Spatial floating-base robots therefore present a normalized quaternion to the dynamics and return normalized saved states.

Floating-base quaternion integration

Standard Diffrax solvers advance the stored quaternion in its ambient coordinates. Soromox projects the quaternion before dynamics evaluations and on saved output, but these solvers do not perform a Lie-group retraction at each Runge--Kutta stage. Validate the timestep or adaptive tolerances for the floating-base model and operating range.

from diffrax import PIDController, Tsit5

trajectory = robot.rollout_to(
    initial_state=initial_state,
    t1=1.0,
    solver_dt=1e-4,
    save_dt=1e-2,
    solver=Tsit5(),
    stepsize_controller=PIDController(rtol=1e-5, atol=1e-7),
)

Semi-implicit Euler¶

SemiImplicitEuler is a deterministic, first-order kick–drift integrator for Soromox systems whose state is ordered as

\[ y = \begin{bmatrix}q & qd & auxiliary\end{bmatrix}^{\mathsf T}. \]

For a fixed step \(h\), it evaluates the complete forward dynamics at the old state and updates velocity before configuration:

\[ qd_{n+1} = qd_n + h\,qdd(q_n, qd_n), \qquad q_{n+1} = \operatorname{retract}(q_n, h\,qd_{n+1}). \]

This kick–drift ordering is generally more stable for oscillatory mechanical systems than forward Euler. It remains a conditionally stable, first-order method, so the timestep must still be validated for the model family and operating range.

from soromox.simulation import SemiImplicitEuler
from soromox.systems import SystemState

trajectory = robot.rollout_to(
    initial_state=SystemState(t=0.0, y=y0),
    t1=1.0,
    solver_dt=1e-5,
    save_dt=1e-2,
    solver=SemiImplicitEuler(robot),
)

Soromox-specific state and dynamics contract

The primary integrated state must follow the system's [q, qd, auxiliary] contract. The solver calls the system's state and retraction helpers, which support unequal floating-base configuration and velocity dimensions and preserve spatial quaternion geometry. It is not a general-purpose semi-implicit method for arbitrary Diffrax PyTrees. Unlike diffrax.SemiImplicitEuler, which requires a separable two-term vector field, the Soromox implementation evaluates the full derivative because its acceleration can depend on both configuration and velocity. It then applies the mechanical kick--drift update to the [q, qd, auxiliary] leaf. Optional Diffrax-level control and environment states are advanced with forward Euler.

soromox.simulation.SemiImplicitEuler ¶

Bases: AbstractSolver


              flowchart TD
              soromox.simulation.SemiImplicitEuler[SemiImplicitEuler]

              

              click soromox.simulation.SemiImplicitEuler href "" "soromox.simulation.SemiImplicitEuler"

Kick-drift Euler for a Soromox [q, qd, auxiliary] state.

Diffrax's generic :class:diffrax.SemiImplicitEuler requires a separable two-term vector field. Soromox dynamics generally include velocity-dependent forces, so this solver evaluates the full state derivative at (q_n, qd_n) and applies

qd_{n+1} = qd_n + dt * qdd_n and q_{n+1} = retract(q_n, dt * qd_{n+1}).

The solver is intended for :class:soromox.systems.SoftRobot rollouts. It uses the supplied system's state helpers rather than splitting the state in half, so unequal floating-base configuration and velocity dimensions, manifold retraction, and trailing system auxiliary states are supported. Diffrax-level control and environment state leaves are advanced with explicit Euler.

Attributes:

Name	Type	Description
`system`	`SoftRobot`	Soft robot system that supplies `split_state`, `pack_state`, and `retract_configuration`.

order ¶

order(terms: AbstractTerm) -> int

Return the deterministic convergence order.

Parameters:

Name	Type	Description	Default
`terms`	`AbstractTerm`	Diffrax term describing the vector field. The order is one for every compatible term.	required

Returns:

Type	Description
`int`	The deterministic convergence order, which is one.

init ¶

init(terms: AbstractTerm, t0: Any, t1: Any, y0: Any, args: Any) -> None

Initialize the solver state.

Parameters:

Name	Type	Description	Default
`terms`	`AbstractTerm`	Diffrax term describing the vector field.	required
`t0`	`Any`	Initial time of the integration interval.	required
`t1`	`Any`	Final time of the integration interval.	required
`y0`	`Any`	Initial Diffrax state.	required
`args`	`Any`	Additional arguments supplied to the vector field.	required

Returns:

Type	Description
`None`	`None` because the method has no persistent solver state.

step ¶

step(terms: AbstractTerm, t0: Any, t1: Any, y0: Any, args: Any, solver_state: Any, made_jump: Any) -> tuple[Any, None, dict[str, Any], None, RESULTS]

Advance the state by one kick--drift step.

Parameters:

Name	Type	Description	Default
`terms`	`AbstractTerm`	Diffrax term describing the vector field.	required
`t0`	`Any`	Start time of the step.	required
`t1`	`Any`	End time of the step.	required
`y0`	`Any`	Diffrax state at `t0`.	required
`args`	`Any`	Additional arguments supplied to the vector field.	required
`solver_state`	`Any`	Persistent solver state. This method does not use it.	required
`made_jump`	`Any`	Whether the solution jumped at `t0`. This method does not require special jump handling.	required

Returns:

Type	Description
`Any`	A Diffrax step tuple containing the state at `t1`, no local error
`None`	estimate, linear-interpolation data, no persistent solver state,
`dict[str, Any]`	and a successful result code.

func ¶

func(terms: AbstractTerm, t0: Any, y0: Any, args: Any) -> Any

Evaluate the underlying vector field.

Parameters:

Name	Type	Description	Default
`terms`	`AbstractTerm`	Diffrax term describing the vector field.	required
`t0`	`Any`	Time at which to evaluate the vector field.	required
`y0`	`Any`	Diffrax state at `t0`.	required
`args`	`Any`	Additional arguments supplied to the vector field.	required

Returns:

Type	Description
`Any`	Vector-field evaluation with the same PyTree structure as `y0`.