Bodies

These are the highest-level single time step models included in Feedbax. They are compositions of neural networks, mechanical models, delayed channels, and other components. Typically they model a single step through an entire sensory-motor loop.

A stereotypical example is 1) a forward pass of a neural network (AKA controller, agent, or policy) results in motor commands (AKA controls, or actions), which are 2) sent to a biomechanical model (AKA plant, or environment) whose state is updated by integrating some differential equations, and 3) sensory feedback about the state is returned so that it can bee passed to the neural network on the next iteration.

Currently, there is only one such model: SimpleFeedback.

---

feedbax.bodies.SimpleFeedbackState (Module)

Type of state PyTree operated on by SimpleFeedback instances.

Attributes:

Name	Type	Description
mechanics	MechanicsState	The state PyTree for a Mechanics instance.
net	NetworkState	The state PyTree for a staged neural network.
feedback	PyTree[ChannelState]	A PyTree of state PyTrees for each feedback channel.

feedbax.bodies.SimpleFeedback (AbstractStagedModel[SimpleFeedbackState])

Model of one step around a feedback loop between a neural network and a mechanical model.

Attributes:

Name	Type	Description
net	AbstractModel[NetworkState]	The neural network that outputs commands for the mechanical model.
mechanics	Mechanics	The discretized model of plant dynamics.
feedback_channels	MultiModel[ChannelState]	A PyTree of feedback channels which may be delayed and noisy.

model_spec: OrderedDict[str, ModelStage[Self, SimpleFeedbackState]] property

Specifies the stages of the model in terms of state operations.

memory_spec: PyTree[bool] property

Specifies which states should typically be remembered.

For example, ForgetfulIterator stores trajectories of states. However it doesn't usually make sense to store states.feedback.queue for every timestep, because it contains info that is already available if states.mechanics is stored at every timestep. If the feedback delay is 5 steps, ForgetfulIterator could end up with 5 extra copies of all the parts of states.mechanics that are part of the feedback. So it may be better not to store states.feedback.queue.

This property will be used by ForgetfulIterator, but will be ignored by Iterator, which remembers the full state indiscriminately—it's faster, but may use more memory.

bounds: PyTree[StateBounds] property

Specifies the bounds of the state.

__call__ (input: ModelInput,state: StateT,key: PRNGKeyArray) -> StateT

Inherited from feedbax._staged.AbstractStagedModel

__init__ (net: AbstractModel[NetworkState],mechanics: Mechanics,feedback_spec: Union[PyTree[ChannelSpec], PyTree[Mapping[str, Any]]] = ChannelSpec(where=lambda mechanics_state: mechanics_state.plant.skeleton),motor_delay: int = 0,motor_noise_func: Callable[[PRNGKeyArray, Array], Array] = Normal(),intervenors: Optional[ArgIntervenors] = None,*,key: Optional[PRNGKeyArray] = None)

Parameters:

Name	Type	Description	Default
net	AbstractModel[NetworkState]	The neural network that outputs commands for the mechanical model.	required
mechanics	Mechanics	The discretized model of plant dynamics.	required
feedback_spec	Union[PyTree[ChannelSpec], PyTree[Mapping[str, Any]]]	Specifies the sensory feedback channel(s); i.e. the delay and noise on the states available to the neural network.	ChannelSpec(where=lambda mechanics_state: skeleton)
motor_delay	int	The number of time steps to delay the neural network output sent to the mechanical model.	0
motor_noise_std		The standard deviation of the Gaussian noise added to the neural network's output.	required
intervenors	Optional[ArgIntervenors]	Intervenors to add to the model at construction time.	None

init (*,key: PRNGKeyArray) -> SimpleFeedbackState

Return a default state for the model.

get_nn_input_size (task: AbstractTask,mechanics: Mechanics,feedback_spec: Union[PyTree[ChannelSpec[MechanicsState]], PyTree[Mapping[str, Any]]] = ChannelSpec[MechanicsState](where=lambda mechanics_state: mechanics_state.plant.skeleton)) -> int

staticmethod

Determine how many scalar input features the neural network needs.

This is a static method because its logic (number of network inputs = number of task inputs + number of feedback inputs from mechanics) is related to the structure of SimpleFeedback. However, it is not an instance method because we want to construct the network before we construct SimpleFeedback.

state_consistency_update (state: SimpleFeedbackState) -> SimpleFeedbackState

Returns a corrected initial state for the model.

1. Update the plant configuration state, given that the user has initialized the effector state.
2. Fill the feedback queues with the initial feedback states. This is less problematic than passing all zeros until the delay elapses for the first time.