# Official Implementation of Symmetry Regularization

## Rotating Chair dataset
To generate the data:
```python generate_data.py```

To train encoder:
```python main_chair.py --model_name symmetryreg ```

## Atari Experiments
To run the representation learning using SymReg in Pong:\
```python atari_rl.py --num-epochs 20 --env-name Pong --rl-algo DQN --model-type symmetrypretrained```

To train the world model and evalute:\
```python atari_rl.py --num-epochs 100 --env-name Pong --model-type symmetrypretrained```

## Mujoco Experiments
To run the training on Inverted Pendulum for the finetuning experiments use the following: \
```python mujoco_rl.py --model-type symmetrypretrained --rl-algo PPO --env-name InvertedPendulum --rl-training-steps 1000000 --num-discretized-actions 100 --num-eval-episodes 5```

To run the training on Inverted Pendulum for the representation learning experiments use the following: \
```python mujoco_rl.py --model-type symmetrydecoupled --rl-algo PPO --env-name InvertedPendulum --rl-training-steps 1000000 --num-discretized-actions 100 --num-eval-episodes 5```

Note that the code of the other experiments are provided in ipython_notebooks folder. 