🚀 Quick Start¶

For construction, immutable parameter replacement, and differentiable identification workflows, see Parameters, Updates, and Optimization.

Get up and running with SoRoMoX in minutes! This hands-on guide walks you through your first soft robot simulation with step-by-step examples.

🎯 Your First Simulation¶

Let's dive right in with a classic example - simulating a double pendulum to understand the basics using the Pendulum class API:

Tabbed Interface

This example is broken down into four tabs. Click through each tab to see the complete workflow: Setup → Initialize → Simulate → Analyze. All code blocks build on previous tabs.

🔧 Setup🏗️ Initialize⚡ Simulate📊 Analyze

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from soromox.systems import Pendulum, PendulumParams, SystemState
from soromox.rendering import RendererConfig, GeometryConfig, ViserRenderer

# Define parameters for a double pendulum
num_links = 2
params = PendulumParams(
    length=jnp.array([0.5, 0.3]),
    center_of_mass_length=jnp.array([0.25, 0.15]),
    mass=jnp.array([1.0, 0.5]),
    moment_inertia=jnp.array([0.1, 0.05]),
    gravity=jnp.array([0.0, -9.81]),
    joint_stiffness=jnp.zeros((num_links, num_links)),
    joint_damping=jnp.zeros((num_links, num_links)),
    joint_rest_configuration=jnp.zeros((num_links,)),
    radius=jnp.array([0.025, 0.015]),
)

# Create the system using the class API
robot = Pendulum(params=params)

# Set initial conditions
q0 = jnp.array([jnp.pi/4, jnp.pi/6])   # Initial angles [rad]
qd0 = jnp.array([0.0, 0.0])            # Initial velocities [rad/s]
initial_state = SystemState(t=0.0, y=jnp.concatenate([q0, qd0]))

# Configure simulation
t_span = (0.0, 5.0)  # 5 seconds
solver_dt = 0.01    # 10ms timestep

# Zero control torques
u = jnp.zeros((num_links,))

# Integrate using the class helper (Diffrax under the hood)
trajectory = robot.rollout_to(
    initial_state=initial_state,
    u=u,
    t1=t_span[1],
    solver_dt=solver_dt,
    save_dt=solver_dt,
)
ts = trajectory.t
q_ts, qd_ts = jnp.split(trajectory.y, 2, axis=1)
u_ts = trajectory.u

# Extract results from trajectory

# Compute end-effector (last tip) trajectory over time (using vmap for batched time sample processing)
# forward_kinematics_tips returns [theta, px, py] for each link tip
xee_ts = jax.vmap(robot.forward_kinematics_tips)(q_ts)[:,-1,1:]

# Create visualization with three subplots
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 10))

# Plot 1: Joint angles over time
ax1.plot(t, q_ts[:, 0], label='Joint 1', linewidth=2, color='#1f77b4')
ax1.plot(t, q_ts[:, 1], label='Joint 2', linewidth=2, color='#ff7f0e')
ax1.set_ylabel('Joint Angles [rad]')
ax1.set_title('🔄 Joint Angles vs Time')
ax1.legend()
ax1.grid(True, alpha=0.3)

# Plot 2: Joint velocities over time
ax2.plot(t, qd_ts[:, 0], label='Joint 1', linewidth=2, color='#1f77b4')
ax2.plot(t, qd_ts[:, 1], label='Joint 2', linewidth=2, color='#ff7f0e')
ax2.set_ylabel('Joint Velocities [rad/s]')
ax2.set_title('💨 Joint Velocities vs Time')
ax2.legend()
ax2.grid(True, alpha=0.3)

# Plot 3: End-effector trajectory in workspace
ax3.plot(xee_ts[:, 0], xee_ts[:, 1],  # x and y coordinates of the end-effector
         linewidth=3, color='#2ca02c', alpha=0.8)
ax3.scatter(xee_ts[0, 0], xee_ts[0, 1], 
            s=100, color='green', marker='o', label='Start', zorder=5)
ax3.scatter(xee_ts[-1, 0], xee_ts[-1, 1], 
            s=100, color='red', marker='s', label='End', zorder=5)
ax3.set_xlabel('X Position [m]')
ax3.set_ylabel('Y Position [m]')
ax3.set_title('🎯 End-Effector Trajectory')
ax3.legend()
ax3.grid(True, alpha=0.3)
ax3.axis('equal')

plt.tight_layout()
plt.show()

# Interactive 3D visualization with ViserRenderer
# Opens a web-based interface in your browser for interactive exploration
viser_renderer = ViserRenderer(robot, config=RendererConfig(geometry=GeometryConfig(num_points=50, backbone_style="discrete")))
viser_renderer.render_sequence(
    ts,
    q_ts,
    playback_speed=1.0,
    loop=True,
    autoplay=True,
    plot_configurations=True,
    robot_name="Pendulum",
)
# The visualization will open in your browser automatically
# You can interact with the 3D scene, play/pause animation, and view plots

🎉 Congratulations!

You've just simulated your first robot with SoRoMoX! The pendulum exhibits chaotic motion due to its nonlinear dynamics.

What you learned: - How to define system parameters - How to initialize a robot system - How to run a simulation using rollout_to() - How to extract and visualize results

Next: Explore the Core Concepts below to understand the architecture, or jump to Soft Robot Examples for more advanced simulations.

🧠 Core Concepts¶

🏗️ Object-Oriented Architecture¶

SoRoMoX uses an object-oriented design based on Equinox dataclasses. Systems are instantiated directly as classes:

System Creation

from soromox.systems import (
    ArticulatedSoftRobot,
    ArticulatedSoftRobotParams,
    LinkSpec,
    PCS,
)
# For spatial articulated soft robots
robot = ArticulatedSoftRobot(params=ArticulatedSoftRobotParams(...))

# For soft continuum robots (PCS and GVS)
robot = PCS.from_links([LinkSpec.circular(...)], backend="auto")

backend="auto" is already the default, so it can normally be omitted. It keeps supported operations on JAX/XLA on CPU and selects Warp for supported GPU kinematics and dynamics. The choice follows normal calls to forward_kinematics(), dynamics_terms(), forward_dynamics(), and the corresponding batched methods. A direct query can be compared with robot.dynamics_terms(q, qd, backend="jax") without rebuilding the model. See Execution Devices and Backends for the supported systems, automatic-selection rules, differentiation behavior, and performance settings.

Key Benefits:

Extensibility: Easy to subclass and modify methods (e.g., custom actuation mappings)
JAX Compatibility: Full support for JIT compilation, automatic differentiation, and vectorization
Type Safety: Clear interfaces with static type checking
Performance: Optimized computation graphs compiled at runtime

📋 Typed Parameters¶

Each robot system expects a typed Equinox PyTree params object. Numeric fields are JAX arrays, so same-shape updates can flow through jit, grad, and vmap without changing the compiled structure.

When omitted, the robot constructor's base_pose mounts a fixed robot upright (+y planar, +z spatial), and gravity uses Earth gravity in the negative vertical world direction. Gravity is an inertial-frame vector and is not rotated with the base. Use poses.planar_mounting_pose(...) or poses.spatial_mounting_pose(...) for named horizontal, upright, and hanging mountings. Explicit arrays remain available for custom poses, gravity directions, and zero-gravity models. See Parameters for exact vectors and quaternions.

Parameter Structure

links = [
    LinkSpec.circular(
        length=length,
        radius=radius,
        density=density,
        reference_strain=reference,
        young_modulus=young,
        shear_modulus=shear,
        material_damping_coefficient=material_damping,
    )
    for length, radius, density in zip(link_lengths, radii, densities)
]
robot = PCS.from_links(links)
robot = robot.update_link_params(length=new_lengths)

For routed tendons, pushing rods, muscles, and equivalent pressure chambers, continue with the dedicated Threadlike actuation guide.

Units Matter

Always use consistent SI units (meters, kilograms, seconds) for reliable results.

🔄 State Representation¶

Robot states are represented using the SystemState dataclass, which encapsulates time, state vector, and optional actuation/control state:

SystemState

from soromox.systems import SystemState

# For an n-DOF system:
positions = q      # Shape: (n,)
velocities = qd    # Shape: (n,)
y = jnp.concatenate([q, qd])  # State vector, shape: (2n,)

# Create a SystemState instance
state = SystemState(
    t=0.0,                    # Current time
    y=y,                      # State vector [q, qd]
    u=None,                   # Optional: actuation input
    control_state=None        # Optional: controller state (e.g., integrator terms)
)

# Access state components
time = state.t
state_vector = state.y
q, qd = jnp.split(state.y, 2)  # Split into positions and velocities

🤖 Soft Robot Example¶

Ready for something more advanced? Let's simulate a soft continuum robot:

🌊 Continuum Robot📐 Jacobian Analysis

from soromox.systems import LinkSpec, PlanarPCS

# Create a 3-segment soft robot
num_segments = 3
links = [
    LinkSpec.circular(
        length=0.1,
        radius=0.02,
        density=1070.0,
        reference_strain=[0.0, 1.0, 0.0],
        young_modulus=2e3,
        shear_modulus=1e3,
        material_damping_coefficient=318.0,
    )
    for _ in range(num_segments)
]
# Note: Damping helps stabilize simulations and represents material dissipation.

# Initialize the PCS robot
# "auto" uses Warp for supported GPU kinematics and dynamics, and JAX otherwise.
# It is the default, but is shown explicitly here for clarity.
robot = PlanarPCS.from_links(links, backend="auto")

# Define configuration (strains)
# Each segment has 3 strain components: [curvature, shear_x, shear_y]
q = jnp.array([0.1, 0.0, 0.0,  # Segment 1: [curvature, shear_x, shear_y]
               0.2, 0.0, 0.0,  # Segment 2
               0.3, 0.0, 0.0]) # Segment 3

# Compute forward kinematics along the robot
s_values = jnp.linspace(0, robot.L.sum(), 100)
backbone_shape = robot.forward_kinematics_abscissa_batched(q, s_values)

# Extract positions for plotting
# The batched method returns [theta, px, py] at each abscissa.
x_positions = backbone_shape[:, 1]  # X coordinates
y_positions = backbone_shape[:, 2]  # Y coordinates

# Visualize the robot shape
plt.figure(figsize=(10, 6))
plt.plot(x_positions, y_positions, 'b-', linewidth=4, label='Robot Backbone')
plt.scatter(x_positions[0], y_positions[0], s=150, color='green', 
            marker='o', label='Base', zorder=5)
plt.scatter(x_positions[-1], y_positions[-1], s=150, color='red', 
            marker='s', label='Tip', zorder=5)
plt.xlabel('X Position [m]')
plt.ylabel('Y Position [m]')
plt.title('🌊 Soft Continuum Robot Shape')
plt.legend()
plt.grid(True, alpha=0.3)
plt.axis('equal')
plt.show()

# Compute Jacobian at the end-effector
s_tip = robot.L.sum()  # End of the robot
J = robot.jacobian(q, s_tip)

print(f"🔍 Jacobian shape: {J.shape}")
print(f"📏 Jacobian matrix:\n{J}")

# Analyze manipulability
manipulability = jnp.sqrt(jnp.linalg.det(J @ J.T))
print(f"💪 Manipulability index: {manipulability:.4f}")

# Compute workspace boundary
theta_range = jnp.linspace(0, 2*jnp.pi, 100)
workspace_boundary = []

for theta in theta_range:
    # Unit direction in task space
    direction = jnp.array([jnp.cos(theta), jnp.sin(theta), 0.0])

    # Compute maximum reach in this direction
    # (This is a simplified analysis - real workspace computation is more complex)
    max_reach = jnp.linalg.norm(J @ direction)
    workspace_boundary.append(max_reach * direction[:2])

workspace_boundary = jnp.array(workspace_boundary)

# Plot workspace
plt.figure(figsize=(8, 8))
plt.plot(workspace_boundary[:, 0], workspace_boundary[:, 1], 
         'r--', linewidth=2, label='Approximate Workspace')
plt.scatter(x_positions[-1], y_positions[-1], s=150, color='blue', 
            marker='o', label='Current Tip Position')
plt.xlabel('X Position [m]')
plt.ylabel('Y Position [m]')
plt.title('🎯 Robot Workspace Analysis')
plt.legend()
plt.grid(True, alpha=0.3)
plt.axis('equal')
plt.show()

🎯 Next Steps¶

📚 Explore Examples: Dive deeper with examples and tutorials covering all robot types and use cases.
📖 API Reference: Complete documentation of all classes, methods, and functions.
🤝 Contributing: Learn how to contribute to SoRoMoX.
🔬 Advanced Topics: Control theory, optimization, and machine learning applications with SoRoMoX.

💡 Pro Tips¶

Performance Optimization

🚀 JIT Compilation📊 Vectorization🎯 Gradient Computation

import jax

# Compile functions for faster execution
fast_kinematics = jax.jit(robot.forward_kinematics)
fast_jacobian = jax.jit(robot.jacobian)

# Process multiple configurations at once
q_batch = jnp.array([[0.1, 0.0, 0.0],
                    [0.2, 0.0, 0.0],
                    [0.3, 0.0, 0.0]])

# Vectorized computation
batch_kinematics = jax.vmap(
    lambda q: robot.forward_kinematics(q, s_tip)
)(q_batch)

# Automatic differentiation for optimization
def objective(q):
    pos = robot.forward_kinematics(q, s_tip)
    target = jnp.array([1.0, 0.5, 0.0])  # Target pose
    return jnp.sum((pos - target)**2)

# Compute gradients
grad_fn = jax.grad(objective)
gradient = grad_fn(q)

Common Pitfalls

Numerical Precision: Enable double precision with jax.config.update("jax_enable_x64", True)
Unstable Simulations: Use a smaller timestep and a higher order solver to prevent numerical instability and improve accuracy of the simulation.
Slow Batched Simulations: If batched simulations are slow, try using a smaller batch size and execute on a desktop-size GPU.
Automatic Differentiation: Please use automatic differentiation with care as it can be very slow and memory-intensive - in particular, when differentiating through the simulation integration.