Refractor - Rodent Imitation/Online RL - Distributed Training - Visio…

…n Encoder - Intention Network Arch (#2) * ignore data folder * conda env creation script for TF cuda support * added requirement.txt * modify the environment before the warpper. * injected rodent class * ignore checkpoint file * modify for rodent * added my note * render expet * modify to work with rodent env * ignore logging * initialized * modified * modified * ignore the training * update notes * added flag todo * chagne hyperparam * rendering * added wandb logging for learner + configuration file support * attempt to do Rodent SingleClipTreacking * added config file & logging support * ready for merge * added preprocessing step * change to vast * move to vast * distributed ray support * ready for merge * Support rodent imitation now * scheduling strategies * distributed training support from flybody code * update my notes * update progress on genearlist training * update my notes * distributed training * ready for merge * black formatted * support intention networks with wrapper on the observation space * generalist training ready * update my notes * make network smaller * update notes * kickstarting training * added back fly and humanoid imitation, load mocap data set, and added intention network * finish kickstarting and general training * added notes * support multiple replay servers for both multi-task/single task training. * bug * support wandb rendering for rollout during training * humanoid tracking working * clean up imitation learning pipeline * curriculum learning dev in progress * fix api between online & imitation * tuning for imitaion learning * added average report of the evaluator * updates my note * update * clean up repo * debug decoder arch * added checkpointer for decoder * added better logging terms * support custom reward scaling term * add option to remove skin * clean up unused imitation visnet * undo remove skin -> modify xml file * implement encoder snapshotter * more robust rendering * formats * refractor + fix bugs in eval * Simplify API * added iterative lifting * modify default action std * rename & fix NN shape issue * refractor intention network API - more general * adapt logger API * API Changes * add reward composition rendering * expose reward composition weight to config * added upright reward & contact terminations * support mixed experience replay server * reward rendering util * implemented visual intention network * explicit froze_decoder option when loading checkpoints * remove redundant comments * implement reward term logging for online tasks + reward termination logics * added config file * renamed folder + package name * rename package reference name * delete unused rodent file * added config option for discount * add hierarchical config file for better readability * added reward termination/more reward terms for bowl escape * implemented multiple layer of encoder * change reward condition * updates readme --------- Co-authored-by: KevinBian107 <[email protected]>
talmolab · Sep 25, 2024 · 71dfc10 · 71dfc10
1 parent 2268560
commit 71dfc10
Show file tree

Hide file tree

Showing 279 changed files with 510,544 additions and 57,376 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,7 +1,17 @@
 __pycache__/
-
 *.ipynb
 .ipynb_checkpoints
-
+clips/
+cps/
+demos/
+*.out
 *.egg-info
-
+nohup.out
+wandb/
+policy/
+outputs/
+training/
+outputs/
+*.mp4
+policy/
+data/
diff --git a/README.md b/README.md
@@ -1,132 +1,3 @@
-# `flybody`: fruit fly body model for MuJoCo physics
+# `VNL-Ray`: Distributed RL Training with MuJoCo Model
 
-[![preprint](https://img.shields.io/badge/preprint-bioRxiv-B31B1B)][paper]
-
-<img src="fly-white.png" width="65%">
-
-![pytest workflow](https://github.com/TuragaLab/flybody/actions/workflows/pytest.yml/badge.svg)
-![python versions](https://github.com/TuragaLab/flybody/actions/workflows/pyversions.yml/badge.svg)
-![lint](https://github.com/TuragaLab/flybody/actions/workflows/lint.yml/badge.svg)
-![tf](https://github.com/TuragaLab/flybody/actions/workflows/tf-test.yml/badge.svg)
-
-`flybody` is an anatomically-detailed body model of the fruit fly _Drosophila melanogaster_ for [MuJoCo][mujoco] physics simulator and reinforcement learning applications. 
-
-The fly model was developed in a collaborative effort by Google DeepMind and HHMI Janelia Research Campus. 
-
-We envision our model as a platform for fruit fly biophysics simulations and for modeling neural control of sensorimotor behavior in an embodied context; see our accompanying [publication][paper].
-
-## Getting Started
-
-The fruit fly body model lives in [this directory][fly-home]. To visualize it, you can drag-and-drop `fruitfly.xml` or `floor.xml` to MuJoCo's `simulate` viewer.
-
-Beginning interacting with the model via Python is as simple as:
-
-```python
-from dm_control import mujoco
-
-physics = mujoco.Physics.from_xml_path('flybody/fruitfly/assets/fruitfly.xml')  # Load model.
-physics.step()  # Step simulation.
-```
-
-The quickest way to get started with `flybody` is to take a look at a [tutorial notebook][tutorial] or [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)][tutorial-colab].
-
-Also, [this notebook][envs] shows examples of the flight, walking, and vision-guided flight RL task environments. 
-
-To train the fly, try the [distributed RL training script][ray-script], which uses [Ray][ray] to parallelize the [DMPO][dmpo] agent training.
-
-[fly-home]: https://github.com/TuragaLab/flybody/tree/main/flybody/fruitfly/assets
-[mujoco]: https://github.com/google-deepmind/mujoco
-[tutorial]: https://github.com/TuragaLab/flybody/blob/main/docs/getting-started.ipynb
-[tutorial-colab]: https://colab.research.google.com/github/TuragaLab/flybody/blob/main/docs/getting-started.ipynb
-[dmpo]: https://github.com/google-deepmind/acme/tree/master/acme/agents/tf/dmpo
-[envs]: https://github.com/TuragaLab/flybody/blob/main/docs/fly-env-examples.ipynb
-[ray-script]: https://github.com/TuragaLab/flybody/blob/main/flybody/train_dmpo_ray.py
-[paper]: https://www.biorxiv.org/content/10.1101/2024.03.11.584515
-[ray]: https://github.com/ray-project/ray
-[tf]: https://github.com/tensorflow/tensorflow
-[acme]: https://github.com/google-deepmind/acme
-[mujoco-rendering]: https://github.com/google-deepmind/dm_control/tree/main?tab=readme-ov-file#rendering
-
-## Installation
-
-Follow these steps to install `flybody`:
-
-### Option 1: Installation from cloned local repo
-
-1. Clone this repo and create a new conda environment:
-   ```bash
-   git clone https://github.com/TuragaLab/flybody.git
-   cd flybody
-   conda env create -f flybody.yml
-   conda activate flybody
-   ```
-   `flybody` can be installed in one of the three modes described next. Also, for installation in editable (developer) mode, use the commands as shown. For installation in regular, not editable, mode, drop the `-e` flag.
-
-2. **Core installation**: minimal installation for experimenting with the
-   fly model in MuJoCo or prototyping task environments. ML dependencies such as [Tensorflow][tf] and [Acme][acme] are not included and policy rollouts and training are not automatically supported.
-   ```bash
-   pip install -e .
-   ```
-
-3. **ML extension (optional)**: same as core installation, plus ML dependencies (Tensorflow, Acme) to allow running
-   policy networks, e.g. for inference or for training using third-party agents not included in this library.
-   ```bash
-   pip install -e .[tf]
-   ```
-
-4. **Ray training extension (optional)**: same as core installation and ML extension, plus [Ray][ray] to also enable
-   distributed policy training in the fly task environments.
-   ```bash
-   pip install -e .[ray]
-   ```
-
-### Option 2: Installation from remote repo
-1. Create a new conda environment:
-   ```bash
-   conda create --name flybody python=3.10 pip ipython cudatoolkit cudnn=8.2.1=cuda11.3_0
-   conda activate flybody
-   ```
-   Proceed with installation in one of the three modes (described above):
-2. **Core installation**:
-   ```bash
-   pip install git+https://github.com/TuragaLab/flybody.git
-   ```
-3. **ML extension (optional)**:
-   ```bash
-   pip install "flybody[tf] @ git+https://github.com/TuragaLab/flybody.git"
-   ```
-5. **Ray training extension (optional)**:
-   ```bash
-   pip install "flybody[ray] @ git+https://github.com/TuragaLab/flybody.git"
-   ```
-
-### Additional configuring
-
-1. You may need to set [MuJoCo rendering][mujoco-rendering] environment varibles, e.g.:
-   ```bash
-   export MUJOCO_GL=egl
-   export MUJOCO_EGL_DEVICE_ID=0
-   ```
-2. Also, for the ML and Ray extensions, `LD_LIBRARY_PATH` may require an update, e.g.:
-   ```bash
-   export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/your/path/to/miniconda3/envs/flybody/lib
-   ```
-
-3. You may want to run `pytest` to test the main components of the `flybody` installation.
-
-## Citing `flybody`
-See our accompanying [publication][paper]. Thank you for your interest in our fly model:)
-```bibtex
-@article{flybody,
-  title = {Whole-body simulation of realistic fruit fly locomotion with
-           deep reinforcement learning},
-  author = {Roman Vaxenburg and Igor Siwanowicz and Josh Merel and Alice A Robie and
-            Carmen Morrow and Guido Novati and Zinovia Stefanidi and Gwyneth M Card and
-            Michael B Reiser and Matthew M Botvinick and Kristin M Branson and
-            Yuval Tassa and Srinivas C Turaga},
-  journal = {bioRxiv},
-  doi = {https://doi.org/10.1101/2024.03.11.584515},
-  url = {https://www.biorxiv.org/content/10.1101/2024.03.11.584515},
-  year = {2024},
-}
-```
+`VNL-Ray`, built on [Turaga Lab's Flybody](https://github.com/TuragaLab/flybody), enables distributed training of diverse biomechanical models using Ray. We incorporated a stochastic intention network to model policy intentions.
diff --git a/docs/fly-env-examples.ipynb b/docs/fly-env-examples.ipynb
diff --git a/docs/fly-on-ball-minimal.ipynb b/docs/fly-on-ball-minimal.ipynb
diff --git a/docs/getting-started.ipynb b/docs/getting-started.ipynb
diff --git a/docs/imgs/image.png b/docs/imgs/image.png