examples

Default training

To train the model, run these commands:

# to train the pendulum example
python pend_lag_cavae_trainer.py 
# to train the fully-actuated cartpole example
python cart_lag_cavae_trainer.py 
# to train the fully-actuated acrobot example
python acro_lag_cavae_trainer.py 

These commands train model with standard batch generation, where each batch is generated by randomly sampling all the training data. See below for more details.

Training with homogeneous control batch generation and annealing

As expained in the paper (Seciton S3.2), all the results in the main paper are trained with homogeneous control batch generation, where each batch is generated by randomly sampling a part of the training dataset so that all the control inputs in a batch are the same. In practice, an annealing scheme on weight might improve training in this case. To train the models with homogeneous control batch generation and annealing, please run

# to train the pendulum example
python pend_lag_cavae_trainer.py --batch_size 256 --reload_dataloaders_every_epoch True --max_epochs 3000 --homo_u --annealing
# to train the fully-actuated cartpole example
python cart_lag_cavae_trainer.py --batch_size 256 --reload_dataloaders_every_epoch True --max_epochs 3000 --homo_u --annealing
# to train the fully-actuated acrobot example
python acro_lag_cavae_trainer.py --batch_size 256 --reload_dataloaders_every_epoch True --max_epochs 3000 --homo_u --annealing

With these commands, you should be able to get results that are close to the figures (Fig. 3 and 4) shown in the paper.

As for prediction, the difference between homogeneous control batch generation and the standard batch generation in our Lagrangian_cavae model visually indistinguishable, as stated in the paper. I only find that in the Acrobot case the model trained by standard batch generation is hard to cooperate with an energy-based controller, indicating that the physics quantities are not well-learned. For the ablation models, those trained by the standard batch generation seems to perform worse than those trained by homogeneous control batch generation.

Miscellaneous

Due to a bug in torch.nn.functional.grid_sample, you might encounter a segmentation fault if you train the model on GPUs. This bug has not been fixed in the latest PyTorch version (1.6.0) when this work has been done.

However, I successfully trained the pendulum example on GPU without error. Thanks to PyTorch-Lightning, you can train it on GPU with

python examples/pend_lag_cavae_trainer.py --gpus 1

I also find that the tensorboard log files generated by pytorch-lightning 0.9.0 cannot be parsed by tensorboard. Because of this, I did all the training with pytorch-lightning 0.8.5, hopefully they have fixed this issue in the latest release.