The official JAX implementation of Action-Angle Networks.
Clone the repository:
git clone https://github.com/google-research/google-research.git
cd google-research/action_angle_networks
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 configs/
:
python main.py --workdir=./tmp --config=configs/action_angle_flow.py
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 main.py --workdir=./tmp --config=configs/action_angle_flow.py \
--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.
Ameya Daigavane