Parameters, Updates, and Optimization¶

SoRoMoX separates convenient construction specifications, dynamic numeric parameters, and static model structure. This guide shows how to construct PCS and GVS systems, update their parameter PyTrees immutably, and use the same interfaces in gradient-based optimization.

For the mechanical meaning of links, joints, cross-sections, stiffness, damping, and isotropic materials, see Continuum Robot Components.

Choosing the right parameter layer¶

All update methods return new objects. They do not mutate the original robot or parameter PyTree.

Constructing PCS and GVS¶

The shared LinkSpec accepts either isotropic material properties or explicit generalized matrices. Isotropic stiffness may be parameterized by Young's and shear moduli or by Young's modulus and Poisson's ratio:

import jax.numpy as jnp

from soromox.systems import PCS, LinkSpec

young = 1.0e6
reference = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0]

pcs = PCS.from_links([
    LinkSpec.circular(
        length=0.20,
        radius=0.012,
        density=1000.0,
        young_modulus=young,
        poisson_ratio=0.45,
        material_damping_coefficient=1.0e4,
        reference_strain=reference,
    )
])

GVS uses the same link description and adds joint, strain-basis, and quadrature specifications:

from soromox.systems import (
    GVS,
    GVSSegment,
    JointSpec,
    StrainBasisSpec,
)

gvs = GVS.from_segments([
    GVSSegment(
        link=LinkSpec.circular(
            length=0.20,
            radius=0.012,
            density=1000.0,
            young_modulus=young,
            poisson_ratio=0.45,
            material_damping_coefficient=1.0e4,
            reference_strain=reference,
        ),
        joint=JointSpec.revolute(
            axis="z",
            stiffness=jnp.array([[0.3]]),
            damping=jnp.array([[0.02]]),
        ),
        basis=StrainBasisSpec(
            type="legendre",
            strain_selector=("kappa_y", "sigma_x"),
            basis_order=1,
        ),
        num_gauss_points=7,
    )
])

An explicit stiffness= or damping= on LinkSpec bypasses the corresponding material construction:

explicit_pcs = PCS.from_links([
    LinkSpec.circular(
        length=0.15,
        radius=0.01,
        density=1000.0,
        stiffness=jnp.diag(
            jnp.array([0.2, 0.8, 0.8, 100.0, 30.0, 30.0])
        ),
        damping=jnp.diag(
            jnp.array([0.01, 0.02, 0.02, 0.5, 0.1, 0.1])
        ),
        reference_strain=reference,
    )
])

Immutable runtime updates¶

Use the narrowest update method that owns a parameter. PCS, PlanarPCS, and GVS share update_link_params:

stiffness = pcs.params.link.stiffness
stiffness = stiffness.at[0].set(1.1 * stiffness[0])

pcs = pcs.update_link_params(
    stiffness=stiffness,
    damping=0.9 * pcs.params.link.damping,
)

GVS additionally exposes joint updates:

joint_damping = gvs.params.joint.damping
joint_damping = joint_damping.at[0, 0, 0].set(0.03)
gvs = gvs.update_joint_params(damping=joint_damping)

Use nested .replace(...) and with_params(...) when replacing multiple components atomically:

replacement = gvs.params.replace(
    link=gvs.params.link.replace(
        density=1.05 * gvs.params.link.density,
        stiffness=1.1 * gvs.params.link.stiffness,
    ),
    joint=gvs.params.joint.replace(
        stiffness=1.2 * gvs.params.joint.stiffness,
        damping=1.1 * gvs.params.joint.damping,
    ),
)
gvs = gvs.with_params(replacement)

Top-level values use update_params:

pcs = pcs.update_params(gravity=jnp.zeros(3))

A fixed mounting pose is a dynamic model leaf rather than a field in the physical parameter PyTree. Update it independently with with_fixed_base_pose(...):

candidate_base_pose = pcs.fixed_base_pose.at[4].add(0.01)
mounted_pcs = pcs.with_fixed_base_pose(candidate_base_pose)

This keeps one physical parameter PyTree reusable across different mountings and between fixed and floating models without making the fixed pose static or nondifferentiable.

Floating systems instead carry their base pose in the runtime configuration and reject with_fixed_base_pose(...).

Updates may change numeric values but not static layout. Changing the number of links, GVS basis orders, active strain selectors, joint types, matrix padding, or quadrature padding requires reconstruction.

Compiled update contract¶

with_params(...) can participate in JAX transformations when the model's static structure remains fixed. Close over the constructed robot and pass its complete parameter PyTree as a dynamic argument:

import equinox as eqx

q = jnp.zeros(pcs.num_coordinates)

@eqx.filter_jit
def energy(params, configuration):
    candidate = pcs.with_params(params)
    return candidate.potential_energy(configuration)

baseline = energy(pcs.params, q)
candidate_params = pcs.params.replace(
    link=pcs.params.link.replace(
        density=1.05 * pcs.params.link.density,
        stiffness=0.95 * pcs.params.link.stiffness,
    )
)
candidate = energy(candidate_params, q)

Fixed mounting poses can likewise be supplied as dynamic optimization variables without reconstructing the model or recompiling for each value:

import equinox as eqx
import jax
import jax.numpy as jnp

target_position = jnp.array([0.1, -0.05, 0.25])

@eqx.filter_jit
def mounted_loss_and_grad(base_pose, configuration):
    def loss(pose):
        candidate = pcs.with_fixed_base_pose(pose)
        tip_pose = candidate.forward_kinematics(configuration, candidate.length)
        return jnp.sum((tip_pose[:3, 3] - target_position) ** 2)

    return jax.value_and_grad(loss)(base_pose)

loss, base_pose_gradient = mounted_loss_and_grad(pcs.fixed_base_pose, q)

Repeated calls share one compiled executable while the pose retains the same shape and dtype. Use a constrained rotation parameterization or normalize the quaternion when optimizing a spatial mounting. To identify physical parameters and mounting jointly, pass both params and base_pose as dynamic arguments and chain with_params(params).with_fixed_base_pose(base_pose).

The two calls share one compiled executable when every parameter leaf keeps the same PyTree position, shape, and dtype. Numeric system, environment, actuator, and passive-element values are dynamic under this contract. Updates remain differentiable and vectorizable with grad, value_and_grad, and vmap.

Installed components use the corresponding immutable delegates: with_actuator_params(...) and with_passive_element_params(...). The same fixed-layout rule applies to their parameter PyTrees.

The distinction between dynamic params and static structure is described in the parameter API Overview. Changing structure, or changing a parameter leaf's shape, dtype, or None/array presence, requires reconstruction and normally a new compilation.

Eager updates perform both structural and value validation. During tracing, only validation based on statically observable types and shapes is possible; runtime numeric candidates must therefore remain in their physical domains. Use constrained parameterizations where appropriate, and validate externally supplied values eagerly before a compiled optimization or simulation loop. Parameter subclasses express this boundary through separate tracer-safe validate_structure() and concrete-only validate_values() hooks. The framework owns the composed validate() and dispatched validate_for_update() methods. Closed-over concrete values encountered during JAX tracing are checked at trace time, while value checks for dynamic tracer leaves remain deferred.

Start from robot.params so constructor-resolved defaults already have their runtime PyTree representation.

Isotropic material updates¶

Generalized stiffness and damping matrices are the canonical runtime values. Young's modulus, either shear modulus or Poisson's ratio, and optional material damping remain convenient construction and optimization variables in a separate caller-owned PyTree:

from soromox.systems import IsotropicMaterialParams

material = IsotropicMaterialParams(
    young_modulus=jnp.array([1.0e6]),
    shear_modulus=jnp.array([3.4e5]),
    material_damping_coefficient=jnp.array([1.0e4]),
)

material = material.replace(
    young_modulus=1.1 * material.young_modulus,
    shear_modulus=0.95 * material.shear_modulus,
    material_damping_coefficient=(
        1.2 * material.material_damping_coefficient
    ),
)
updated_pcs = pcs.with_isotropic_material(material)

The same call works for PlanarPCS and GVS. Scalar fields are broadcast; arrays must contain one value per link. The robot stores only the generated canonical matrices, not material.

Use poisson_ratio instead of shear_modulus to optimize the alternative isotropic parameterization. Exactly one must be supplied:

poisson_material = IsotropicMaterialParams(
    young_modulus=jnp.array([1.0e6]),
    poisson_ratio=jnp.array([0.45]),
)

Here material damping is omitted, so applying the material produces zero link damping. The presence of the optional fields is part of the PyTree structure: JIT compilation specializes on the chosen stiffness parameterization and on whether damping is active, while gradients are traced through only the active numeric fields.

Inspect the mapping without updating the robot:

stiffness, damping = pcs.link_matrices_from_material(material)
updated = pcs.with_isotropic_material(material)

assert jnp.allclose(updated.params.link.stiffness, stiffness)
assert jnp.allclose(updated.params.link.damping, damping)

Differentiating parameter updates¶

Typed params and robot updates are JAX PyTrees, so a loss can construct a candidate robot without mutation. This example differentiates all active isotropic material fields:

import jax

robot = pcs.with_isotropic_material(material)
target_stiffness = 1.05 * robot.stiffness_matrix()
target_damping = 0.9 * robot.damping_matrix(jnp.zeros(robot.num_coordinates))

def material_loss(candidate_material):
    candidate = robot.with_isotropic_material(candidate_material)
    stiffness_error = candidate.stiffness_matrix() - target_stiffness
    damping_error = (
        candidate.damping_matrix(jnp.zeros(candidate.num_coordinates))
        - target_damping
    )
    return (
        jnp.mean(stiffness_error**2)
        + jnp.mean(damping_error**2)
    )

value, gradient = jax.value_and_grad(material_loss)(material)
print(value)
print(gradient.young_modulus)
print(gradient.shear_modulus)
print(gradient.material_damping_coefficient)

Cached unit-response operators reduce every material evaluation to batched scalar-matrix combinations:

K = E K_E + G K_G
D = eta D_eta

The same immutable pattern applies to other numeric fields. For example, canonical stiffness can be differentiated directly:

def matrix_loss(link_stiffness):
    candidate = robot.update_link_params(stiffness=link_stiffness)
    return jnp.mean(
        (candidate.stiffness_matrix() - target_stiffness) ** 2
    )

matrix_value, matrix_gradient = jax.value_and_grad(matrix_loss)(
    robot.params.link.stiffness
)

Positive parameterization¶

Optimize material properties in log space when they must remain nonnegative. This constrains the material scalars without imposing positive-definiteness on the canonical generalized matrices. The following applies to the Young/shear/damping parameterization:

log_material = jax.tree.map(jnp.log, material)

def decode(log_values):
    return IsotropicMaterialParams(
        young_modulus=jnp.exp(log_values.young_modulus),
        shear_modulus=jnp.exp(log_values.shear_modulus),
        material_damping_coefficient=jnp.exp(
            log_values.material_damping_coefficient
        ),
    )

For the Young/Poisson parameterization, constrain Poisson's ratio to its open physical interval with a sigmoid transform, for example poisson_ratio = -1 + 1.5 * jax.nn.sigmoid(raw_poisson_ratio). Omitted damping remains None and is not an optimization leaf.

Optax optimization loop¶

The following loop uses the project's examples dependency and keeps the optimized variable as the log-material PyTree:

import optax

optimizer = optax.adam(1.0e-2)
optimizer_state = optimizer.init(log_material)

def log_loss(log_values):
    candidate = robot.with_isotropic_material(decode(log_values))
    return jnp.mean(
        (candidate.stiffness_matrix() - target_stiffness) ** 2
    )

value_and_grad = jax.jit(jax.value_and_grad(log_loss))
for _ in range(100):
    value, gradients = value_and_grad(log_material)
    updates, optimizer_state = optimizer.update(
        gradients,
        optimizer_state,
        log_material,
    )
    log_material = optax.apply_updates(log_material, updates)

final_material = decode(log_material)
optimized_robot = robot.with_isotropic_material(final_material)

robot may be a PCS, PlanarPCS, or GVS instance.

Geometry and material co-optimization¶

Length or cross-section updates refresh unit-response operators but do not silently overwrite canonical matrices. Apply geometry first and material second:

new_section = robot.params.link.cross_section.replace(
    coefficients=(
        1.05 * robot.params.link.cross_section.coefficients
    )
)
geometry_robot = robot.update_link_params(cross_section=new_section)

assert jnp.allclose(
    geometry_robot.params.link.stiffness,
    robot.params.link.stiffness,
)
cooptimized_robot = geometry_robot.with_isotropic_material(material)

In a geometry/material optimization loop, construct geometry_robot from the candidate geometry first, then apply the candidate material. This evaluates the material mapping with the refreshed geometry operators while preserving the explicit update order.

Material variables or direct matrices?¶

Prefer isotropic material optimization when Young's modulus, shear modulus or Poisson's ratio, and optional material damping have physical meaning in the identification problem. Optimize canonical params.link.stiffness and params.link.damping directly for anisotropic, coupled, or learned constitutive models.

PCS matrices have shape (N, 6, 6), PlanarPCS matrices (N, 3, 3), and GVS matrices (N, max_dof, max_dof) with zero padding beyond each link's active coordinates. The update and optimization pattern is otherwise identical.