Argus

A vision-based pose estimator for the in-hand manipulation hardware setup in the AMBER Lab.

Installation

1. Clone this repository and cd into the repo root.
2. Manage the cudatoolkit version and python dependencies using conda by running the following

conda env create --name <your_env_name> --file=environment.yml

# pick the install command that is most relevant to you!
pip install -e .  # if you don't need unity or dev dependencies
pip install -e .[unity]  # installs ml-agents, forces numpy and scipy to be old
pip install -e .[dev]  # use [dev] for dev tooling and testing + unity - else, don't need it

3. If you installed the dev dependencies, activate the pre-commit checks with

   pre-commit install
   

   Now, when you commit files, the checks will be run first.

4. We use wandb to log the training process. To use it, you need to set up an account and run

   wandb login
   

   When prompted, enter your API key.

Usage

Below are quickstart instructions. Most of the configs are managed via tyro, so you can easily check what other configurable fields are exposed using the -h option for those commands if you want something different.

Data Generation

The data generation pipeline has 2 steps.

1. Simulated mjpc data is generated by code in the cube_rotation repository. This data is then saved out into a json file. Suppose its path is /path/to/mjpc/data.json
2. The data_generation.py script reads in this data and renders images of the scene to generate a new dataset used to train the estimator. It also needs a Unity executable built from the environment with the LEAP hand in it.

This second step requires Unity to be installed. Here are concrete steps.

• Install some version of Unity from 2023 and open a project from the LeapProject directory.
• The scene that pops up should have the LEAP hand and a cube in it. Build out an executable by navigating to File > Build Settings > Build. Suppose you save it to /path/to/unity/exe.extension.

Once you have these two directories, you can run the data generation pipeline with default settings by running the followin command from the repo root:

python argus/data_generation.py --env-exe-path /path/to/unity/exe.extension --mjpc-data-path /path/to/mjpc/data.json

This will save out an hdf5 file used to train the estimator to the directory /REPO_ROOT/outputs/data/cube_unity_data.hdf5.

Training and Validation

To train the estimator with default settings run (for your choice of small, medium, or large dataset)

python argus/train.py --dataset-config.dataset-path outputs/data/cube_unity_data_<small,medium,large>

This will train an estimator and save out a model to /REPO_ROOT/outputs/models/<wandb_id>.pth. Multi-GPU training is also implemented via torch.DistributedDataParallel. To use this, simply run

python argus/train.py --dataset-config.dataset-path outputs/data/cube_unity_data_<small,medium,large> --multigpu

To run some visual evaluations as sanity checks, run

python argus/validate.py --model-path /home/albert/research/argus/outputs/models/<checkpoint>.pth

This will save out images showing qualitative performance on the validation set to the directory /REPO_ROOT/outputs/validation_visuals/<checkpoint>.

To run visual evaluations on real-world data, run

python argus/validate_real.py --model-path /home/albert/research/argus/outputs/models/<checkpoint>.pth --dataset-config.dataset-path outputs/data/cube_unity_data_<small,medium,large>

Testing

To run tests, run pytest in the repo root.