Migrate to UV

.github/workflows/ci.yaml

@@ -19,13 +19,9 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: "3.11"
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v3
       - uses: pre-commit/action@v3.0.1
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
-          pip install poetry
-          poetry install --with dev
-      - name: Test with pytest
-        run: |
-          source $(poetry env info --path)/bin/activate
-          pytest -v
+          uv sync --dev

README.md

@@ -11,7 +11,7 @@ A collection of algorithms and experiment tools for safe sim to real transfer an
 ## Requirements
 
 - Python == 3.11.6
-- `venv` or `Poetry`
+- [`uv`](https://docs.astral.sh/uv/) (recommended) or the built-in `venv`
 
 ## Installation
 
@@ -25,15 +25,70 @@ source venv/bin/activate
 pip install -e .
 ````
 
-### Using Poetry
+### Using uv
+
+Install uv if it is not already available:
+
+```bash
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+Create a project environment and install dependencies:
 
 ```bash
 git clone https://github.com/yardenas/safe-learning
 cd safe-learning
-poetry install
-poetry shell
+uv sync
+uv run python --version  # sanity check, optional
 ```
 
+### Installing `madrona_mjx` (optional, required for Madrona backend)
+
+Some benchmarks (e.g., the MJX-based pick-and-place tasks) require the custom
+[`madrona_mjx`](https://github.com/shacklettbp/madrona_mjx) fork. Build and
+install it inside the UV environment you created above:
+
+1. From the parent directory of `safe-sim2real`, clone the repository and check
+   out the tested commit:
+
+   ```bash
+   git clone https://github.com/shacklettbp/madrona_mjx.git
+   cd madrona_mjx
+   git checkout c34f3cf6d95148dba50ffeb981aea033b8a4d225
+   git submodule update --init --recursive
+   ```
+
+2. Configure and build (disable Vulkan if you do not have it available):
+
+   ```bash
+   mkdir -p build
+   cd build
+   cmake -DLOAD_VULKAN=OFF ..
+   cmake --build . -j
+   cd ..
+   ```
+
+3. While having your environment activated, install the Python bindings into your UV environment:
+
+   ```bash
+   uv pip install -e .
+   ```
+
+Refer to the upstream repository for platform-specific prerequisites (CUDA,
+Vulkan, compiler versions). Re-run `uv pip install -e .` whenever you rebuild the
+library.
+
+**Troubleshooting tips**
+
+- If you see CUDA OOMs immediately after the build, try `export
+  MADRONA_DISABLE_CUDA_HEAP_SIZE=1` before launching training.
+- Populate the kernel caches to avoid recompilation on every run:
+
+  ```bash
+  export MADRONA_MWGPU_KERNEL_CACHE=/path/to/cache/mwgpu
+  export MADRONA_BVH_KERNEL_CACHE=/path/to/cache/bvh
+  ```
+
 ## Usage
 
 Our code uses [Hydra](https://hydra.cc/) to configure experiments. Each experiment is defined as a `yaml` file in `ss2r/configs/experiments`. For example, to train a Unitree Go1 policy with a constraint on joint limit:
@@ -65,4 +120,3 @@ If you find our repository useful in your work, please consider citing:
 * **Project Webpage**: 
 * **Paper**:
 * **Contact**: 
-

euler_guide.md

@@ -5,16 +5,16 @@
 
 ## Installing Madrona-MJX on Euler
 1. Make sure that the steps above have been completed on the current terminal.
-2. Go to the parent directory of the safe-sim2real repository: `cd ..`
-3. Activate the poetry environment: `eval $(poetry env activate)`
-4. Move out of `safe-sim2real` to your home folder.
+2. Go to the parent directory of the safe-learning repository: `cd ..`
+3. Activate the environment: `source .venv/bin/activate`
+4. Move out of `safe-learning` to your home folder.
 5. `git clone https://github.com/shacklettbp/madrona_mjx.git && git checkout c34f3cf6d95148dba50ffeb981aea033b8a4d225`
 6. `cd madrona_mjx`
 7. `git submodule update --init --recursive`
 8. Load cmake `module load cmake/3.27.7`
 9. Point to cuda CMAKE witout loading cuda `export CMAKE_PREFIX_PATH="/cluster/software/stacks/2024-06/spack/opt/spack/linux-ubuntu22.04-x86_64_v3/gcc-12.2.0/cuda-12.1.1-5znnrjb5x5xr26nojxp3yhh6v77il7ie/:${CMAKE_PREFIX_PATH}"`
 10. `mkdir build && cd build && cmake -DLOAD_VULKAN=OFF .. && cd ..`
 11. Compile on a compute node `srun --ntasks=1 --cpus-per-task=20 --gpus=rtx_4090:1 --time=0:30:00 --mem-per-cpu=10240 make -j`
-12. Change back to the root of the safe-sim2real repository: `cd .. && cd safe-sim2real`
+12. Change back to the root of the safe-learning repository: `cd .. && cd safe-learning`
 13. `poetry install --with madrona-mjx`
 14. Test `srun --ntasks=1 --cpus-per-task=20 --gpus=rtx_4090:1 --time=0:30:00 --mem-per-cpu=10240 python train_brax.py +experiment=cartpole_vision`