Learning Physical Dynamics with Action-Angle Networks

The official JAX implementation of Action-Angle Networks from our paper "Learning Integrable Dynamics with Action-Angle Networks".

Instructions

Clone the repository:

git clone https://github.com/ameya98/ActionAngleNetworks.git
cd ActionAngleNetworks

Create and activate a virtual environment:

python -m venv .venv && source .venv/bin/activate

Install dependencies with:

pip install --upgrade pip && pip install -r requirements.txt

Start training with a configuration defined under action_angle_networks/configs/:

python -m action_angle_networks.main --config=action_angle_networks/configs/harmonic_motion/action_angle_flow.py --workdir=./tmp/

Changing Hyperparameters

Since the configuration is defined using config_flags, you can override hyperparameters. For example, to change the number of training steps and the batch size:

python -m action_angle_networks.main --config=action_angle_networks/configs/harmonic_motion/action_angle_flow.py --workdir=./tmp/ \
--config.num_train_steps=10 --config.batch_size=100

For more extensive changes, you can directly edit the configuration files, and even add your own.

Note

This is a fork of my original implementation here.

Code Authors

Ameya Daigavane