dair_pll.drake_experiment

Wrappers for Drake/ContactNets multibody experiments.

class dair_pll.drake_experiment.DrakeSystemConfig(urdfs=<factory>)[source]

Bases: SystemConfig

urdfs: typing.Dict[str, str]

class dair_pll.drake_experiment.MultibodyLosses(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

PREDICTION_LOSS = 1

CONTACTNETS_LOSS = 2

class dair_pll.drake_experiment.MultibodyLearnableSystemConfig(urdfs=<factory>, loss=MultibodyLosses.PREDICTION_LOSS)[source]

Bases: DrakeSystemConfig

loss: dair_pll.drake_experiment.MultibodyLosses = 1: Whether to use ContactNets or prediction loss.

class dair_pll.drake_experiment.DrakeMultibodyLearnableExperimentConfig(data_config=<factory>, base_config=<factory>, learnable_config=<factory>, optimizer_config=<factory>, storage='./', run_name='experiment_run', run_wandb=True, wandb_project=None, full_evaluation_period=1, full_evaluation_samples=5, update_geometry_in_videos=False, visualize_learned_geometry=True)[source]

Bases: SupervisedLearningExperimentConfig

visualize_learned_geometry: bool = True: Whether to use learned geometry in trajectory overlay visualization.

class dair_pll.drake_experiment.DrakeExperiment(config)[source]

Bases: SupervisedLearningExperiment, ABC

visualization_system: typing.Optional[dair_pll.drake_system.DrakeSystem]

base_drake_system: typing.Optional[dair_pll.drake_system.DrakeSystem]

get_drake_system()[source]

Return type:: DrakeSystem

get_base_system()[source]

Abstract callback function to construct base system from system config.

Return type:: System
Returns:: Experiment’s base system.

get_learned_drake_system(learned_system)[source]

If possible, constructs a DrakeSystem -equivalent model of the given learned system, such as when the learned system is a MultibodyLearnableSystem.

Parameters:: learned_system (System) – System being learned in experiment.
Return type:: Optional[DrakeSystem]
Returns:: Drake version of learned system.

visualizer_regeneration_is_required()[source]

Checks if visualizer should be regenerated, e.g. if learned geometries have been updated and need to be pushed to the visulizer.

Return type:: bool

get_visualization_system(learned_system)[source]

Generate a dummy DrakeSystem for visualizing comparisons between trajectories generated by the base system and something else, e.g. data.

Implemented as a thin wrapper of vis_utils.generate_visualization_system(), which generates a drake system where each model in the base DrakeSystem has a duplicate, and visualization elements are repainted for visual distinction.

Parameters:: learned_system (System) – Current trained learnable system.
Return type:: DrakeSystem
Returns:: New DrakeSystem with doubled state and repainted elements.

base_and_learned_comparison_summary(statistics, learned_system)[source]

Extracts a SystemSummary that compares the base system to the learned system.

For Drake-based experiments, this comparison is implemented as overlaid videos of corresponding ground-truth and predicted trajectories. The nature of this video is described further in dair_pll.vis_utils.

Parameters:

statistics (Dict) – Dictionary of training statistics.
learned_system (System) – Most updated version of learned system during training.

Return type:

SystemSummary

Returns:

Summary containing overlaid video(s).

class dair_pll.drake_experiment.DrakeDeepLearnableExperiment(config)[source]: Bases: DrakeExperiment, DeepLearnableExperiment

class dair_pll.drake_experiment.DrakeMultibodyLearnableExperiment(config)[source]

Bases: DrakeExperiment

get_learned_system(_)[source]

Abstract callback function to construct learnable system for experiment.

Optionally, learned system can be initialized to depend on the training dataset.

Parameters:: train_states – (*, space.n_x) batch of all states in training set.
Return type:: MultibodyLearnableSystem
Returns:: Experiment’s learnable system.

visualizer_regeneration_is_required()[source]

Checks if visualizer should be regenerated, e.g. if learned geometries have been updated and need to be pushed to the visulizer.

Return type:: bool

get_learned_drake_system(learned_system)[source]

If possible, constructs a DrakeSystem -equivalent model of the given learned system, such as when the learned system is a MultibodyLearnableSystem.

Parameters:: learned_system (System) – System being learned in experiment.
Return type:: Optional[DrakeSystem]
Returns:: Drake version of learned system.

contactnets_loss(x_past, x_future, system, keep_batch=False)[source]

LossCallbackCallable which applies the ContactNets [1] loss to the system.

References

[1] S. Pfrommer*, M. Halm*, and M. Posa. “ContactNets: Learning Discontinuous Contact Dynamics with Smooth, Implicit Representations,” Conference on Robotic Learning, 2020, https://proceedings.mlr.press/v155/pfrommer21a.html

Return type:: Tensor