This repository contains the code to reproduce experiments of the ICLR 2024 paper Identifying Policy Gradient Subspaces by Jan Schneider, Pierre Schumacher, Simon Guist, Le Chen, Daniel Häufle, Bernhard Schölkopf, and Dieter Büchler.

Download MuJoCo version 1.50. and extract the zip archive to ~/.mujoco. It has to be version 1.50. since the code builds upon stable-baselines3==1.8.0, which depends on an old gym version and thus requires an old version of mujoco-py and consequently MuJoCo.

Make sure that the system packages libosmesa6-dev and patchelf are installed (required to build mujoco-py).

sudo apt-get install libosmesa6-dev
sudo apt-get install patchelf

Install the pg_subspaces package.

cd iclr_2024_pg_subspaces
pip install .

To train an agent, use the following command (the working directory should be iclr_2024_pg_subspaces).

python -m pg_subspaces.scripts.train [arg=value ...]

Important arguments include

To view the full list of arguments and their default values, have a look at pg_subspaces/scripts/conf/train.yaml.

The training script will create a log directory under logs/training/{env}/{date}/{time}. Use tensorboard to view the learning curve (scalar: eval/mean_reward).

tensorboard --logdir /path/to/trainlogs

To run the gradient subspace fraction analysis, use the following command

python -m pg_subspaces.scripts.analyze train_logs=/path/to/trainlogs [arg=value ...]

where /path/to/trainlogs should be replaced by the path to the log directory created by the train command from section Training an agent.

Important arguments include

The analysis results will be added to the tensorboard logs of the training run (scalars: high_curvature_subspace_analysis/default/gradient_subspace_fraction_*/*).

To compute the subspace overlap metric, run the following command (need to run after the gradient subspace fraction analysis)

python -m pg_subspaces.scripts.compute_subspace_overlaps train_logs=/path/to/trainlogs [arg=value ...]

Important arguments include

The analysis results will be added to the tensorboard logs of the training run (scalars: high_curvature_subspace_analysis/default/overlaps_*/*).

@inproceedings{schneider2024identifying,
    title={Identifying Policy Gradient Subspaces},
    author={Schneider, Jan and Schumacher, Pierre and Guist, Simon and Chen, Le and H{\"a}ufle, Daniel and Sch{\"o}lkopf, Bernhard and B{\"u}chler, Dieter},
    booktitle={The Twelfth International Conference on Learning Representations},
    year={2024}
}