Head Stabilization¶

Model¶

class flygym.examples.head_stabilization.JointAngleScaler¶

Bases: object

A class for standardizing joint angles (i.e., using mean and standard deviation.)

Attributes:

meannp.ndarray: The mean values used for scaling.
stdnp.ndarray: The standard deviation values used for scaling.

classmethod from_data(joint_angles: ndarray)¶

Create a JointAngleScaler instance from joint angle data. The mean and standard deviation values are calculated from the data.

Parameters:

joint_anglesnp.ndarray: The joint angle data. The shape should be (n_samples, n_joints) where n_samples is, for example, the length of a time series of joint angles.

Returns:

JointAngleScaler: A JointAngleScaler instance.

classmethod from_params(mean: ndarray, std: ndarray)¶

Create a JointAngleScaler instance from predetermined mean and standard deviation values.

Parameters:

meannp.ndarray: The mean values. The shape should be (n_joints,).
stdnp.ndarray: The standard deviation values. The shape should be (n_joints,).

Returns:

JointAngleScaler: A JointAngleScaler instance.

class flygym.examples.head_stabilization.WalkingDataset(sim_data_file: Path, contact_force_thr: Tuple[float, float, float] = (0.5, 1, 3), joint_angle_scaler: Callable | None = None, ignore_first_n: int = 200, joint_mask=None)¶

Bases: Dataset

PyTorch Dataset class for walking data.

Parameters:

sim_data_filePath: The path to the simulation data file.
contact_force_thrTuple[float, float, float], optional: The threshold values for contact forces, by default (0.5, 1, 3).
joint_angle_scalerOptional[Callable], optional: A callable object used to scale joint angles, by default None.
ignore_first_nint, optional: The number of initial data points to ignore, by default 200.
joint_maskOptional, optional: A mask to apply on joint angles, by default None.

Attributes:

gaitstr: The type of gait.
terrainstr: The type of terrain.
subsetstr: The subset of the data, i.e., “train” or “test”.
dn_drivestr: The DN drive used to generate the data.
contact_force_thrnp.ndarray: The threshold values for contact forces.
joint_angle_scalerCallable: The callable object used to scale joint angles.
ignore_first_nint: The number of initial data points to ignore.
joint_maskOptional: The mask applied on joint angles. This is used to zero out certain DoFs to evaluate which DoFs are likely more important for head stabilization.
contains_fly_flipbool: Indicates if the simulation data contains fly flip errors.
contains_physics_errorbool: Indicates if the simulation data contains physics errors.
roll_pitch_tsnp.ndarray: The optimal roll and pitch correction angles. The shape is (n_samples, 2).
joint_anglesnp.ndarray: The scaled joint angle time series. The shape is (n_samples, n_joints).
contact_masknp.ndarray: The contact force mask (i.e., 1 if leg touching the floor, 0 otherwise). The shape is (n_samples, 6).

class flygym.examples.head_stabilization.ThreeLayerMLP¶

Bases: LightningModule

A PyTorch Lightning module for a three-layer MLP that predicts the head roll and pitch correction angles based on proprioception and tactile information.

configure_optimizers()¶: Use the Adam optimizer.

forward(x)¶

Forward pass through the model.

Parameters:

xtorch.Tensor: The input tensor. The shape should be (n_samples, 42 + 6) where 42 is the number of joint angles and 6 is the number of contact masks.

training_step(batch, batch_idx)¶: Training step of the PyTorch Lightning module.

validation_step(batch, batch_idx)¶: Validation step of the PyTorch Lightning module.

class flygym.examples.head_stabilization.HeadStabilizationInferenceWrapper(model_path: Path, scaler_param_path: Path, contact_force_thr: Tuple[float, float, float] = (0.5, 1, 3))¶

Bases: object

Wrapper for the head stabilization model to make predictions on observations. Whereas data are collected in large tensors during training, this class provides a “flat” interface for making predictions one observation (i.e., time step) at a time. This is useful for deploying the model in closed loop.

Utilities¶

flygym.examples.head_stabilization.util.get_head_stabilization_model_paths() → Tuple[Path, Path]¶

Get the paths to the head stabilization models.

Returns:

Path: Path to the head stabilization model checkpoint.
Path: Path to the pickle file containing joint angle scaler parameters.