├── requirements.txt ├── .gitmodules ├── .gitignore ├── .colorize ├── train_many ├── datasets ├── imagepairs.py ├── images.py └── pendulum.py ├── train_vae_simple ├── models ├── train_many ├── simple.py ├── vae_trajectory.py ├── vae.py ├── sympy2torch.py ├── pendulum_energy.py └── stabledynamics.py ├── train_pendulum_simple ├── train_vae_stable ├── train_pendulum_stable ├── train_pendulum_stable_knownV ├── README.md ├── render.py ├── render_vae ├── plot_data.py ├── trajectory.py ├── vae_trajmodel.py ├── test_lstm.py ├── util.py ├── pendulum_error.py ├── train_lstm.py └── train.py /requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | numpy 3 | torch 4 | torchvision 5 | tensorflow 6 | tensorboardX 7 | scikit-learn 8 | sympy 9 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "models/ritheshkumar95_vqvae"] 2 | path = models/ritheshkumar95_vqvae 3 | url = git@github.com:ritheshkumar95/pytorch-vqvae.git 4 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .python-version 2 | .vscode/ 3 | *.log 4 | 5 | youtube 6 | experiments/ 7 | pendulum-cache/ 8 | runs/ 9 | 10 | **/__pycache__/ 11 | .ipynb_checkpoints 12 | -------------------------------------------------------------------------------- /.colorize: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ("$@" 2>&1>&3 | sed -e $'s,\[[0-9:]* INFO.*,\e[93m&\e[m,' -e $'s,\[[0-9:]* INFO @root].*,\e[36m&\e[m,' -e $'s,\[[0-9:]* WARNING.*,\e[91m&\e[m,' -e $'s,\[[0-9:]* ERROR.*,\e[41m&\e[m,' >&2) 3>&1 | sed $'s,.*,\e[m&\e[m,' 4 | -------------------------------------------------------------------------------- /train_many: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | for ALPHA in 0.001 0.0001 4 | do 5 | for INNER_EPSILON in 0.01 0.005 0.001 6 | do 7 | for SMOOTH_V in 0 8 | do 9 | for REHU in 0.005 0.001 0.0005 10 | do 11 | ./train_pendulum_stable stable "$1" "$ALPHA" NN-REHU "$INNER_EPSILON" "$SMOOTH_V" "$REHU" & 12 | done 13 | wait 14 | done 15 | done 16 | done 17 | -------------------------------------------------------------------------------- /datasets/imagepairs.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | from . import images 4 | import torch 5 | 6 | class SeqPairs(torch.utils.data.Dataset): 7 | def __init__(self, image_dataset): 8 | super().__init__() 9 | self.image_dataset = image_dataset 10 | 11 | def __getitem__(self, index): 12 | q1, _ = self.image_dataset[index] 13 | q2, _ = self.image_dataset[index + 1] 14 | return ((q1, q2), (q1, q2)) 15 | 16 | def __len__(self): 17 | return len(self.image_dataset) - 1 18 | 19 | def build(props): 20 | return SeqPairs(images.build(props)) 21 | -------------------------------------------------------------------------------- /train_vae_simple: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | DATASET="$1" 4 | LR="0.0005" 5 | W1="-0.25" 6 | INNER="PSICNN" 7 | 8 | if [ -z "$DATASET" ]; then 9 | echo "NO DATASET" 10 | exit 11 | fi 12 | 13 | OUTDIR="experiments/vae_simple/${DATASET}" 14 | mkdir -p "$OUTDIR" 15 | MODEL="vae_trajectory[w=$W1,projfn=$INNER]" 16 | 17 | date >> "$OUTDIR/progress.txt" 18 | ./.colorize ./train.py \ 19 | --log-to "runs/$OUTDIR" \ 20 | --batch-size 300 \ 21 | --learning-rate "$LR" \ 22 | --epochs 150 \ 23 | imagepairs[files=youtube/$DATASET/*] \ 24 | "$MODEL" \ 25 | "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt" 26 | 27 | ./render_vae "$DATASET" "$MODEL" "$OUTDIR" 28 | -------------------------------------------------------------------------------- /models/train_many: -------------------------------------------------------------------------------- 1 | ./train_pendulum_stable_knownV knownv 1 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05 2 | ./train_pendulum_stable_knownV knownv 2 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05 3 | ./train_pendulum_stable_knownV knownv 4 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05 4 | ./train_pendulum_stable_knownV knownv 6 0.001 ActualPendulumEnergy NO_EPSILON 0 0.05 5 | 6 | ./train_pendulum_simple 1 7 | ./train_pendulum_simple 2 8 | ./train_pendulum_simple 4 9 | ./train_pendulum_simple 6 10 | ./train_pendulum_simple 8 11 | 12 | ./train_pendulum_stable stable 1 0.001 PSD-REHU 0.001 0 0.005 13 | ./train_pendulum_stable stable 2 0.001 PSD-REHU 0.001 0 0.005 14 | ./train_pendulum_stable stable 4 0.001 PSD-REHU 0.001 0 0.005 15 | ./train_pendulum_stable stable 6 0.001 PSD-REHU 0.001 0 0.005 16 | ./train_pendulum_stable stable 8 0.001 PSD-REHU 0.001 0 0.005 17 | -------------------------------------------------------------------------------- /models/simple.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import logging 4 | 5 | import torch 6 | from torch import nn 7 | 8 | logger = logging.getLogger(__name__) 9 | 10 | global SIZE_A, SIZE_B, model 11 | SIZE_A = 320 12 | SIZE_B = 320 13 | 14 | model = None 15 | 16 | loss_ = nn.MSELoss() 17 | loss = lambda Ypred, Yactual, X, **kw: loss_(Ypred, Yactual) 18 | 19 | def loss_flatten(l): 20 | return [l] 21 | 22 | def loss_labels(): 23 | return ["loss"] 24 | 25 | def summary(*a, **kw): 26 | pass 27 | 28 | def configure(props): 29 | global SIZE_A, SIZE_B, model 30 | if "a" in props: 31 | SIZE_A = int(props["a"]) 32 | 33 | if "b" in props: 34 | SIZE_B = int(props["b"]) 35 | 36 | model = nn.Sequential( 37 | nn.Linear(SIZE_A, SIZE_B), nn.ReLU(), 38 | nn.Linear(SIZE_B, SIZE_B), nn.ReLU(), 39 | nn.Linear(SIZE_B, SIZE_A)) 40 | 41 | logger.info(f"Set layer sizes to {SIZE_A} -> {SIZE_B} -> {SIZE_B} -> {SIZE_A}") 42 | 43 | -------------------------------------------------------------------------------- /train_pendulum_simple: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | N="$1" 5 | 6 | # Latent space dim: 7 | LSD=$(( 2 * N )) 8 | 9 | OUTDIR="experiments/pendulum-simple/${N}" 10 | MODEL="simple[a=$LSD,b=116]" 11 | 12 | mkdir -p "$OUTDIR" 13 | echo $MODEL > "$OUTDIR/model" 14 | 15 | date >> "$OUTDIR/progress.txt" 16 | ./.colorize ./train.py \ 17 | --log-to "runs/$OUTDIR" \ 18 | --batch-size 2000 \ 19 | --learning-rate "0.001" \ 20 | --epochs 1000 \ 21 | --test-with "pendulum[n=$N,test]" \ 22 | "pendulum[n=$N]" \ 23 | "$MODEL" \ 24 | "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt" 25 | 26 | ./.colorize ./pendulum_error.py \ 27 | pendulum[n=$N,test] \ 28 | "$MODEL" \ 29 | "$OUTDIR/checkpoint_*.pth" \ 30 | 1000 | tee -a "$OUTDIR/eval.txt" 31 | -------------------------------------------------------------------------------- /train_vae_stable: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e 4 | 5 | EXPERIMENT_ID="$1" 6 | DATASET="$2" 7 | W1="$3" 8 | W2="$4" 9 | INNER="$5" 10 | ALPHA="0.0001" 11 | REHU="0.0001" 12 | SCALE_FX="$6" 13 | 14 | if [ -z "$EXPERIMENT_ID" ]; then 15 | echo "NO EXPERIMENT ID" 16 | exit 17 | elif [ -z "$DATASET" ]; then 18 | echo "NO DATASET" 19 | exit 20 | fi 21 | 22 | OUTDIR="experiments/${EXPERIMENT_ID}/${DATASET}_${W1}_${W2}_${INNER}_${ALPHA}_${REHU}_scaled${SCALE_FX}" 23 | MODEL="vae_trajectory[stable,w=$W1,bce=$W2,a=$ALPHA,projfn=$INNER,h=320,hp=200,rehu=$REHU,scale_fx=${SCALE_FX}]" 24 | 25 | mkdir -p "$OUTDIR" 26 | echo $MODEL > "$OUTDIR/model" 27 | 28 | date >> "$OUTDIR/progress.txt" 29 | ./.colorize ./train.py \ 30 | --log-to "runs/$OUTDIR" \ 31 | --batch-size 260 \ 32 | --learning-rate "0.0002" \ 33 | --epochs 150 \ 34 | imagepairs[files=youtube/$DATASET/*] \ 35 | "$MODEL" \ 36 | "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt" 37 | 38 | ./render_vae "$DATASET" "$MODEL" "$OUTDIR" 39 | -------------------------------------------------------------------------------- /train_pendulum_stable: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | EXP_ID="$1" 5 | N="$2" 6 | ALPHA="$3" 7 | INNER="$4" 8 | INNER_EPSILON="$5" 9 | SMOOTH_V="$6" 10 | REHU="$7" 11 | 12 | # Latent space dim: 13 | LSD=$(( 2 * N )) 14 | 15 | OUTDIR="experiments/pendulum-${EXP_ID}/${N}_${ALPHA}_${INNER}_${INNER_EPSILON}_${SMOOTH_V}_${REHU}" 16 | MODEL="stabledynamics[latent_space_dim=$LSD,a=$ALPHA,projfn=$INNER,projfn_eps=$INNER_EPSILON,smooth_v=$SMOOTH_V,hp=60,h=100,rehu=$REHU]" 17 | 18 | mkdir -p "$OUTDIR" 19 | echo $MODEL > "$OUTDIR/model" 20 | 21 | date >> "$OUTDIR/progress.txt" 22 | ./.colorize ./train.py \ 23 | --log-to "runs/$OUTDIR" \ 24 | --batch-size 2000 \ 25 | --learning-rate "0.001" \ 26 | --epochs 1000 \ 27 | --test-with "pendulum[n=$N,test]" \ 28 | "pendulum[n=$N]" \ 29 | "$MODEL" \ 30 | "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt" 31 | 32 | ./.colorize ./pendulum_error.py \ 33 | pendulum[n=$N,test] \ 34 | "$MODEL" \ 35 | "$OUTDIR/checkpoint_*.pth" \ 36 | 1000 | tee -a "$OUTDIR/eval.txt" 37 | -------------------------------------------------------------------------------- /train_pendulum_stable_knownV: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | EXP_ID="$1" 5 | N="$2" 6 | ALPHA="$3" 7 | INNER="ActualPendulumEnergy" 8 | INNER_EPSILON="$5" 9 | SMOOTH_V="$6" 10 | REHU="$7" 11 | 12 | # Latent space dim: 13 | LSD=$(( 2 * N )) 14 | 15 | OUTDIR="experiments/pendulum-${EXP_ID}/${N}_${ALPHA}_${INNER}_${INNER_EPSILON}_${SMOOTH_V}_${REHU}" 16 | MODEL="stabledynamics[latent_space_dim=$LSD,a=$ALPHA,projfn=$INNER,pendulum_n=$N,smooth_v=$SMOOTH_V,hp=60,h=100,rehu=$REHU]" 17 | 18 | mkdir -p "$OUTDIR" 19 | echo $MODEL > "$OUTDIR/model" 20 | 21 | date >> "$OUTDIR/progress.txt" 22 | ./.colorize ./train.py \ 23 | --log-to "runs/$OUTDIR" \ 24 | --batch-size 2000 \ 25 | --learning-rate "0.005" \ 26 | --epochs 1000 \ 27 | --test-with "pendulum[n=$N,test]" \ 28 | "pendulum[n=$N]" \ 29 | "$MODEL" \ 30 | "$OUTDIR/checkpoint_{epoch:0>5}.pth" | tee -a "$OUTDIR/progress.txt" 31 | 32 | ./.colorize ./pendulum_error.py \ 33 | pendulum[n=$N,test] \ 34 | "$MODEL" \ 35 | "$OUTDIR/checkpoint_*.pth" \ 36 | 1000 | tee -a "$OUTDIR/eval.txt" 37 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Learning Stable Deep Dynamics Models 2 | 3 | Companion code to "Learning Stable Deep Dynamics Models" (Manek and Kolter, 2019) 4 | 5 | ## Installation 6 | 7 | You need Python 3.6 or later, with packages listed in `requirements.txt`. Ensure that the directories `experiments/` and `runs/` are writable. 8 | 9 | ## Running 10 | 11 | All training commands produce output in the `experiments/` folder, automatically named with the hyperparameters used in training. You can track training progress using `tensorboard runs/`. 12 | 13 | ### Pendulum 14 | 15 | To train the models for an ``-link pendulum, you should first run the command: 16 | 17 | ```sh 18 | ./train_pendulum_simple 19 | ``` 20 | 21 | which will create and cache the training data and the evaluation data. (Warning: this is a cpu-heavy task that may take hours to complete.) Once that is complete, you can train multiple models in parallel with: 22 | 23 | ```sh 24 | ./train_many 25 | ``` 26 | 27 | You can concurrently train train models with different ``. 28 | 29 | Ensure that `pendulum-cache/` is writable. 30 | 31 | ### VAE 32 | 33 | To train the VAE, you need to convert videos to sequentially-numbered frames in `youtube/