# Official Implementation of NGT with CUDA Graphs

Official PyTorch implementation of the **Noise-Guided Transport (NGT)** algorithm
along with baselines.
It uses PyTorch's
[CUDA Graphs](https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/).

## Prerequisites

- **GPU**: NVIDIA GPU with CUDA support
- **CUDA Version**: CUDA 12.5 or higher (for GPU support)
- **NVIDIA Drivers**: Compatible with your CUDA version (e.g. 555)
- (Optional) For Docker:
  - **Docker**: version 27.3.1
  - **NVIDIA Container Toolkit** for GPU support
- **Expert Demonstrations**: Download from <dataset_url> (**redacted here for anonymity**)

## Installation

The project is set up with **Poetry**,
with the option to containerize with **Docker**
instead of a local Poetry install.

### Clone the Repository (**redacted here for anonymity**)

```bash
git clone <repo_url> # redacted
cd ngt-pytorch
```

### Setup wandb

Add `your_wandb_api_key` as the environment variable `WANDB_API_KEY`
in the `.env` file (the file will be created if it does not exist):

```bash
echo "WANDB_API_KEY=your_wandb_api_key" >> .env
```

### Install Poetry

Install Poetry directly with the official installer:
```bash
curl -sSL https://install.python-poetry.org | python3 -
```

Also install **Poethepoet**.
Poethepoet enables the creation of macros,
which are defined in `pyproject.toml`.
The macros defined there facilitate both pure-Poetry and Docker usages.
The following command installs it globally, i.e. available across all projects:

```bash
poetry self add poethepoet
```

Also install **poetry-dotenv-plugin**.
This plugin handles `.env` files seamlessly,
no need to export nor specify any env-file.
Just put every environment variable you need in the `.env` file
at the root of the project.
In particular, having you **wandb API key** in the file as described above
is enough for the login to be successful,
whether you choose to use Docker or not.
The following command installs it globally, i.e. available across all projects:

```bash
poetry self add poetry-dotenv-plugin
```

Make sure that the `poetry` and `poe` bash command are on the PATH.
Run this directly or add this to your `.bashrc`/`.zshrc`:

```bash
export PATH="$HOME/.local/bin:$PATH"
```

Double-check that they are available with:

```bash
poetry --version
poe --version
```

### Install project dependencies

#### Option A: Poetry

```bash
poetry install
```

If that command hangs, use the verbose version: `poetry install -vvv`.
If you see that it hangs at "Using keyring backend 'SecretService Keyring'",
you can configure Poetry to not use the keyring for managing credentials:

```bash
export PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring
```

#### Option B: Docker

(Optional) Build and push the Docker image to Dockerhub:

```bash
poe dbp
```

N.B.: optional because the image is already on Dockerhub.

## Usage

You can either **train** or **evaluate**.
What follows describes how to run either task
on either pure-Poetry or Docker.

N.B. 1: `-O` (capital "o") turns the value of the constant `__debug__` to `False`,
making Python run in **optimized mode**.
In this mode, Python ignores the `assert` statements. **Use with caution when prototyping**.
Plus, in this mode, the `beartype` decorator is turned off, hence reducing its tiny overhead time.
Conversely, running the command without the `-O` option activates `beartype` real-time type-checking functionalities.

N.B. 2: The **evaluate** task evaluates a model
by retrieving its best performing checkpoint stored on Weights & Biases servers.
What should be given as value to the `--load_ckpt` argument
is the "Run Path" on the Weighs & Biases overview page of the run to evaluate.
The parameters of the best model are downloaded in a temporary directory.

### Option A: Poetry

#### Train

```bash
poetry run python -O main.py train --base_cfg="tasks/defaults/base.yml" --override_cfg="tasks/defaults/overrides/your_override.yml" --env_id="Hopper-v4" --seed=0 --num_demos=4 --subsampling_rate=20 --expert_path="./experts" --description="teeny description"
```

#### Evaluate

```bash
poetry run python -O main.py evaluate --base_cfg="tasks/defaults/base.yml" --override_cfg="tasks/defaults/overrides/your_override.yml" --env_id="Hopper-v4" --seed=0 --load_ckpt="wandb_run_path" --description="teeny description"
```

### Option B: Docker

Run the macro:

```bash
poe docker
```

You are now inside the container.

#### Train

```bash
python -O main.py train --base_cfg="tasks/defaults/base.yml" --override_cfg="tasks/defaults/overrides/your_override.yml" --env_id="Hopper-v4" --seed=0 --num_demos=4 --subsampling_rate=20 --expert_path="./experts" --description="teeny description"
```

#### Evaluate

```bash
python -O main.py evaluate --base_cfg="tasks/defaults/base.yml" --override_cfg="tasks/defaults/overrides/your_override.yml" --env_id="Hopper-v4" --seed=0 --load_ckpt="wandb_run_path" --description="teeny description"
```
## Extras

### Expert Performance

To print the performance (mean and standard deviation)
of the experts across all environments and demonstrations stored in `.h5` files, use:

```bash
poetry run python -O inspect_experts.py --expert_path="./experts"
```

### Download progress files from wandb

To download the `progress.json` and `progress.csv` files that have
tracked the performance of agents during training, the following will
download all those files for a given wandb _group_:

```bash
poetry run python -O retrieve_wandb.py --wandb_id="wand_id" --wandb_project="wand_project" --group_name="group_name" --download_dir="download_dir"
```

## Advanced: Spawn Arrays of Jobs (**removed `spawner.py` because it contains cluster names and partition names that would break anonymity**)

The `spawner.py` enables the creation (and launch) of an array of experiments
on a **Slurm-administrated cluster**,
or locally in a new `tmux` _session_, with _one_ experiment running _per window_ in the session.

Here is how it can be used:

```bash
poetry run python spawner.py --base_cfg="tasks/defaults/base.yml" --override_cfg="tasks/defaults/overrides/your_override.yml" --env_bundle="huma" --deployment="slurm" --num_seeds=3 --num_demos="[1,4]" --subsampling_rate="[1,20]" --expert_path="./experts" --deploy_now  --description="teeny description"
```

To run the same command on a list of bundles, use this:

```bash
echo "walker cheetah ant huma" | tr ' ' '\n' | parallel 'poetry run python spawner.py --base_cfg="tasks/defaults/base.yml" --override_cfg="tasks/defaults/overrides/your_override.yml" --env_bundle={} --deployment="slurm" --num_seeds=4 --num_demos="[1, 4]" --subsampling_rate="[1, 20]" --expert_path="./experts" --deploy_now --debug --description="teeny description"'
```

This command uses [GNU `parallel`](https://www.gnu.org/software/parallel/) which by default
determines the number of jobs based on the length of the input list or based on available CPU resources.
To force parallel to use a certain number of cores, use the `-j` option, e.g. `parallel -j 4` for 4 cores.
The software is usually available by default on HPC clusters.

To create the scripts _without_ deploying them immediately, use `--nodeploy_now` instead of `--deploy_now`.
This logic applies to all the boolean options since we are using [google/python-fire](https://github.com/google/python-fire).