Structure of mechanical models: effector states, and aggregation of dynamic and static updates

[mlprt]

I'm uncertain about the design of AbstractPlant and Mechanics.

In brief, an AbstractPlant is a subclass of

1. AbstractDynamicalSystem: AbstractPlant implements a general vector_field method, which aggregates the vector_field methods of all the AbstractDynamicalSystem instances referred to by the abstract property AbstractPlant.dynamics_spec. Final subclasses of AbstractPlant must implement this property.
2. AbstractStagedModel: final subclasses of AbstractPlant must implement a model_spec property that defines the sequence of kinematic/static updates.

On the other hand, a Mechanics instance wraps an AbstractPlant instance in order to discretize it, associate it with a Diffrax solver, and execute all its kinematic and dynamic updates. Currently, it also handles the effector states.

In this issue I discuss this arrangement in more detail. At the end, I raise a couple of potential issues:

• Should the vector fields be aggregated, and updated through a single Diffrax solver step?
• Which class should manage the effector states?

AbstractPlant

AbstractPlant objects provide most of the components of a biomechanical model. In particular, they aggregate the differential equations as well as the kinematic updates the mechanical components.

For example, an AbstractPlant object that models a muscled arm (e.g. feedbax.mechanics.plant.MuscledArm) may bring together:

1. a differential equation describing the physics of skeletal movement due to torques generated by muscles, and belonging to an AbstractSkeleton object;
2. another differential equation describing the dynamics of muscle activation, belonging to an AbstractMuscle object;
3. some kinematic methods belonging to the AbstractPlant object itself, that directly calculate the muscle length and contraction velocity given the current state of the skeleton;
4. a method of the same AbstractMuscle object as in (2) that determines the forces generated by the muscle, given its current length and velocity. These are converted into torques after being passed back to the AbstractPlant object.

AbstractPlant subclasses must implement the property dynamics_spec, which aggregates all the AbstractDynamicalSystem components that should be included in the dynamics step for the entire plant. For example, MuscledArm includes a first-order filter for the muscle activation, as well as an ODE describing the skeletal dynamics, and both of these are referred to in its dynamics_spec:

feedbax/feedbax/mechanics/plant.py

Lines 387 to 404 in 147fb42

	@cached_property
	def dynamics_spec(self) -> Mapping[str, DynamicsComponent[PlantState]]:
	"""Specifies the components of the muscled arm dynamics."""
	return dict(
	{
	"muscle_activation": DynamicsComponent(
	dynamics=self.activator,
	where_input=lambda input, state: input,
	where_state=lambda state: state.muscles.activation,
	),
	#! is this applying the torques twice? since arm will do `input_torque + state.torque`
	"skeleton": DynamicsComponent(
	dynamics=self.skeleton,
	where_input=lambda input, state: state.skeleton.torque,
	where_state=lambda state: state.skeleton,
	),
	}
	)

The entries in dynamics_spec are automatically aggregated into a single field by the method AbstractPlant.vector_field, which fulfils AbstractPlant's implementation of AbstractDynamicalSystem:

feedbax/feedbax/mechanics/plant.py

Lines 101 to 112 in 147fb42

	d_state = jax.tree_map(jnp.zeros_like, state)
	for component in self.dynamics_spec.values():
	d_state = eqx.tree_at(
	component.where_state,
	d_state,
	component.dynamics.vector_field(
	t, component.where_state(state), component.where_input(input, state)
	),
	)
	return d_state

AbstractPlant also inherits from AbstractStagedModel and therefore defines a model_spec, which composes the various static operations that are performed on the plant state (e.g. calculating muscle length and velocity, or clipping states to bounds).

Mechanics

An instance of AbstractPlant describes continuous dynamics, as do the vector fields defined in dynamics_spec and aggregated in vector_field.

Mechanics associates the continuous dynamics of an AbstractPlant instance with a Diffrax solver and a time step, thereby discretizing them. It includes a property _term which returns diffrax.ODETerm(self.plant.vector_field), which is used in the Diffrax solver call:

https://github.com/mlprt/feedbax/blob/147fb42f8793db55d06c8e432e4627a3679a801a/feedbax/mechanics/mechanics.py#L127C1-L135C10

Mechanics inherits solely from AbstractStagedModel[MechanicsState], and when called it:

• updates the plant state by 1) calling the AbstractPlant instance (i.e. the kinematic operations in its model_spec), then 2) executing a step of the dynamics solver.
• syncs the effector state with the configuration state. For example, the configuration state of a typical arm model consists of joint angles and angular velocities, whereas the effector state is typically the Cartesian state of some point on the arm -- the very end of it. The actual methods for doing these calculations are part of the skeleton object.

Here is how these operations are defined in Mechanics.model_spec:

feedbax/feedbax/mechanics/mechanics.py

Lines 83 to 112 in 147fb42

	@property
	def model_spec(self) -> OrderedDict[str, ModelStage]:
	"""Specifies the stages of the model."""
	return OrderedDict(
	{
	"convert_effector_force": ModelStage(
	callable=lambda self: self.plant.skeleton.update_state_given_effector_force,
	where_input=lambda input, state: state.effector.force,
	where_state=lambda state: state.plant.skeleton,
	),
	"kinematics_update": ModelStage(
	# the `plant` module directly implements non-ODE operations
	callable=lambda self: self.plant,
	where_input=lambda input, state: input,
	where_state=lambda state: state.plant,
	),
	"dynamics_step": ModelStage(
	callable=lambda self: self.dynamics_step,
	where_input=lambda input, state: input,
	where_state=lambda state: state,
	),
	"get_effector": ModelStage(
	callable=lambda self: wrap_stateless_callable(
	self.plant.skeleton.effector, pass_key=False
	),
	where_input=lambda input, state: state.plant.skeleton,
	where_state=lambda state: state.effector,
	),
	}
	)

Potential issues

Aggregation of vector fields -> single solver step

Should all the component vector fields be explicitly aggregated in AbstractPlant objects, to be associated with a single Diffrax solver step in Mechanics? It's not generally the case that the fields are coupled, and in principle they could be stepped separately.

Sometimes aggregating the dynamics can lead to confusing situations. For example, see the admonition in the docstring for AbstractMuscle.__call__ -- it's a little conceptually weird to have to include the muscle's ActivationFilter (which models something that would really occur inside a muscle fibre) as part of the AbstractPlant, and not the AbstractMuscle. One potential solution is to have AbstractMuscle inherit from both AbstractStagedModel and AbstractDynamicalSystem, just like AbstractPlant does.

Defining and updating the effector

The loss function is generally calculated in effector (or operational) coordinates. For example, a reaching task defines its reach goals in Cartesian coordinates. Similarly, force perturbations may be applied to the effector as linear forces, which must be converted to torques. However, only the configuration state (e.g. joint angles/angular velocities) is part of the ODE describing the skeleton dynamics. It does not make sense to return derivatives for the effector state through the Diffrax solver.

That is why the effector state is currently part of MechanicsState and not of AbstractSkeletonState. It may be best to change that, as all the relevant kinematic methods belong to an AbstractSkeleton, and the effector state is rightly just an alternative representation of the skeleton state specifically (assuming we won't be modeling soft tissue). Of course, AbstractSkeletonState could be separated into configuration and effector fields, and dynamics_spec could ensure that only state.configuration is passed to the skeleton's vector_field.

Once effector state is removed from MechanicsState, it will only contain two fields: plant and solver. I think it's kind of annoying that the user has to refer to state.mechanics.plant.skeleton instead of state.mechanics.skeleton, and I'd prefer if Mechanics could work as a kind of wrapper that operates on PlantState directly. However, the solver state needs to be passed from each time step to the next. The solver could be moved to PlantState, but that creates a bit of an asymmetry between the model and state hierarchies, given that numerical solution is the responsibility of of Mechanics and not AbstractPlant.

Which point on the skeleton is considered the effector may vary between tasks or experiments. So far, it's always been the endpoint of an arm, or (trivially) a point mass. There may also be models with multiple effectors. One way to deal with complexity here is to 1) keep all the Cartesian states for the skeleton, and 2) use generalized target specifications (#10) so the user can define losses for arbitrary parts of those states.