MLP Dynamics

Neural-network-based vehicle dynamics model.

Notes

Implements a simple MLP dynamics model registered via ModelSpec for integration with the Physics-Constrained Learning (PCT) optimization stack.

Param

Parameters for the neural network vehicle dynamics model.

Supports scalar parameters; create can generate sized defaults.

Attributes:

Name	Type	Description
W1	Array	First layer weights shaped (hidden_size, STATE_DIM + CONTROL_DIM).
b1	Array	First layer bias shaped (hidden_size,).
W2	Array	Second layer weights shaped (STATE_DIM, hidden_size).
b2	Array	Second layer bias shaped (STATE_DIM,).
alpha	Array	Activation scaling factor.
timestep	float	Physics integration timestep.
timestep_ratio	int	Simulation steps per control step (static field).
hidden_size	int	Hidden layer width (static field).

create classmethod

create(
    *,
    hidden_size=DEFAULT_HIDDEN_SIZE,
    alpha=1.0,
    timestep=0.05,
    timestep_ratio=1,
)

Construct a parameter bundle with the requested hidden size.

build_base_params_with_hidden

build_base_params_with_hidden(
    timestep_hundredths, *, hidden_size=DEFAULT_HIDDEN_SIZE
)

Construct a Param instance with the desired hidden size.

default_params

default_params(hidden_size=DEFAULT_HIDDEN_SIZE)

Return zero-initialized network parameters for the requested hidden size.

param_bounds

param_bounds(hidden_size=DEFAULT_HIDDEN_SIZE)

Return lower/upper bounds for each parameter.

param_shapes

param_shapes(hidden_size=DEFAULT_HIDDEN_SIZE)

Return parameter shapes keyed by parameter name.

vehicle_dynamics_mlp

vehicle_dynamics_mlp(x_and_u, params)

Evaluate the neural network vehicle dynamics model.

Parameters:

Name	Type	Description	Default
x_and_u	Array	Concatenated state-control vector of shape (STATE_DIM + CONTROL_DIM,).	required
params	Param	Parameters containing the network weights and offsets.	required

Returns:

Type	Description
Array	Concatenated time derivative of state and control vectors.

yaw_normalized_loss

yaw_normalized_loss(output_states, target_states)

Yaw-normalized loss between output and target states.

Handles angular differences by wrapping them to [-pi, pi] and uses L1 loss for non-angular components.

Parameters:

Name	Type	Description	Default
output_states	Array	Predicted states of shape (..., 5).	required
target_states	Array	Target states of shape (..., 5).	required

Returns:

Type	Description
float	Total loss summed over all elements.

yaw_normalized_loss_per_element

yaw_normalized_loss_per_element(
    output_states, target_states
)

Yaw-normalized loss per state dimension.

Parameters:

Name	Type	Description	Default
output_states	Array	Predicted states of shape (B, T, 5).	required
target_states	Array	Target states of shape (B, T, 5).	required

Returns:

Type	Description
Array	Loss per dimension of shape (B, 5).

yaw_normalized_loss_per_item

yaw_normalized_loss_per_item(output_states, target_states)

Yaw-normalized loss per item in the batch.

Assumes input shape is (B, T, D) where D=5.

Parameters:

Name	Type	Description	Default
output_states	Array	Predicted states of shape (B, T, 5).	required
target_states	Array	Target states of shape (B, T, 5).	required

Returns:

Type	Description
Array	Loss per batch item of shape (B,).

yaw_normalized_loss_weighted

yaw_normalized_loss_weighted(
    output_states, target_states, weights=None
)

Weighted yaw-normalized loss for state dimensions.

Parameters:

Name	Type	Description	Default
output_states	Array	Predicted states of shape (..., 5).	required
target_states	Array	Target states of shape (..., 5).	required
weights	Array	Per-dimension weights of shape (5,). Defaults to ones.	None

Returns:

Type	Description
float	Weighted loss value.

Raises:

Type	Description
ValueError	If weights has an unexpected shape.