Readme updates

README.md

@@ -1,80 +1,90 @@
 # Retinal RL
 
-## Setting up the development environment
+A deep learning framework for vision research using deep reinforcement learning.
 
-We provide a singularity / apptainer container which should always be up to date and allow to run the code immediately. You do not need to build it yourself (but can, of course), you can just pull it!
+## Apptainer Environment
 
-### Install Singularity / Apptainer
+Retinal-Rl is designed to run in a containerized environment using [Apptainer](https://apptainer.org/docs/user/latest/).
 
-First you need to install [apptainer](https://github.com/apptainer/apptainer/) (or singularity) in order to run code.
+### Installation
 
-### Get the container
+1. [Install](https://apptainer.org/docs/admin/main/installation.html) Apptainer to run the containerized environment.
 
-Once you have apptainer installed, you can simply pull the container
+2. Get the container:
 
+- Either pull the pre-built container:
 ```bash
 apptainer pull retinal-rl.sif oras://ghcr.io/berenslab/retinal-rl:singularity-image
 ```
-
-or try to build it on your own (no advantages of doing that, except you want to change some dependency in the .def file):
-
+- or build from source:
 ```bash
 apptainer build retinal-rl.sif resources/retinal-rl.def
 ```
 
-### Prepare config directory for experiments
-
-The repository comes with some example configuration files, which you find under 'resources/config_templates'. For running experiments however, they need to be in 'config'.
-You can either copy them there by hand or run the following script from the top-level directory:
+### Running Experiments
 
+The `scan` command prints info about the proposed neural network architecture:
 ```bash
-bash tests/ci/copy_configs.sh
+apptainer exec retinal-rl.sif python main.py +experiment="{experiment}" command=scan
+```
+The experiment must always be specified with the `+experiment` flag. To train a
+model, use the `train` command:
+```bash
+apptainer exec retinal-rl.sif python main.py +experiment="{experiment}" command=train
 ```
 
-### Test basic functionality
+`apptainer` commands can typically be replaced with `singularity` if the latter is rather used.
 
-Now you are basically ready to run experiments!
-To test that everything is working fine, you can run:
+## Hydra Configuration
 
-```bash
-bash tests/ci/scan_configs.sh
-```
+The project uses [Hydra](https://hydra.cc/) for configuration management.
 
-The script loops over all experiments defined in config/experiment and runs a "scan" on them.
-If instead you want to run a single experiment file, run:
+### Directory Structure
+
+The structure of the `./config/` directory is as follows:
 
-```bash
-apptainer exec retinal-rl.sif python main.py +experiment="$experiment" command=scan system.device=cpu
+```
+base/config.yaml     # General and system configurations
+user/
+├── brain/           # Neural network architectures
+├── dataset/         # Dataset configurations
+├── optimizer/       # Training optimizers
+└── experiment/      # Experiment configurations
 ```
 
-## Running retinal RL simulations [DEPRECATED]
+### Default Configuration
 
-There are three main scripts for working with `retinal-rl`:
+Template configs are available under `./resources/config_templates/user/...`, which also provide documentation of the configuration variables themselves. Consult the hydra documentation for more information on [configuring your project](https://hydra.cc/docs/intro/).
 
-- `train.py`: Train a model.
-- `analyze.py`: Generate some analyses.
-- `enjoy.py`: Watch a real time simulation of a trained agent.
+### Configuration Management
 
-Each script can be run by python in `python -m {script}`, where {script} is the name of the desired script (without the `.py` extension), followed by a number of arguments. Note that `train.py` must always be run first to create the necessary files and folders, and once run will permanently set most (all?) of the arguments of the simulation, and will ignore changes to these arguments if training is resumed.
+1. Configuration templates may be copied to the user directory by running:
+```bash
+bash tests/ci/copy_configs.sh
+```
 
-Certain arguments must always be provided, regardless of script, namely:
+2. Template and custom configurations can be sanity-checked with:
+```bash
+bash tests/ci/scan_configs.sh
+```
+which runs the `scan` command for all experiments.
 
-- `--env`: Specifies the desired map. This will always have the form `retinal_{scenario}`, where scenario is the shared name of one of the `.wad`/`.cfg` file pairs in the `scenarios` directory.
-- `--algo`: The training algorithm; for now this should always be `APPO`.
-- `--experiment`: The directory under the `train_dir` directory where simulation results are saved.
+## Weights & Biases Integration
 
-The following argument should always be set when training for the first time:
+Retinal-RL supports logging to [Weights & Biases](https://wandb.ai/site) for experiment tracking.
 
-- `--encoder_custom`: The options are `simple`, which is a small, hard-coded network that still tends to perform well, and `lindsey`, which has a number of tuneable hyperparameters.
+### Basic Configuration
 
-For specifying the form of the `lindsey` network, the key arguments are:
+By default plots and analyses are saved locally. To enable Weights & Biases logging, add the `logging.use_wandb: True` flag to the command line:
+```bash
+apptainer exec retinal-rl.sif python main.py +experiment="{experiment}" logging.use_wandb=True command=train
+```
 
-- `--global_channels`: The number of channels in each CNN layers, except for the bottleneck layer.
-- `--retinal_bottleneck`: Number of channels in the retinal bottleneck.
-- `--vvs_depth`: Number of CNN layers in the ventral stream network.
-- `--kernel_size`: Size of the kernels.
+### Parameter Sweeps
 
-Finally, when training a model there are a number of additional parameters for controlling the reinforcement learning brain, and adjusting simulation parameters. The key ones to worry about are
+Wandb [sweeps](https://docs.wandb.ai/guides/sweeps) can be added to `user/sweeps/{sweep}.yaml` and launched from the command line:
+```bash
+apptainer exec retinal-rl.sif python main.py +experiment="{experiment}" +sweep="{sweep}" command=sweep
+```
 
-- `--hidden_size`: The size of the hidden/latent state used to represent the RL problem.
-- `--num_workers`: This is the number of simulation threads to run. This shouldn't be more than the number of cores on the CPU, and can be less if the simulation is GPU bottlenecked.
+Typically the only command line arguments that need a `+` prefix will be `+experiment` and `+sweep`. Also note that `.yaml` extensions are dropped at the command line.