Skip to content

Commit f4faf19

Browse files
author
Bourmindge Liey
committed
Merge remote-tracking branch 'origin/main' into notebooks-mpc
2 parents 0a58312 + 113e229 commit f4faf19

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+3912
-1255
lines changed

‎.devcontainer/Dockerfile

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
FROM nvidia/cuda:12.6.3-devel-ubuntu22.04
2+
3+
# Prevent interactive prompts during package installation
4+
ENV DEBIAN_FRONTEND=noninteractive
5+
6+
# Install system dependencies
7+
RUN apt-get update && apt-get install -y --no-install-recommends \
8+
python3.11 \
9+
python3-pip \
10+
python3.11-venv \
11+
build-essential \
12+
git \
13+
curl \
14+
ssh \
15+
mesa-utils \
16+
&& rm -rf /var/lib/apt/lists/* && apt-get clean
17+
18+
# Create symbolic links for python
19+
RUN ln -sf /usr/bin/python3.11 /usr/bin/python \
20+
&& ln -sf /usr/bin/python3.11 /usr/bin/python3
21+
22+
# Create non-root user
23+
ARG USERNAME=vscode
24+
ARG USER_UID=1000
25+
ARG USER_GID=$USER_UID
26+
27+
RUN groupadd --gid $USER_GID $USERNAME \
28+
&& useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \
29+
&& apt-get update \
30+
&& apt-get install -y sudo \
31+
&& echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
32+
&& chmod 0440 /etc/sudoers.d/$USERNAME
33+
34+
# Set environment variables for CUDA
35+
ENV PATH="/usr/local/cuda-12.0/bin:${PATH}"
36+
ENV LD_LIBRARY_PATH="/usr/local/cuda-12.0/lib64:${LD_LIBRARY_PATH}"
37+
38+
ENV NVIDIA_VISIBLE_DEVICES ${NVIDIA_VISIBLE_DEVICES:-all}
39+
ENV NVIDIA_DRIVER_CAPABILITIES ${NVIDIA_DRIVER_CAPABILITIES:+$NVIDIA_DRIVER_CAPABILITIES,}graphics
40+
41+
# Create and set proper permissions for workspace directory
42+
RUN mkdir -p /workspace && chown -R $USERNAME:$USERNAME /workspace && sudo chmod -R 777 /workspace
43+
44+
# Create virtual environment with proper permissions
45+
ENV VIRTUAL_ENV=/home/vscode/venv
46+
RUN python -m venv $VIRTUAL_ENV && \
47+
chown -R $USERNAME:$USERNAME $VIRTUAL_ENV
48+
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
49+
50+
# Switch to non-root user
51+
USER $USERNAME
52+
53+
# Install development tools
54+
RUN pip install --no-cache-dir ruff
55+
56+
# Create a temporary directory for installation with proper permissions
57+
WORKDIR /tmp/install
58+
59+
# Copy pyproject.toml with correct ownership
60+
COPY --chown=$USERNAME:$USERNAME pyproject.toml .
61+
62+
# Install dependencies as non-root user
63+
RUN pip install --no-cache-dir -e ".[test,gpu]"
64+
65+
# Set final working directory and ensure proper permissions
66+
WORKDIR /workspace

‎.devcontainer/devcontainer.linux.json

+58
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
{
2+
"name": "Crazyflow Development",
3+
"build": {
4+
"dockerfile": "Dockerfile",
5+
"context": ".."
6+
},
7+
"customizations": {
8+
"vscode": {
9+
"extensions": [
10+
"ms-python.python",
11+
"ms-python.pylance",
12+
"charliermarsh.ruff",
13+
"ms-vscode.test-adapter-converter",
14+
"ms-python.pytest-adapter",
15+
"nvidia.nsight-vscode-edition",
16+
"ms-azuretools.vscode-docker",
17+
"ms-toolsai.jupyter-renderers",
18+
"ms-toolsai.jupyter",
19+
"ms-toolsai.jupyter-keymap"
20+
],
21+
"settings": {
22+
"python.defaultInterpreterPath": "/home/vscode/venv/bin/python",
23+
"python.testing.pytestEnabled": true, // test framework
24+
"python.testing.pytestArgs": [
25+
"-m",
26+
"not render"
27+
],
28+
"python.testing.unittestEnabled": false, // test framework
29+
"python.testing.nosetestsEnabled": false, // test framework
30+
"editor.formatOnSave": true,
31+
"editor.codeActionsOnSave": {
32+
"source.fixAll.ruff": true,
33+
"source.organizeImports.ruff": true
34+
},
35+
"[python]": {
36+
"editor.defaultFormatter": "charliermarsh.ruff",
37+
"editor.formatOnSave": true,
38+
"editor.rulers": [
39+
100
40+
]
41+
}
42+
}
43+
}
44+
},
45+
"containerEnv": {
46+
"DISPLAY": "${localEnv:DISPLAY}",
47+
"XAUTHORITY": "${localEnv:XAUTHORITY}"
48+
},
49+
"runArgs": [
50+
"--interactive",
51+
"--net=host",
52+
"--ipc=host",
53+
// "--gpus=all", //use only with GPU
54+
// "--runtime=nvidia" //use only with GPU
55+
],
56+
"postCreateCommand": "python -m pip install -e '.[test,gpu]'",
57+
"remoteUser": "vscode"
58+
}

‎.devcontainer/devcontainer.wsl2.json

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
{
2+
"name": "Crazyflow Development",
3+
"build": {
4+
"dockerfile": "Dockerfile",
5+
"context": ".."
6+
},
7+
"customizations": {
8+
"vscode": {
9+
"extensions": [
10+
"ms-python.python",
11+
"ms-python.pylance",
12+
"charliermarsh.ruff",
13+
"ms-vscode.test-adapter-converter",
14+
"ms-python.pytest-adapter",
15+
"nvidia.nsight-vscode-edition",
16+
"ms-azuretools.vscode-docker",
17+
"ms-toolsai.jupyter-renderers",
18+
"ms-toolsai.jupyter",
19+
"ms-toolsai.jupyter-keymap"
20+
],
21+
"settings": {
22+
"python.defaultInterpreterPath": "/home/vscode/venv/bin/python",
23+
"python.testing.pytestEnabled": true, // test framework
24+
"python.testing.pytestArgs": [
25+
"-m",
26+
"not render"
27+
],
28+
"python.testing.unittestEnabled": false, // test framework
29+
"python.testing.nosetestsEnabled": false, // test framework
30+
"editor.formatOnSave": true,
31+
"editor.codeActionsOnSave": {
32+
"source.fixAll.ruff": true,
33+
"source.organizeImports.ruff": true
34+
},
35+
"[python]": {
36+
"editor.defaultFormatter": "charliermarsh.ruff",
37+
"editor.formatOnSave": true,
38+
"editor.rulers": [
39+
100
40+
]
41+
}
42+
}
43+
}
44+
},
45+
// mounts required for WSL2 X11
46+
"mounts": [
47+
"source=/tmp/.X11-unix,target=/tmp/.X11-unix,type=bind",
48+
"source=/mnt/wslg,target=/mnt/wslg,type=bind",
49+
"source=/usr/lib/wsl,target=/usr/lib/wsl,type=bind"
50+
],"containerEnv": {
51+
"DISPLAY": "${localEnv:DISPLAY}",
52+
"XAUTHORITY": "${localEnv:XAUTHORITY}",
53+
"WAYLAND_DISPLAY": "${localEnv:WAYLAND_DISPLAY}", // WSL2
54+
"XDG_RUNTIME_DIR": "${localEnv:XDG_RUNTIME_DIR}", // WSL2
55+
"PULSE_SERVER": "${localEnv:PULSE_SERVER}" // WSL2
56+
},
57+
"runArgs": [
58+
"--interactive",
59+
"--net=host",
60+
"--ipc=host"
61+
// "--gpus=all", //use only with GPU
62+
// "--runtime=nvidia" //use only with GPU
63+
],
64+
"postCreateCommand": "python -m pip install -e '.[test,gpu]'",
65+
"remoteUser": "vscode"
66+
}

‎.gitignore

+5-2
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,16 @@
1-
.vscode/
21
**.egg-info/
32
**/__pycache__/
43
logs
54
secrets
65
**/*.secret
76
saves
87
.pytest_cache
9-
**/*.json
8+
*.json
109
**/*.csv
1110
build
1211
.venv
12+
!/.devcontainer/devcontainer.json
13+
!/.devcontainer/devcontainer.linux.json
14+
!/.devcontainer/devcontainer.wsl2.json
15+
!/.vscode/launch.json
1316
c_generated_code

‎.vscode/launch.json

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
{
2+
// Use IntelliSense to learn about possible attributes.
3+
// Hover to view descriptions of existing attributes.
4+
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
5+
"version": "0.2.0",
6+
"configurations": [
7+
{
8+
"name": "Python Debugger: Current File",
9+
"type": "debugpy",
10+
"request": "launch",
11+
"program": "${file}",
12+
"console": "integratedTerminal",
13+
"justMyCode": false
14+
}
15+
]
16+
}

‎README.md

+34
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,37 @@
1+
### Using the project with VSCode devcontainers
2+
3+
**Running on CPU**: by default the containers run on CPU. You don't need to take any action.
4+
5+
**Running on GPU**: The devcontsainers can easily run using your computer's NVIDIA GPU on Linux and Windows. This makes sense if you want to accelerate simulation by running thousands of simulation in parallel. In order to work you need to install the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=WSL-Ubuntu&target_version=2.0&target_type=deb_local), [NVIDIA Container runtime](https://developer.nvidia.com/container-runtime) for your computer. Finally, enable GPU access to the devcontainers by setting the commented out `"--gpus=all"` and `"--runtime=nvidia"` flags in `devcontainer.json`.
6+
7+
8+
**Linux**
9+
10+
1. Install [Docker](https://docs.docker.com/engine/install/) (, and make sure Docker Daemon is running)
11+
2. Install [VSCode](https://code.visualstudio.com/), with [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
12+
3. Clone this project's code. Rename `/.devcontainer/devcontainer.linux.json` to `/.devcontainer/devcontainer.json`.
13+
4. Open this project in VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
14+
15+
**Windows** (requires Windows 10 or later)
16+
17+
For windows, we require WSL2 to run the devcontainer. (So its actually Linux with extra steps.) Full instructions can be found [in the official docs](https://code.visualstudio.com/blogs/2020/07/01/containers-wsl#_getting-started). Here are the important steps:
18+
1. Install [Docker](https://docs.docker.com/desktop/setup/install/windows-install/), and WSL2, and Ubuntu 22.04 LTS (, and make sure Docker Daemon is running)
19+
2. Docker will recognize that you have WSL installed and prompt you via Windows Notifications to enable WSL integration -> confirm this with `Enable WSL integration`. If not, open `Docker Desktop`, navigate to the settings, and manually enable WSL integration. (There are TWO setting options for this. Make sure to enable BOTH!)
20+
3. Install [VSCode](https://code.visualstudio.com/), with the [WSL extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-wsl), [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
21+
4. Clone the source code for the exercises in the WSL2 file system to `/home` (`~`), or wherever you like. (Performance when working on the WSL file system is much better compared to Windows file system). You can access the WSL filesystem either by starting a WSL/Ubuntu console, or via the Windows File Explorer at `\\wsl.localhost\Ubuntu\home` (replace `Ubuntu` with your distro, if necessary).
22+
7. Rename `/.devcontainer/devcontainer.windows.json` to `/.devcontainer/devcontainer.json`.
23+
8. Open this project in VSCode. The easiest way to do so is by starting a WSL/Ubuntu shell, navigating via `cd` to the source code, then type `code .` to open VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
24+
25+
26+
**MacOS**
27+
28+
Should work like Linux. Untested.
29+
30+
____________
31+
32+
Known Issues:
33+
- if building docker container fails at `RUN apt-get update`, make sure your host systems time is set correct: https://askubuntu.com/questions/1511514/docker-build-fails-at-run-apt-update-error-failed-to-solve-process-bin-sh
34+
135
# crazyflow
236
Fast, parallelizable simulations of Crazyflies with JAX and MuJoCo.
337

‎benchmark/main.py

+15-20
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from ml_collections import config_dict
88

99
import crazyflow # noqa: F401, ensure gymnasium envs are registered
10-
from crazyflow.sim.core import Sim
10+
from crazyflow.sim import Sim
1111

1212

1313
def analyze_timings(times: list[float], n_steps: int, n_worlds: int, freq: float) -> None:
@@ -19,8 +19,8 @@ def analyze_timings(times: list[float], n_steps: int, n_worlds: int, freq: float
1919
tmax, idx_tmax = np.max(times), np.argmax(times)
2020

2121
# Check for significant variance
22-
if tmax / tmin > 5:
23-
print("Warning: step time varies by more than 5x. Is JIT compiling during the benchmark?")
22+
if tmax / tmin > 10:
23+
print("Warning: step time varies by more than 10x. Is JIT compiling during the benchmark?")
2424
print(f"Times: max {tmax:.2e} @ {idx_tmax}, min {tmin:.2e} @ {idx_tmin}")
2525

2626
# Performance metrics
@@ -43,28 +43,23 @@ def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, devic
4343
device = jax.devices(device)[0]
4444

4545
envs = gymnasium.make_vec(
46-
"DroneReachPos-v0",
47-
time_horizon_in_seconds=2,
48-
return_datatype="numpy",
49-
num_envs=sim_config.n_worlds,
50-
**sim_config,
46+
"DroneReachPos-v0", time_horizon_in_seconds=3, num_envs=sim_config.n_worlds, **sim_config
5147
)
5248

5349
# Action for going up (in attitude control)
5450
action = np.zeros((sim_config.n_worlds, 4), dtype=np.float32)
5551
action[..., 0] = 0.3
5652
# Step through env once to ensure JIT compilation
57-
envs.reset_all(seed=42)
58-
envs.step(action)
53+
envs.reset(seed=42)
5954
envs.step(action)
6055

61-
jax.block_until_ready(envs.unwrapped.sim.states.pos) # Ensure JIT compiled dynamics
56+
jax.block_until_ready(envs.unwrapped.sim.data) # Ensure JIT compiled dynamics
6257

6358
# Step through the environment
6459
for _ in range(n_steps):
6560
tstart = time.perf_counter()
6661
envs.step(action)
67-
jax.block_until_ready(envs.unwrapped.sim.states.pos)
62+
jax.block_until_ready(envs.unwrapped.sim.data)
6863
times.append(time.perf_counter() - tstart)
6964

7065
envs.close()
@@ -83,14 +78,14 @@ def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
8378

8479
sim.reset()
8580
sim.attitude_control(cmd)
86-
sim.step()
87-
jax.block_until_ready(sim.states.pos) # Ensure JIT compiled dynamics
81+
sim.step(sim.freq // sim.control_freq)
82+
jax.block_until_ready(sim.data) # Ensure JIT compiled dynamics
8883

8984
for _ in range(n_steps):
9085
tstart = time.perf_counter()
9186
sim.attitude_control(cmd)
92-
sim.step()
93-
jax.block_until_ready(sim.states.pos)
87+
sim.step(sim.freq // sim.control_freq)
88+
jax.block_until_ready(sim.data)
9489
times.append(time.perf_counter() - tstart)
9590

9691
analyze_timings(times, n_steps, sim.n_worlds, sim.freq)
@@ -102,16 +97,16 @@ def main():
10297
sim_config = config_dict.ConfigDict()
10398
sim_config.n_worlds = 1
10499
sim_config.n_drones = 1
105-
sim_config.physics = "sys_id"
100+
sim_config.physics = "analytical"
106101
sim_config.control = "attitude"
107-
sim_config.controller = "emulatefirmware"
102+
sim_config.attitude_freq = 500
108103
sim_config.device = device
109104

110105
print("Simulator performance")
111-
profile_step(sim_config, 100, device)
106+
profile_step(sim_config, 1000, device)
112107

113108
print("\nGymnasium environment performance")
114-
profile_gym_env_step(sim_config, 100, device)
109+
profile_gym_env_step(sim_config, 1000, device)
115110

116111

117112
if __name__ == "__main__":

0 commit comments

Comments
 (0)