hyperdrone 0.1.2__tar.gz → 0.1.4__tar.gz

hyperdrone-0.1.4/.gitignore

@@ -0,0 +1 @@
+/dist/

hyperdrone-0.1.4/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.5
+Name: hyperdrone
+Version: 0.1.4
+Summary: The fastest drone simulator and renderer (raytracer).
+Author-email: Jonas Eschmann <jonas.eschmann@gmail.com>
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: nanobind>=2.0
+Requires-Dist: numpy
+Provides-Extra: examples
+Requires-Dist: imageio; extra == 'examples'
+Requires-Dist: imageio-ffmpeg; extra == 'examples'
+Requires-Dist: matplotlib; extra == 'examples'
+Requires-Dist: scipy; extra == 'examples'
+Provides-Extra: gym
+Requires-Dist: gymnasium>=1.0; extra == 'gym'

hyperdrone-0.1.4/README.md

@@ -0,0 +1,265 @@
+# hyperdrone
+
+Drone simulation stack for RLtools: raytracing renderer, L2F multirotor dynamics, and
+environment setup under one Python package.
+
+| Subpackage | What it is |
+|---|---|
+| `hyperdrone.render` | Raytracing renderer (OptiX / Metal / Vulkan / generic CPU). Self-contained — no dependency on drone dynamics. |
+| `hyperdrone.dynamics` | Vectorized, stateful L2F multirotor simulator (`Sim`); cpu and cuda variants. |
+| `hyperdrone.env` | The RL environment: the C++ `MultiEnvironment<hyperdrone::World>` driven through the exact rl_tools batch verbs. |
+| `hyperdrone.gym` | Optional Gymnasium `VectorEnv` adapter over `hyperdrone.env` (`pip install "hyperdrone[gym]"`). |
+| `hyperdrone.jit` | Shared compile-and-cache infrastructure both domain packages build on. |
+| `hyperdrone.cuda` | CUDA staging helpers (`upload` → DLPack tensor sets). |
+| `hyperdrone.conta` | Client for the sha1 content-addressed asset store — protocol twin of the C++ `include/conta/conta.h`, sharing one cache across both languages. |
+
+Native components are JIT-compiled per set of compile-time constants (resolution, camera
+count, drone count, ...) into a persistent cache; later uses load the cached library
+directly. `render` and `dynamics` are peers that never import each other — the coupling
+surface is a plain DLPack tensor of packed camera bases, device-resident on OptiX + CUDA.
+
+## Install
+
+From an rl-tools checkout (development):
+
+```bash
+.venv/bin/pip install -e interface/python/hyperdrone
+```
+
+From an sdist (self-contained — bundles the rl-tools headers and raytracing sources, no
+checkout needed):
+
+```bash
+pip install hyperdrone-<version>.tar.gz
+```
+
+Install the dependencies used by the packaged examples with:
+
+```bash
+pip install "hyperdrone[examples]"
+```
+
+Requirements: CMake >= 3.24, a C++17 compiler, and per feature: assimp (system package)
+for GLB loading, CUDA + OptiX driver for the OptiX render backend and cuda dynamics,
+Vulkan dev + glslang for the VULKAN backend. The WEBGPU backend fetches a hash-pinned
+wgpu-native prebuilt at first build (no extra system packages). Nothing beyond the
+compiler for GENERIC + cpu.
+
+## Environment
+
+| Variable | Meaning |
+|---|---|
+| `HYPERDRONE_RENDER_BACKEND` | `OPTIX` \| `METAL` \| `VULKAN` \| `WEBGPU` \| `GENERIC` \| `AUTO` (default: Metal on macOS, OptiX elsewhere) |
+| `HYPERDRONE_DYNAMICS_DEVICE` | `CPU` \| `CUDA` \| `AUTO` (default: cuda when available) |
+| `HYPERDRONE_CACHE_DIR` | root for CMake build trees and downloaded build dependencies (default `~/.cache/hyperdrone`) |
+| `HYPERDRONE_RLTOOLS_ROOT` | rl-tools source root override (default: enclosing checkout, else the vendored tree) |
+| `HYPERDRONE_PROCTHOR_PATH` | local ProcTHOR example scene override (otherwise downloaded and cached) |
+| `HYPERDRONE_SKIP_BUILD` | skip the CMake staleness check when the artifacts already exist |
+| `HYPERDRONE_OFFLINE` | forbid network during builds (requires seeded/vendored dependencies) |
+| `HYPERDRONE_BUILD_JOBS` | parallel build jobs (default 5) |
+
+Each component and variant gets its own build tree (`render-optix-*`, `dynamics-cuda-*`,
+...); switching never invalidates another's cache. All configure/build steps run under a
+per-tree file lock, so many worker processes can share one cache safely. CMake
+FetchContent sources and native build state live in `<cache>/.dependencies`; the native
+build does not write generated files into an editable checkout or installed package.
+
+Whenever a renderer is successfully allocated, RLtools reports the selected backend on
+stderr, for example `#rl_tools::rendering::raytracing: backend=metal`. The
+`Renderer.backend` property provides the same lowercase name programmatically.
+
+## Rendering
+
+```python
+import math, numpy as np
+from hyperdrone import render
+from hyperdrone.examples.data import procthor_scene_path
+
+scene = render.load_scene(procthor_scene_path(), fidelity="high")
+
+renderer = render.Renderer(width=320, height=240, num_cameras=4, output="rgbd", fidelity="high")
+renderer.init(scene)
+
+camera = renderer.camera(position=(0, 0, 1.5), look_at=(1, 0, 1.5), fov=math.radians(80))
+renderer.set_cameras(np.repeat(camera[None], 4, axis=0))
+renderer.render()
+
+rgb   = renderer.frame()    # (4, 240, 320, 4) uint8
+depth = renderer.depth()    # (4, 240, 320) float32
+```
+
+The renderer operates in the L2F FLU frame: +X forward, +Y left, +Z up; GLB assets are
+swizzled at load time. Procedural scenes (`render.Scene()` + `render.Object` /
+`render.Mesh` / `render.SceneLight`), segmentation output, overlays (per-camera dynamic
+content via `AssetPool` + `spawn`/`attach`), motion blur, and anti-aliasing follow the
+same API as before under `hyperdrone.render.*`.
+
+`num_probes` sizes the collision-probe pass (`generate_probe_directions()`, `render("collision")`, `collisions()`; one probe per camera by default). `num_probes=0` compiles that pass out entirely, the counterpart of the probe-ray macro the C++ benchmark targets set: no probe buffers are allocated, the default `render()` target runs only the image channels, and the probe methods raise.
+
+### Zero-copy I/O (DLPack)
+
+```python
+renderer.frame()                  # snapshot copy (safe to keep)
+renderer.frame(copy=False)        # numpy view of the staging buffer (refreshed in place)
+rgba = torch.from_dlpack(renderer.frame_dlpack())
+# live uint8 (num_cameras, height, width, 4), zero-copy GPU tensor on OptiX
+packed = torch.from_dlpack(renderer.frame_raw_dlpack())
+# same memory as packed uint32 (num_cameras, height, width)
+```
+
+Inputs accept any DLPack producer. CUDA-resident camera input (OptiX):
+`set_cameras` dispatches on `__dlpack_device__`, so a CUDA tensor — a torch GPU tensor,
+a set pre-uploaded with `hyperdrone.cuda.upload`, or `Sim.camera_bases()` — is handed
+over device-to-device on the render stream, fully async; `stream=` takes the producer's
+cudaStream_t handle for event-ordered hand-off.
+
+## Dynamics
+
+```python
+from hyperdrone import dynamics
+
+sim = dynamics.Sim(num_drones=4096, model="x500_sim", device="cuda")
+sim.reset(seed=0)                     # deterministic; host-sampled, identical on cpu/cuda
+sim.step(actions)                     # (N, 4) float32 in [-1, 1], host or CUDA tensor
+sim.state["position"]                 # zero-copy DLPack views, CUDA-resident on cuda
+sim.state.numpy("position")           # host copy
+sim.observe()                         # (N, 18): position, rotation matrix, velocities
+sim.parameters["mass"] = masses       # per-drone runtime parameters (until the next reset)
+
+cameras = sim.camera_bases(fov=math.radians(100), aspect=renderer.aspect)
+renderer.set_cameras(cameras, stream=sim.stream)   # device-resident hand-off
+```
+
+`num_drones` and domain randomization are compile-time (the JIT key); the model preset
+(`crazyflie`, `x500_real`, `x500_sim`, `mrs`, ...), integration `dt`, and all physical
+parameters are runtime. Reward and termination live in the
+environment (`hyperdrone.env`), not the simulator. One process can use one dynamics
+variant (cpu or cuda).
+
+## RL environment
+
+```python
+import numpy as np
+from hyperdrone.env import EnvConfig, MultiEnvironment
+
+config = EnvConfig(num_environments=1, instances=1024, cam_width=32, cam_height=32, task="target_frame")
+env = MultiEnvironment("scenes/", config=config, seed=0)
+
+mask = np.ones(env.total_instances, dtype=np.uint8)
+env.reset(mask)                    # resample parameters + states where mask is set
+env.render(mask)                   # render the FPV cameras (mask marks fresh episodes)
+observations = env.observe()       # (total, observation_dim) float32
+env.step(actions)                  # (total, action_dim) float32 in [-1, 1]
+rewards, terminated = env.rewards(), env.terminated()
+critic_input = env.observe_privileged()
+env.rotate_scene()                 # deterministic scene rotation; reset all instances after
+env.observation_layout             # named blocks: which channels/values mean what
+```
+
+`reset(mask)` samples parameters and states immediately. `observe()` renders as needed,
+so explicit `render(mask)` calls are optional. The training loop owns episode counters and
+chooses when to reset; the Gymnasium adapter below supplies same-step autoreset and time limits.
+
+All environment semantics — reset, reward, termination, scene scheduling, observation
+composition — live on the C++ side (`rl_tools::rl::environments::hyperdrone::MultiEnvironment<World>`);
+the binding marshals tensors and nothing else, and a seeded rollout is pinned bit-exact
+against the C++ verbs by a golden test. The scenes argument is a directory of `.glb`
+scenes (sorted corpus) or a list of references — `.glb` paths, `conta:HASH` strings, or
+conta store manifest entries `{"description": ..., "hash": ...}` — used as the corpus in
+list order and partitioned across environments; conta references (scenes, `drone_asset=`,
+`gate_asset=`) resolve through `hyperdrone.conta` into the shared content-addressed
+cache, so a config pins the exact corpus on any machine. Configuration follows the C++
+extension ladder:
+`preset=` names the platform (`"crazyflie"`, `"x500_fpv"` — SELF_VISIBLE, pass a
+body/prop_* GLB via `drone_asset=` — and `"x500_fpv_imu"`, the same platform stepping at
+IMU rate), `task=` names the wrapper (`"target_frame"`, `"moving_gate"`,
+`"visual_inertial_localization"`), `n_agents=` enables multi-agent, and `EnvConfig(spec_header=...)` pins
+an arbitrary C++ specification (a header defining `hyperdrone_env_user::WORLD`, hashed
+into the JIT key). A spec header can go beyond constants to a full user-authored task
+wrapper — verb overloads registering and moving entities, compiled reward, termination:
+`hyperdrone/examples/Orbiter.ipynb` walks through one, and its contract is pinned by
+`tests/env/user_task_header.h`. The render backend for the environment follows
+`HYPERDRONE_ENV_BACKEND` (default: the render backend selection).
+
+Manual composition of `Sim` and `Renderer` (no MDP — rendering research, data
+generation) remains first-class: see `python -m hyperdrone.examples.drone_flythrough`
+for the wiring. `FreeSpaceSampler` lives in `hyperdrone.render` — it rides the
+renderer's collision probes and is deterministic given a seed.
+
+### Visual-inertial localization
+
+`task="visual_inertial_localization"` on `preset="x500_fpv_imu"` is the localization
+benchmark: fixed-length synchronized episodes stepping at IMU rate (`env.dt`), a camera
+frame every `env.frame_stride` steps, and the IMU sample of each step from
+`env.observe_imu()` (named by `env.observation_layout_imu`: accelerometer, gyroscope,
+frame age, new-frame flag). The task is autonomous: the RAPTOR autopilot built into it
+flies a waypoint route (its checkpoint is pinned by content hash and fetched through
+conta on first use), so `action_dim` is 0 and `step()` takes no actions. The privileged
+observation carries the current `waypoint` next to the dynamics state; the drone's own
+body is not rendered, matching the C++ benchmark harness. Resets must be all-or-none per
+environment.
+
+```python
+config = EnvConfig(instances=2, cam_width=160, cam_height=120, preset="x500_fpv_imu", task="visual_inertial_localization")
+env = MultiEnvironment(scenes, config=config, seed=0)
+env.reset(); env.render(np.ones(env.total_instances, dtype=np.uint8))
+for step in range(env.episode_step_limit):
+    if step > 0: env.render()
+    if step % env.frame_stride == 0:
+        frames = env.frames()          # the estimator's camera frame
+    env.step()
+    imu = env.observe_imu()            # the estimator's IMU sample for this step
+    truth = env.observe_privileged()   # position / rotation matrix / velocities / waypoint
+```
+
+`python -m hyperdrone.examples.visual_inertial_localization` runs an episode and scores
+IMU dead reckoning against the ground truth, the same protocol as the C++ demo.
+
+### Gymnasium
+
+```python
+from hyperdrone.gym import VectorEnv   # pip install "hyperdrone[gym]"
+env = VectorEnv("scenes/", config=EnvConfig(instances=1024), seed=0)
+observations, infos = env.reset()
+observations, rewards, terminations, truncations, infos = env.step(actions)
+```
+
+The adapter resets instances that terminate or reach `episode_step_limit` within the same
+step. Returned observations start the new episode for those instances; no final observation
+is exposed. `truncations` flags time limits that are not terminations. The core packages
+never import gymnasium.
+
+End-to-end example: `python -m hyperdrone.examples.drone_flythrough`; renderer benchmark:
+`python -m hyperdrone.examples.benchmark` (flag-compatible with the C++ benchmark
+counterpart).
+
+## Publishing
+
+From the repository root, build and upload a release with:
+
+```bash
+interface/python/hyperdrone/scripts/publish.sh
+```
+
+Set the release version in `pyproject.toml` and `hyperdrone/__init__.py` first. The
+script uses the repository's `.venv` (creating it if needed), installs the build and
+upload tools, builds a bundled source archive and a `py3-none-any` wheel from that
+archive, checks both with Twine, and uploads both to PyPI. Authenticate through your
+existing Twine configuration (`~/.pypirc`, keyring, or `TWINE_PASSWORD`), or enter a
+PyPI API token when prompted.
+
+Use `interface/python/hyperdrone/scripts/publish.sh --build-only` to build and check
+without uploading. Each run keeps its two artifacts in a fresh
+`interface/python/hyperdrone/dist/release.*` directory, so older builds are never
+included in an upload. Native components remain JIT-compiled on first use.
+
+## Tests
+
+```bash
+HYPERDRONE_RENDER_BACKEND=GENERIC .venv/bin/python -m pytest interface/python/hyperdrone/tests -v
+```
+
+`tests/jit` exercises the build infrastructure against a toy component in seconds;
+`tests/test_architecture.py` enforces the import DAG (render and dynamics never see each
+other); `scripts/test_sdist.sh` is the self-containment rot guard (build sdist → clean
+venv → vendored-tree render).

hyperdrone-0.1.4/hatch_build.py

@@ -0,0 +1,39 @@
+"""sdist self-containment: snapshot the required subset of the rl-tools repository into
+hyperdrone/_vendor/rl-tools so `pip install hyperdrone` works without a checkout. The
+manifest is explicit — never "copy everything". Wheels built directly from a checkout skip
+vendoring (they are for local use where the checkout exists); release wheels are built
+from the sdist, where the vendored tree is already on disk."""
+from pathlib import Path
+
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+
+VENDOR_MANIFEST = (
+    "include/rl_tools",
+    "include/conta",
+    "src/rendering/raytracing",
+    "CMakeLists.txt",
+    "cmake",
+    "LICENSE",
+)
+
+
+class VendorRlToolsHook(BuildHookInterface):
+    def initialize(self, version, build_data):
+        if self.target_name != "sdist":
+            return
+        repository = self._repository_root()
+        if repository is None:
+            return  # building from an sdist: the vendored tree is already in place
+        force_include = build_data.setdefault("force_include", {})
+        for entry in VENDOR_MANIFEST:
+            source = repository / entry
+            if source.exists():
+                force_include[str(source)] = f"hyperdrone/_vendor/rl-tools/{entry}"
+
+    def _repository_root(self):
+        candidate = Path(self.root)
+        for _ in range(8):
+            if (candidate / "include" / "rl_tools").is_dir():
+                return candidate
+            candidate = candidate.parent
+        return None

hyperdrone-0.1.4/hyperdrone/__init__.py

@@ -0,0 +1,31 @@
+"""hyperdrone — drone simulation stack for RLtools.
+
+Subpackages:
+  hyperdrone.render   — raytracing renderer (self-contained, no drone dynamics)
+  hyperdrone.dynamics — L2F multirotor simulator (vectorized, stateful Sim; pure dynamics)
+  hyperdrone.env      — the RL environment: the C++ MultiEnvironment behind the batch verbs
+  hyperdrone.gym      — optional Gymnasium VectorEnv adapter (pip install "hyperdrone[gym]")
+  hyperdrone.jit      — shared compile-and-cache infrastructure
+  hyperdrone.cuda     — CUDA staging helpers (DLPack tensor sets)
+
+The three domain packages are independent: render and dynamics couple only through a
+DLPack tensor of camera bases; env owns no Python-side semantics — it drives the C++
+environment. Native components are JIT-compiled per set of compile-time constants and
+cached under ~/.cache/hyperdrone (override with HYPERDRONE_CACHE_DIR). Importing
+hyperdrone (or any subpackage) is cheap; the first build happens on first use.
+"""
+import importlib
+
+__version__ = "0.1.4"
+
+_SUBPACKAGES = ("render", "dynamics", "env", "gym", "jit", "cuda")
+
+
+def __getattr__(name):
+    if name in _SUBPACKAGES:
+        return importlib.import_module(f".{name}", __name__)
+    raise AttributeError(f"module 'hyperdrone' has no attribute {name}")
+
+
+def __dir__():
+    return sorted(list(globals()) + list(_SUBPACKAGES))

hyperdrone-0.1.4/hyperdrone/_native/dynamics/CMakeLists.txt

@@ -0,0 +1,68 @@
+cmake_minimum_required(VERSION 3.24)
+project(hyperdrone_dynamics LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+set(RLTOOLS_ROOT "" CACHE PATH "Path to the rl-tools source root")
+set(HYPERDRONE_VARIANT "CPU" CACHE STRING "Dynamics device (CPU|CUDA)")
+
+if(NOT RLTOOLS_ROOT)
+    message(FATAL_ERROR "RLTOOLS_ROOT must point to the rl-tools source root")
+endif()
+
+set(RL_TOOLS_DISABLE_TARGETS ON CACHE BOOL "" FORCE)
+set(RL_TOOLS_DISABLE_TESTS ON CACHE BOOL "" FORCE)
+set(RL_TOOLS_DISABLE_GIT_DIFF ON CACHE BOOL "" FORCE)
+set(RL_TOOLS_RL_ENVIRONMENTS_DISABLE_MUJOCO ON CACHE BOOL "" FORCE)
+set(RL_TOOLS_RENDERING_DISABLE_RAYTRACING ON CACHE BOOL "" FORCE)
+set(RL_TOOLS_WARNINGS_AS_ERRORS OFF CACHE BOOL "" FORCE)
+
+if(HYPERDRONE_VARIANT STREQUAL "CUDA")
+    enable_language(CUDA)
+    set(CMAKE_CUDA_STANDARD 17)
+    set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+        set(CMAKE_CUDA_ARCHITECTURES native)
+    endif()
+endif()
+
+add_subdirectory(${RLTOOLS_ROOT} rl_tools EXCLUDE_FROM_ALL)
+
+find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
+execute_process(
+    COMMAND ${Python_EXECUTABLE} -m nanobind --cmake_dir
+    OUTPUT_VARIABLE nanobind_DIR
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    COMMAND_ERROR_IS_FATAL ANY
+)
+find_package(nanobind CONFIG REQUIRED)
+
+nanobind_add_module(hyperdrone_dynamics_core host.cpp)
+target_include_directories(hyperdrone_dynamics_core PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../common)
+target_link_libraries(hyperdrone_dynamics_core PRIVATE ${CMAKE_DL_LIBS})
+set_target_properties(hyperdrone_dynamics_core PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
+
+set(HYPERDRONE_DYNAMICS_EXTRA_LINK "")
+if(HYPERDRONE_VARIANT STREQUAL "CUDA")
+    find_package(CUDAToolkit REQUIRED)
+    target_compile_definitions(hyperdrone_dynamics_core PRIVATE HYPERDRONE_DYNAMICS_CUDA=1)
+    target_link_libraries(hyperdrone_dynamics_core PRIVATE CUDA::cudart)
+    add_compile_definitions(HYPERDRONE_DYNAMICS_CUDA=1)
+    set_source_files_properties(impl.cpp PROPERTIES LANGUAGE CUDA)
+else()
+    # per-drone parallelism on the CPU variant; drones are independent, so this stays
+    # deterministic for a fixed seed
+    find_package(OpenMP QUIET)
+    if(OpenMP_CXX_FOUND)
+        set(HYPERDRONE_DYNAMICS_EXTRA_LINK OpenMP::OpenMP_CXX)
+    endif()
+endif()
+
+include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/hyperdrone_jit.cmake)
+hyperdrone_jit_add_configs(
+    COMPONENT dynamics
+    SOURCES impl.cpp
+    LINK RLtools::RLtools ${HYPERDRONE_DYNAMICS_EXTRA_LINK}
+)

hyperdrone-0.1.4/hyperdrone/_native/dynamics/host.cpp

@@ -0,0 +1,174 @@
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/string.h>
+
+#include "iface.h"
+#include "jit_host.h"
+
+#include <cstring>
+#include <stdexcept>
+#include <string>
+
+#if defined(HYPERDRONE_DYNAMICS_CUDA)
+#include "cuda_staging.h"
+#include <cuda_runtime.h>
+#endif
+
+namespace nb = nanobind;
+namespace hdd = hyperdrone::dynamics;
+
+struct JitSim {
+    hyperdrone::JitInstance<hdd::Sim> instance;
+
+    JitSim(const std::string& path, const std::string& expected_config)
+        : instance(path, "hyperdrone_dynamics", HYPERDRONE_DYNAMICS_IFACE_VERSION, expected_config) {}
+
+    hdd::Sim* operator->() const { return instance.instance; }
+
+    size_t component_dim(int component) const {
+        if(component < 0 || component > 4){
+            throw std::invalid_argument("hyperdrone: unknown state component");
+        }
+        return component == 4 ? instance.instance->config().action_dim
+                              : (size_t)hdd::state_component_dims[component];
+    }
+
+    void copy_out(const float* source, float* destination, size_t count) const {
+#if defined(HYPERDRONE_DYNAMICS_CUDA)
+        instance.instance->synchronize();
+        cudaMemcpy(destination, source, count * sizeof(float), cudaMemcpyDeviceToHost);
+#else
+        std::memcpy(destination, source, count * sizeof(float));
+#endif
+    }
+};
+
+using HostFloatArray = nb::ndarray<const float, nb::c_contig, nb::device::cpu>;
+
+NB_MODULE(hyperdrone_dynamics_core, m){
+    m.doc() = "hyperdrone dynamics core: JIT loader for the vectorized L2F multirotor simulator";
+
+    auto jit_sim_class = nb::class_<JitSim>(m, "JitSim");
+    jit_sim_class
+        .def(nb::init<const std::string&, const std::string&>(), nb::arg("library_path"), nb::arg("expected_config"))
+        .def_prop_ro("num_drones", [](const JitSim& jit){ return jit->config().num_drones; })
+        .def_prop_ro("action_dim", [](const JitSim& jit){ return jit->config().action_dim; })
+        .def_prop_ro("observation_dim", [](const JitSim& jit){ return jit->config().observation_dim; })
+        .def_prop_ro("domain_randomization", [](const JitSim& jit){ return jit->config().domain_randomization; })
+        .def_prop_ro("device_name", [](const JitSim& jit){ return std::string(jit->device_name()); })
+        .def_prop_ro("buffer_device_type", [](const JitSim& jit){ return jit->buffer_device_type(); })
+        .def_prop_ro("stream", [](const JitSim& jit){ return (uintptr_t)jit->stream(); })
+        .def("synchronize", [](JitSim& jit){
+            nb::gil_scoped_release release;
+            jit->synchronize();
+        })
+        .def("set_model", [](JitSim& jit, const std::string& name){ return jit->set_model(name.c_str()); }, nb::arg("name"))
+        .def("set_dt", [](JitSim& jit, float dt){ jit->set_dt(dt); }, nb::arg("dt"))
+        .def_prop_ro("dt", [](const JitSim& jit){ return jit->dt(); })
+        .def("reset", [](JitSim& jit, uint64_t seed, bool sample_parameters, bool sample_states){
+            nb::gil_scoped_release release;
+            jit->reset(seed, sample_parameters, sample_states);
+        }, nb::arg("seed"), nb::arg("sample_parameters"), nb::arg("sample_states"))
+        .def("step", [](JitSim& jit, HostFloatArray actions){
+            const hdd::Config config = jit->config();
+            if(actions.ndim() != 2 || actions.shape(0) != config.num_drones || actions.shape(1) != config.action_dim){
+                throw std::invalid_argument("hyperdrone: actions must have shape (num_drones, action_dim) and dtype float32");
+            }
+            nb::gil_scoped_release release;
+            jit->step(actions.data(), 1, 0);
+        }, nb::arg("actions"))
+        .def("state_dlpack", [](JitSim& jit, int component){
+            const hdd::Config config = jit->config();
+            const size_t dim = jit.component_dim(component);
+            return nb::ndarray<>((void*)jit->state_buffer(component), {config.num_drones, dim}, nb::find(&jit), {},
+                                 nb::dtype<float>(), jit->buffer_device_type(), 0);
+        }, nb::arg("component"))
+        .def("read_state", [](JitSim& jit, int component, nb::ndarray<float, nb::c_contig, nb::device::cpu> out){
+            const hdd::Config config = jit->config();
+            const size_t dim = jit.component_dim(component);
+            if(out.size() != config.num_drones * dim){
+                throw std::invalid_argument("hyperdrone: output array has the wrong size for this state component");
+            }
+            nb::gil_scoped_release release;
+            jit.copy_out(jit->state_buffer(component), out.data(), config.num_drones * dim);
+        }, nb::arg("component"), nb::arg("out"))
+        .def("set_state", [](JitSim& jit, int component, HostFloatArray values){
+            const hdd::Config config = jit->config();
+            const size_t dim = jit.component_dim(component);
+            if(values.ndim() != 2 || values.shape(0) != config.num_drones || values.shape(1) != dim){
+                throw std::invalid_argument("hyperdrone: state component values must have shape (num_drones, dim) and dtype float32");
+            }
+            nb::gil_scoped_release release;
+            jit->set_state_component(component, values.data());
+        }, nb::arg("component"), nb::arg("values"))
+        .def("update_observations", [](JitSim& jit){
+            nb::gil_scoped_release release;
+            jit->update_observations();
+        })
+        .def("observations_dlpack", [](JitSim& jit){
+            const hdd::Config config = jit->config();
+            return nb::ndarray<>((void*)jit->observation_buffer(), {config.num_drones, config.observation_dim}, nb::find(&jit), {},
+                                 nb::dtype<float>(), jit->buffer_device_type(), 0);
+        })
+        .def("read_observations", [](JitSim& jit, nb::ndarray<float, nb::c_contig, nb::device::cpu> out){
+            const hdd::Config config = jit->config();
+            if(out.size() != (size_t)config.num_drones * config.observation_dim){
+                throw std::invalid_argument("hyperdrone: observations output must have num_drones * observation_dim elements");
+            }
+            nb::gil_scoped_release release;
+            jit.copy_out(jit->observation_buffer(), out.data(), (size_t)config.num_drones * config.observation_dim);
+        }, nb::arg("out"))
+        .def("update_camera_bases", [](JitSim& jit, HostFloatArray mount, float fov, float aspect){
+            if(mount.size() != 12){
+                throw std::invalid_argument("hyperdrone: mount transform must have 12 elements (3x4 row-major)");
+            }
+            nb::gil_scoped_release release;
+            jit->update_camera_bases(mount.data(), fov, aspect);
+        }, nb::arg("mount"), nb::arg("fov"), nb::arg("aspect"))
+        .def("camera_bases_dlpack", [](JitSim& jit){
+            const hdd::Config config = jit->config();
+            return nb::ndarray<>((void*)jit->camera_bases_buffer(), {config.num_drones, (size_t)12}, nb::find(&jit), {},
+                                 nb::dtype<float>(), jit->buffer_device_type(), 0);
+        })
+        .def("read_camera_bases", [](JitSim& jit, nb::ndarray<float, nb::c_contig, nb::device::cpu> out){
+            const hdd::Config config = jit->config();
+            if(out.size() != (size_t)config.num_drones * 12){
+                throw std::invalid_argument("hyperdrone: camera bases output must have num_drones * 12 elements");
+            }
+            nb::gil_scoped_release release;
+            jit.copy_out(jit->camera_bases_buffer(), out.data(), (size_t)config.num_drones * 12);
+        }, nb::arg("out"))
+        .def("read_parameter", [](JitSim& jit, const std::string& name, nb::ndarray<float, nb::c_contig, nb::device::cpu> out){
+            if(out.size() != jit->config().num_drones){
+                throw std::invalid_argument("hyperdrone: parameter output must have num_drones elements");
+            }
+            if(!jit->read_parameter(name.c_str(), out.data())){
+                throw std::invalid_argument("hyperdrone: unknown parameter " + name);
+            }
+        }, nb::arg("name"), nb::arg("out"))
+        .def("write_parameter", [](JitSim& jit, const std::string& name, HostFloatArray values){
+            if(values.size() != jit->config().num_drones){
+                throw std::invalid_argument("hyperdrone: parameter values must have num_drones elements");
+            }
+            if(!jit->write_parameter(name.c_str(), values.data())){
+                throw std::invalid_argument("hyperdrone: unknown parameter " + name);
+            }
+        }, nb::arg("name"), nb::arg("values"));
+
+#if defined(HYPERDRONE_DYNAMICS_CUDA)
+    m.attr("HAS_CUDA") = true;
+
+    hyperdrone::register_cuda_staging(m);
+
+    jit_sim_class
+        .def("step_device", [](JitSim& jit, nb::ndarray<const float, nb::c_contig, nb::device::cuda> actions, uintptr_t stream){
+            const hdd::Config config = jit->config();
+            if(actions.ndim() != 2 || actions.shape(0) != config.num_drones || actions.shape(1) != config.action_dim){
+                throw std::invalid_argument("hyperdrone: actions must have shape (num_drones, action_dim) and dtype float32");
+            }
+            jit->step(actions.data(), 2, stream);
+        }, nb::arg("actions"), nb::arg("stream") = 0);
+#else
+    m.attr("HAS_CUDA") = false;
+#endif
+}