# EqR: Equivaraint Representation Learning for RL


## Install 
To install the requirements, follow these steps:
```bash
# PyTorch
conda install pytorch torchvision -c pytorch

# Install requirements
pip install -r requirements.txt

```

## Usage:

* To run EqR E(2) with L_{AET}, L_{GET} and L_R on krull (change the group-type to 'GLn' to run GL (2) or 'Real' to run T(2) and num-blocks to select number of subgroups K)
```bash
python -m scripts.run --game krull --model-type 'Symmetric-MDP' --group-type 'En' --num-blocks 12  --model-group-weight 1. --symmetricMDP-loss-weight 1. --reward-loss-weight 1.
```

* To run EqR E(2) with L_{AET} and L_R on krull (change the group-type to 'GLn' to run GL (2) or 'Real' to run T(2) and num-blocks to select number of subgroups K)
```bash
python -m scripts.run --game krull --model-type 'Simple-Equivariant' --group-type 'En' --num-blocks 12  --model-group-weight 1. --reward-loss-weight 1.
```

To remove reward loss L_R set  ```--reward-loss-weight 0```

When reporting scores, we average across 10 seeds. 
