hyperverse 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/src/rendering/raytracing/example/environment/environment.h

@@ -0,0 +1,96 @@
+#pragma once
+
+#include <rl_tools/rendering/raytracing/renderer.h>
+#include <array>
+
+namespace rl_tools::rl::environments::raytracing_example {
+    template <typename T_T, typename T_TI, T_TI T_NUM_ENVS, T_TI T_CAM_WIDTH = 64, T_TI T_CAM_HEIGHT = 64, T_TI T_NUM_PROBES = 64, typename T_SHADING = rendering::raytracing::Medium, bool T_ENABLE_MOTION_BLUR = false, T_TI T_MOTION_BLUR_SAMPLES = 1, bool T_ENABLE_ANTI_ALIASING = false, T_TI T_ANTI_ALIASING_GRID_SIZE = 1, bool T_OUTPUT_RGB = true, bool T_OUTPUT_DEPTH = false, bool T_OUTPUT_SEGMENTATION = false, bool T_OUTPUT_FLOW = false, bool T_OUTPUT_FLOW_FORWARD = false, T_TI T_NUM_OVERLAYS = 0, T_TI T_MAX_OVERLAY_INSTANCES = 0, T_TI T_MAX_OVERLAYS_PER_CAMERA = 0>
+    struct Specification {
+        using T = T_T;
+        using TI = T_TI;
+        using SHADING = T_SHADING;
+        static constexpr TI NUM_ENVS = T_NUM_ENVS;
+        static constexpr TI CAM_WIDTH = T_CAM_WIDTH;
+        static constexpr TI CAM_HEIGHT = T_CAM_HEIGHT;
+        static constexpr T FOV = 80;
+        static constexpr TI NUM_PROBES = T_NUM_PROBES;
+        static constexpr bool ENABLE_MOTION_BLUR = T_ENABLE_MOTION_BLUR;
+        static constexpr TI MOTION_BLUR_SAMPLES = T_MOTION_BLUR_SAMPLES;
+        static constexpr bool ENABLE_ANTI_ALIASING = T_ENABLE_ANTI_ALIASING;
+        static constexpr TI ANTI_ALIASING_GRID_SIZE = T_ANTI_ALIASING_GRID_SIZE;
+        static constexpr bool HAS_RGB = T_OUTPUT_RGB;
+        static constexpr bool HAS_DEPTH = T_OUTPUT_DEPTH;
+        static constexpr bool ENABLE_DEPTH = HAS_DEPTH;
+        struct RENDERER_CONFIG: rendering::raytracing::config::Default<T, TI>{
+            static constexpr TI CAM_WIDTH = T_CAM_WIDTH, CAM_HEIGHT = T_CAM_HEIGHT, NUM_CAMERAS = T_NUM_ENVS, NUM_PROBES = T_NUM_PROBES;
+            using SHADING = T_SHADING;
+            static constexpr bool OUTPUT_RGB = T_OUTPUT_RGB;
+            static constexpr bool OUTPUT_DEPTH = T_OUTPUT_DEPTH;
+            static constexpr bool OUTPUT_SEGMENTATION = T_OUTPUT_SEGMENTATION;
+            static constexpr bool OUTPUT_FLOW = T_OUTPUT_FLOW;
+            static constexpr bool OUTPUT_FLOW_FORWARD = T_OUTPUT_FLOW_FORWARD;
+            static constexpr TI NUM_OVERLAYS = T_NUM_OVERLAYS;
+            static constexpr TI MAX_OVERLAY_INSTANCES = T_MAX_OVERLAY_INSTANCES;
+            static constexpr TI MAX_OVERLAYS_PER_CAMERA = T_MAX_OVERLAYS_PER_CAMERA;
+            static constexpr bool ENABLE_MOTION_BLUR = T_ENABLE_MOTION_BLUR;
+            static constexpr TI MOTION_BLUR_SAMPLES = T_MOTION_BLUR_SAMPLES;
+            static constexpr bool ENABLE_ANTI_ALIASING = T_ENABLE_ANTI_ALIASING;
+            static constexpr TI ANTI_ALIASING_GRID_SIZE = T_ANTI_ALIASING_GRID_SIZE;
+        };
+        using RAYTRACING_SPEC = rendering::raytracing::Specification<RENDERER_CONFIG>;
+    };
+
+    enum class ObjectID : int {
+        CAMERA_RIG = 0,
+        COUNT
+    };
+
+    template <typename SPEC>
+    struct Parameters {
+        using T = typename SPEC::T;
+        T scene_translation[3];
+        T base_height;
+    };
+
+    template <typename SPEC>
+    struct State {
+        using T = typename SPEC::T;
+        T position[3];
+        T velocity[3];
+        T yaw;
+    };
+
+    template <typename SPEC>
+    struct ObservationRGB {
+        using TI = typename SPEC::TI;
+        static constexpr TI DIM = SPEC::CAM_HEIGHT * SPEC::CAM_WIDTH;
+    };
+
+    template <typename SPEC>
+    struct Environment {
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        using RAYTRACING_SPEC = typename SPEC::RAYTRACING_SPEC;
+        using Renderer = rendering::raytracing::Renderer<RAYTRACING_SPEC>;
+
+        Renderer* renderer = nullptr;
+        rendering::Bundle<T>* bundle = nullptr;
+        rendering::raytracing::AssetPool pool; // overlay assets, filled by the caller before init when NUM_OVERLAYS > 0
+        const char* scene_path = nullptr;
+        T dt = 1.0f / 60.0f;
+        T max_velocity = 1.5f;
+        T acceleration = 3.0f;
+        T look_ahead = 1.0f;
+        T eye_height = 1.6f;
+        bool owns_renderer = false;
+        bool owns_bundle = false;
+
+        // host-resident camera staging for the observe paths, owned like the renderer
+        using CAMERA_STAGING_SPEC = tensor::Specification<rendering::raytracing::Camera<T>, TI, tensor::Shape<TI, SPEC::NUM_ENVS>>;
+        Tensor<CAMERA_STAGING_SPEC> camera_staging;
+
+        static constexpr TI NUM_INITIAL_STATES = SPEC::NUM_ENVS;
+        std::array<State<SPEC>, NUM_INITIAL_STATES> indoor_initial_states{};
+        TI num_indoor_initial_states = 0;
+    };
+}

hyperdrone/_vendor/rl-tools/src/rendering/raytracing/example/environment/operations_cpu.h

@@ -0,0 +1,263 @@
+#pragma once
+
+#include <rl_tools/operations/cpu_mux.h>
+#include "environment.h"
+
+#include <rl_tools/rendering/raytracing/operations_cpu_mux.h>
+#include <rl_tools/rendering/datasets/glb/operations_cpu.h>
+#include <rl_tools/rendering/datasets/procthor/procthor.h>
+#include <rl_tools/rendering/datasets/procthor/operations_cpu.h>
+#include <rl_tools/rendering/datasets/annotations/operations_cpu.h>
+
+#include <cmath>
+#include <string>
+
+namespace rl_tools {
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void precompute_indoor_initial_states(DEVICE& device, rl::environments::raytracing_example::Environment<SPEC>& env) {
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        using ANNOTATIONS_SPEC = rendering::datasets::annotations::FreeSpaceSpecification<T, TI>;
+        rendering::datasets::annotations::FreeSpace<ANNOTATIONS_SPEC> annotations;
+        rendering::datasets::annotations::FreeSpaceParameters<T, TI> free_space_parameters{};
+        rendering::datasets::annotations::annotate(device, annotations, env.bundle->metadata, *env.renderer, free_space_parameters);
+        TI take_n = std::min(annotations.num_positions, rl::environments::raytracing_example::Environment<SPEC>::NUM_INITIAL_STATES);
+        for(TI i = 0; i < take_n; i++){
+            auto& s = env.indoor_initial_states[i];
+            s.position[0] = annotations.positions[i].position[0];
+            s.position[1] = annotations.positions[i].position[1];
+            s.position[2] = annotations.positions[i].position[2];
+            s.velocity[0] = static_cast<T>(0);
+            s.velocity[1] = static_cast<T>(0);
+            s.velocity[2] = static_cast<T>(0);
+            s.yaw = annotations.positions[i].yaw;
+        }
+        env.num_indoor_initial_states = take_n;
+    }
+
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::environments::raytracing_example::Environment<SPEC>& env) {
+        if (env.renderer == nullptr) {
+            env.renderer = new typename rl::environments::raytracing_example::Environment<SPEC>::Renderer{};
+            env.owns_renderer = true;
+            malloc(device, *env.renderer);
+            malloc(device, env.camera_staging);
+        }
+        if (env.bundle == nullptr) {
+            env.bundle = new rendering::Bundle<typename SPEC::T>{};
+            env.owns_bundle = true;
+        }
+    }
+
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::environments::raytracing_example::Environment<SPEC>& env) {
+        if (env.renderer != nullptr) {
+            free(device, *env.renderer);
+            if (env.owns_renderer) {
+                delete env.renderer;
+                free(device, env.camera_staging);
+            }
+            env.renderer = nullptr;
+            env.owns_renderer = false;
+        }
+        if (env.bundle != nullptr) {
+            if (env.owns_bundle) {
+                delete env.bundle;
+            }
+            env.bundle = nullptr;
+            env.owns_bundle = false;
+        }
+    }
+
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, rl::environments::raytracing_example::Environment<SPEC>& env) {
+        utils::assert_exit(device, env.scene_path != nullptr, "raytracing_example::init: scene_path is null");
+        bool loaded;
+        if constexpr (SPEC::RAYTRACING_SPEC::HAS_SEGMENTATION) {
+            rendering::raytracing::ObjectAssembly assembly;
+            loaded = load<typename SPEC::RAYTRACING_SPEC::SHADING, SPEC::RAYTRACING_SPEC::HAS_RGB>(device, assembly, std::string(env.scene_path));
+            if(loaded){
+                add(device, env.bundle->scene, assembly);
+                rendering::datasets::compute_bounds(device, *env.bundle);
+            }
+        }
+        else {
+            loaded = load<typename SPEC::RAYTRACING_SPEC::SHADING, SPEC::RAYTRACING_SPEC::HAS_RGB>(device, *env.bundle, std::string(env.scene_path));
+        }
+        utils::assert_exit(device, loaded, "raytracing_example::init: failed to load scene");
+
+        if constexpr (SPEC::RAYTRACING_SPEC::ENABLE_OVERLAYS) {
+            init(device, *env.renderer, *env.bundle, env.pool);
+        }
+        else {
+            init(device, *env.renderer, *env.bundle);
+        }
+        generate_probe_directions(device, *env.renderer);
+        precompute_indoor_initial_states(device, env);
+    }
+
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void initial_parameters(DEVICE& device, const rl::environments::raytracing_example::Environment<SPEC>& env, rl::environments::raytracing_example::Parameters<SPEC>& parameters) {
+        parameters.scene_translation[0] = 0;
+        parameters.scene_translation[1] = 0;
+        parameters.scene_translation[2] = 0;
+        parameters.base_height = 0;
+    }
+
+    template <typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample_initial_parameters(DEVICE& device, const rl::environments::raytracing_example::Environment<SPEC>& env, rl::environments::raytracing_example::Parameters<SPEC>& parameters, RNG& rng) {
+        initial_parameters(device, env, parameters);
+    }
+
+    template <typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample_initial_state(DEVICE& device, const rl::environments::raytracing_example::Environment<SPEC>& env, rl::environments::raytracing_example::Parameters<SPEC>& parameters, rl::environments::raytracing_example::State<SPEC>& state, RNG& rng) {
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+
+        utils::assert_exit(device, env.num_indoor_initial_states > 0, "raytracing_example::sample_initial_state: no indoor initial states available");
+        const TI index = random::uniform_int_distribution(device.random, static_cast<TI>(0), static_cast<TI>(env.num_indoor_initial_states - 1), rng);
+        state = env.indoor_initial_states[index];
+        state.position[2] = parameters.base_height;
+        state.velocity[0] = static_cast<T>(0);
+        state.velocity[1] = static_cast<T>(0);
+        state.velocity[2] = static_cast<T>(0);
+    }
+
+    template <typename DEVICE, typename SPEC, typename ACTION_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC::T step(DEVICE& device, const rl::environments::raytracing_example::Environment<SPEC>& env, rl::environments::raytracing_example::Parameters<SPEC>& parameters, const rl::environments::raytracing_example::State<SPEC>& state, const Matrix<ACTION_SPEC>& action, rl::environments::raytracing_example::State<SPEC>& next_state, RNG& rng) {
+        using T = typename SPEC::T;
+        static_assert(ACTION_SPEC::ROWS == 1);
+        static_assert(ACTION_SPEC::COLS == 3);
+
+        T desired_vx = get(action, 0, 0) * env.max_velocity;
+        T desired_vy = get(action, 0, 1) * env.max_velocity;
+        T yaw_rate = get(action, 0, 2);
+
+        T alpha = env.acceleration * env.dt;
+        if (alpha > static_cast<T>(1)) {
+            alpha = static_cast<T>(1);
+        }
+
+        next_state.velocity[0] = state.velocity[0] + alpha * (desired_vx - state.velocity[0]);
+        next_state.velocity[1] = state.velocity[1] + alpha * (desired_vy - state.velocity[1]);
+        next_state.velocity[2] = 0;
+
+        next_state.position[0] = state.position[0] + next_state.velocity[0] * env.dt;
+        next_state.position[1] = state.position[1] + next_state.velocity[1] * env.dt;
+        next_state.position[2] = parameters.base_height;
+        next_state.yaw = state.yaw + yaw_rate * env.dt;
+
+        return env.dt;
+    }
+
+    template <typename DEVICE, typename SPEC, typename ACTION_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC::T reward(DEVICE& device, const rl::environments::raytracing_example::Environment<SPEC>& env, rl::environments::raytracing_example::Parameters<SPEC>& parameters, const rl::environments::raytracing_example::State<SPEC>& state, const Matrix<ACTION_SPEC>& action, const rl::environments::raytracing_example::State<SPEC>& next_state, RNG& rng) {
+        using T = typename SPEC::T;
+        const T v = std::sqrt(next_state.velocity[0] * next_state.velocity[0] + next_state.velocity[1] * next_state.velocity[1]);
+        const T control = std::abs(get(action, 0, 0)) + std::abs(get(action, 0, 1)) + static_cast<T>(0.1) * std::abs(get(action, 0, 2));
+        return v - static_cast<T>(0.05) * control;
+    }
+
+    template <typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT bool terminated(DEVICE& device, const rl::environments::raytracing_example::Environment<SPEC>& env, rl::environments::raytracing_example::Parameters<SPEC>& parameters, const rl::environments::raytracing_example::State<SPEC>& state, RNG& rng) {
+        return false;
+    }
+
+    template <typename DEVICE, typename SPEC, typename OBS_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void observe(DEVICE& device, rl::environments::raytracing_example::Environment<SPEC>& env, const rl::environments::raytracing_example::Parameters<SPEC>& parameters, const rl::environments::raytracing_example::State<SPEC>& state, const rl::environments::raytracing_example::ObservationRGB<SPEC>&, Matrix<OBS_SPEC>& observation, RNG& rng) {
+        static_assert(OBS_SPEC::ROWS == 1);
+        static_assert(OBS_SPEC::COLS == rl::environments::raytracing_example::ObservationRGB<SPEC>::DIM);
+
+        using T = typename SPEC::T;
+        const T position[3] = {
+            parameters.scene_translation[0] + state.position[0],
+            parameters.scene_translation[1] + state.position[1],
+            parameters.scene_translation[2] + state.position[2] + env.eye_height
+        };
+        const T look_at[3] = {
+            position[0] + env.look_ahead * std::cos(state.yaw),
+            position[1] + env.look_ahead * std::sin(state.yaw),
+            position[2]
+        };
+        const T up[3] = {0, 0, 1};
+        const auto camera = make_camera_data(position, look_at, up,
+            SPEC::FOV,
+            static_cast<T>(SPEC::CAM_WIDTH) / static_cast<T>(SPEC::CAM_HEIGHT));
+        set(device, env.camera_staging, camera, 0);
+        auto camera_staging_slot = view_range(device, env.camera_staging, 0, tensor::ViewSpec<0, 1>{});
+        auto camera_slot = view_range(device, cameras(device, *env.renderer), 0, tensor::ViewSpec<0, 1>{});
+        copy(device, env.renderer->device, camera_staging_slot, camera_slot);
+        render(device, *env.renderer);
+
+        std::vector<uint32_t> frame_staging((size_t)SPEC::CAM_WIDTH * SPEC::CAM_HEIGHT);
+        auto frame_camera_0 = view(device, frame_buffer(device, *env.renderer), 0);
+        Tensor<tensor::Specification<uint32_t, typename SPEC::TI, typename decltype(frame_camera_0)::SPEC::SHAPE>> frame_alias;
+        frame_alias._data = frame_staging.data();
+        copy(env.renderer->device, device, frame_camera_0, frame_alias);
+        for (typename SPEC::TI i = 0; i < static_cast<typename SPEC::TI>(SPEC::CAM_WIDTH * SPEC::CAM_HEIGHT); i++) {
+            set(observation, 0, i, static_cast<typename OBS_SPEC::T>(frame_staging[i]));
+        }
+    }
+
+    template <typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT rendering::raytracing::Camera<typename SPEC::T> make_camera_for_state(const rl::environments::raytracing_example::Environment<SPEC>& env, const rl::environments::raytracing_example::Parameters<SPEC>& parameters, const rl::environments::raytracing_example::State<SPEC>& state) {
+        using T = typename SPEC::T;
+        const T aspect = static_cast<T>(SPEC::CAM_WIDTH) / static_cast<T>(SPEC::CAM_HEIGHT);
+        const T cy = std::cos(state.yaw);
+        const T sy = std::sin(state.yaw);
+
+        const T position[3] = {
+            parameters.scene_translation[0] + state.position[0],
+            parameters.scene_translation[1] + state.position[1],
+            parameters.scene_translation[2] + state.position[2] + env.eye_height
+        };
+
+        const T look_at[3] = {
+            position[0] + env.look_ahead * cy,
+            position[1] + env.look_ahead * sy,
+            position[2]
+        };
+
+        const T up[3] = {0, 0, 1};
+
+        return make_camera_data(position, look_at, up, SPEC::FOV, aspect);
+    }
+
+    template <typename DEVICE, typename SPEC, typename PARAMETERS_SPEC, typename STATE_SPEC, typename ACTIONS_SPEC, typename NEXT_STATE_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void step_batch(DEVICE& device, const rl::environments::raytracing_example::Environment<SPEC>& env, Tensor<PARAMETERS_SPEC>& parameters, const Tensor<STATE_SPEC>& states, const Matrix<ACTIONS_SPEC>& actions, Tensor<NEXT_STATE_SPEC>& next_states, RNG& rng, typename SPEC::TI num_envs) {
+        static_assert(utils::typing::is_same_v<typename PARAMETERS_SPEC::T, rl::environments::raytracing_example::Parameters<SPEC>>);
+        static_assert(utils::typing::is_same_v<typename STATE_SPEC::T, rl::environments::raytracing_example::State<SPEC>>);
+        static_assert(utils::typing::is_same_v<typename NEXT_STATE_SPEC::T, rl::environments::raytracing_example::State<SPEC>>);
+        static_assert(ACTIONS_SPEC::ROWS == SPEC::NUM_ENVS);
+        static_assert(ACTIONS_SPEC::COLS == 3);
+
+        for (typename SPEC::TI env_i = 0; env_i < num_envs; env_i++) {
+            auto& parameters_env = get_ref(device, parameters, env_i);
+            auto& state_env = get_ref(device, states, env_i);
+            auto action_env = row(device, actions, env_i);
+            auto& next_state_env = get_ref(device, next_states, env_i);
+            step(device, env, parameters_env, state_env, action_env, next_state_env, rng);
+        }
+    }
+
+    template <typename DEVICE, typename SPEC, typename PARAMETERS_SPEC, typename STATE_SPEC, typename OUT_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void observe_batch(DEVICE& device, rl::environments::raytracing_example::Environment<SPEC>& env, const Tensor<PARAMETERS_SPEC>& parameters, const Tensor<STATE_SPEC>& states, typename SPEC::TI num_envs, Tensor<OUT_SPEC>& out_pixels) {
+        static_assert(utils::typing::is_same_v<typename PARAMETERS_SPEC::T, rl::environments::raytracing_example::Parameters<SPEC>>);
+        static_assert(utils::typing::is_same_v<typename STATE_SPEC::T, rl::environments::raytracing_example::State<SPEC>>);
+        static_assert(utils::typing::is_same_v<typename OUT_SPEC::T, uint32_t>);
+        static_assert(get<0>(typename OUT_SPEC::SHAPE{}) == SPEC::NUM_ENVS);
+        static_assert(get<1>(typename OUT_SPEC::SHAPE{}) == SPEC::CAM_HEIGHT);
+        static_assert(get<2>(typename OUT_SPEC::SHAPE{}) == SPEC::CAM_WIDTH);
+
+        if (env.renderer == nullptr || num_envs != SPEC::NUM_ENVS) {
+            return;
+        }
+
+        for (typename SPEC::TI env_i = 0; env_i < num_envs; env_i++) {
+            set(device, env.camera_staging, make_camera_for_state(env, get_ref(device, parameters, env_i), get_ref(device, states, env_i)), env_i);
+        }
+        copy(device, env.renderer->device, env.camera_staging, cameras(device, *env.renderer));
+        render(device, *env.renderer);
+        copy(env.renderer->device, device, frame_buffer(device, *env.renderer), out_pixels);
+    }
+}

hyperdrone/_vendor/rl-tools/src/rendering/raytracing/example/example.cpp

@@ -0,0 +1,273 @@
+#define RL_TOOLS_RENDERING_RAYTRACING_DISABLE_PROBE_RAYS 0
+
+#define RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_RGB 0
+#define RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_RGBD 1
+#define RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH 2
+
+#ifndef RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE
+#define RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_RGB
+#endif
+
+#include <rl_tools/operations/cpu_mux.h>
+
+#include "environment/environment.h"
+#include "environment/operations_cpu.h"
+#include <rl_tools/rendering/raytracing/save_cpu.h>
+
+
+#include <iostream>
+#include <chrono>
+#include <type_traits>
+#include <cmath>
+#include <vector>
+#include <sstream>
+#include <cstdio>
+#include <cstring>
+#include <array>
+
+namespace rlt = rl_tools;
+
+static constexpr bool OUTPUT_RGB = RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE != RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH;
+static constexpr bool OUTPUT_DEPTH = RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE != RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_RGB;
+
+int main(int argc, char** argv) {
+    bool output_video = false;
+    bool no_probe = false;
+    for (int i = 1; i < argc; i++) {
+        if (std::strcmp(argv[i], "--output-video") == 0) {
+            output_video = true;
+        }
+        else if (std::strcmp(argv[i], "--no-probe") == 0) {
+            no_probe = true;
+        }
+        else {
+            std::cerr << "Unknown argument: " << argv[i] << std::endl;
+            return 1;
+        }
+    }
+
+    using T = float;
+    using TI = typename rlt::devices::DEVICE_FACTORY<>::index_t;
+    static constexpr TI NUM_ENVS = 256;
+    constexpr T PI = static_cast<T>(3.14159265358979323846);
+    using SPEC = rlt::rl::environments::raytracing_example::Specification<T, TI, NUM_ENVS, 64, 64, 64, rlt::rendering::raytracing::Low, false, 1, false, 1, OUTPUT_RGB, OUTPUT_DEPTH>;
+
+    static_assert(std::is_standard_layout_v<rlt::rl::environments::raytracing_example::Parameters<SPEC>>);
+    static_assert(std::is_trivially_copyable_v<rlt::rl::environments::raytracing_example::Parameters<SPEC>>);
+    static_assert(std::is_standard_layout_v<rlt::rl::environments::raytracing_example::State<SPEC>>);
+    static_assert(std::is_trivially_copyable_v<rlt::rl::environments::raytracing_example::State<SPEC>>);
+
+    using DEVICE = rlt::devices::DEVICE_FACTORY<>;
+    using RNG = typename DEVICE::SPEC::RANDOM::ENGINE<>;
+
+    DEVICE device;
+    rlt::init(device);
+    RNG rng;
+    rlt::init(device, rng, 0);
+
+    rlt::rl::environments::raytracing_example::Environment<SPEC> env;
+    env.scene_path = "tests/data/ProcTHOR-Train-1.glb";
+
+    using PARAMETERS_SPEC = rlt::tensor::Specification<rlt::rl::environments::raytracing_example::Parameters<SPEC>, TI, rlt::tensor::Shape<TI, NUM_ENVS>>;
+    using STATE_SPEC = rlt::tensor::Specification<rlt::rl::environments::raytracing_example::State<SPEC>, TI, rlt::tensor::Shape<TI, NUM_ENVS>>;
+    using PIXELS_SPEC = rlt::tensor::Specification<uint32_t, TI, rlt::tensor::Shape<TI, NUM_ENVS, SPEC::CAM_HEIGHT, SPEC::CAM_WIDTH>>;
+
+    rlt::Tensor<PARAMETERS_SPEC> parameters;
+    rlt::Tensor<STATE_SPEC> states;
+    rlt::Tensor<PIXELS_SPEC> pixels;
+    std::vector<rlt::rl::environments::raytracing_example::State<SPEC>> base_states(NUM_ENVS);
+
+    rlt::malloc(device, parameters);
+    rlt::malloc(device, states);
+    rlt::malloc(device, pixels);
+
+    rlt::malloc(device, env);
+    rlt::init(device, env);
+
+    for (TI env_i = 0; env_i < NUM_ENVS; env_i++) {
+        rlt::rl::environments::raytracing_example::Parameters<SPEC> p;
+        rlt::sample_initial_parameters(device, env, p, rng);
+        rlt::rl::environments::raytracing_example::State<SPEC> s;
+        rlt::sample_initial_state(device, env, p, s, rng);
+        // const T angle = static_cast<T>(env_i) * static_cast<T>(0.01);
+        // s.position[0] = -0.937;
+        // s.position[1] = 1.690;
+        // s.position[2] = 8.410;
+        // s.velocity[0] = 0;
+        // s.velocity[1] = 0;
+        // s.velocity[2] = 0;
+        // s.yaw = angle + static_cast<T>(3.14159265358979323846) / static_cast<T>(2.0);
+        rlt::set(device, parameters, p, env_i);
+        rlt::set(device, states, s, env_i);
+        base_states[env_i] = s;
+    }
+
+
+    constexpr TI STEPS = 1024*4;
+#if RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE == RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH
+    if (output_video) {
+        std::cerr << "--output-video requires RGB or RGBD output mode." << std::endl;
+        return 1;
+    }
+#else
+    FILE* mp4_pipe = nullptr;
+    std::vector<uint32_t> per_camera_rgba;
+    std::vector<uint32_t> megaframe_rgba;
+    if (output_video) {
+        constexpr TI GRID_WIDTH = SPEC::RAYTRACING_SPEC::GRID_COLS * SPEC::CAM_WIDTH;
+        constexpr TI GRID_HEIGHT = SPEC::RAYTRACING_SPEC::GRID_ROWS * SPEC::CAM_HEIGHT;
+        per_camera_rgba.resize(static_cast<size_t>(NUM_ENVS) * static_cast<size_t>(SPEC::CAM_WIDTH) * static_cast<size_t>(SPEC::CAM_HEIGHT));
+        megaframe_rgba.resize(static_cast<size_t>(GRID_WIDTH) * static_cast<size_t>(GRID_HEIGHT));
+
+        const double frame_rate = env.dt > 0 ? (1.0 / static_cast<double>(env.dt)) : 30.0;
+        std::ostringstream ffmpeg_cmd;
+        ffmpeg_cmd
+            << "ffmpeg -y -f rawvideo -pixel_format rgba "
+            << "-video_size " << GRID_WIDTH << "x" << GRID_HEIGHT << " "
+            << "-framerate " << frame_rate << " -i - "
+            << "-an -c:v libx264 -pix_fmt yuv420p "
+            << "raytracing_example_megaframe.mp4";
+        mp4_pipe = popen(ffmpeg_cmd.str().c_str(), "w");
+        if (mp4_pipe == nullptr) {
+            std::cerr << "Failed to start ffmpeg." << std::endl;
+            return 1;
+        }
+    }
+#endif
+
+    // camera staging fills while the GPU renders; upload_cameras pushes it into the
+    // backend-native cameras tensor once the previous launch has synced
+    rlt::Tensor<typename decltype(env.renderer->cameras)::SPEC> camera_staging;
+    rlt::malloc(device, camera_staging);
+    auto compute_states_and_cameras = [&](TI step_i) {
+        for (TI env_i = 0; env_i < NUM_ENVS; env_i++) {
+            const T yaw_phase = static_cast<T>(0.02 * env_i + 0.08 * step_i);
+            auto s = base_states[env_i];
+            s.velocity[0] = 0;
+            s.velocity[1] = 0;
+            s.velocity[2] = 0;
+            s.yaw = static_cast<T>(0.5) * PI * (static_cast<T>(1) + std::sin(yaw_phase));
+            rlt::set(device, states, s, env_i);
+            rlt::set(device, camera_staging, rlt::make_camera_for_state(env, rlt::get_ref(device, parameters, env_i), s), env_i);
+        }
+    };
+    auto upload_cameras = [&]() {
+        rlt::copy(device, env.renderer->device, camera_staging, rlt::cameras(device, *env.renderer));
+    };
+
+    auto do_video_output = [&](TI step_i) -> int {
+#if RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE == RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH
+        (void)step_i;
+        return 0;
+#else
+        if (mp4_pipe == nullptr) return 0;
+        constexpr size_t pixel_count = (size_t)NUM_ENVS * SPEC::CAM_WIDTH * SPEC::CAM_HEIGHT;
+        static std::vector<uint32_t> frame_staging(pixel_count);
+        rlt::Tensor<typename decltype(env.renderer->frame_buffer)::SPEC> frame_alias;
+        frame_alias._data = frame_staging.data();
+        rlt::copy(env.renderer->device, device, rlt::frame_buffer(device, *env.renderer), frame_alias);
+        const uint32_t* per_camera_rgba_ptr = frame_staging.data();
+        for (TI camera_i = 0; camera_i < NUM_ENVS; camera_i++) {
+            const TI col = camera_i % SPEC::RAYTRACING_SPEC::GRID_COLS;
+            const TI row = camera_i / SPEC::RAYTRACING_SPEC::GRID_COLS;
+            const TI offset_x = col * SPEC::CAM_WIDTH;
+            const TI offset_y = row * SPEC::CAM_HEIGHT;
+            for (TI y = 0; y < SPEC::CAM_HEIGHT; y++) {
+                const uint32_t* src = per_camera_rgba_ptr + static_cast<size_t>(camera_i) * SPEC::CAM_WIDTH * SPEC::CAM_HEIGHT + static_cast<size_t>(y) * SPEC::CAM_WIDTH;
+                uint32_t* dst = megaframe_rgba.data() + static_cast<size_t>(offset_y + y) * (SPEC::RAYTRACING_SPEC::GRID_COLS * SPEC::CAM_WIDTH) + offset_x;
+                std::memcpy(dst, src, static_cast<size_t>(SPEC::CAM_WIDTH) * sizeof(uint32_t));
+            }
+        }
+        const size_t bytes = megaframe_rgba.size() * sizeof(uint32_t);
+        const size_t written = fwrite(megaframe_rgba.data(), 1, bytes, mp4_pipe);
+        if (written != bytes) {
+            std::cerr << "Failed writing frame " << step_i << " to ffmpeg." << std::endl;
+            pclose(mp4_pipe);
+            mp4_pipe = nullptr;
+            return 1;
+        }
+        return 0;
+#endif
+    };
+
+    rlt::synchronize(device, *env.renderer);
+    auto t0 = std::chrono::steady_clock::now();
+
+    // Pipelined loop: overlap CPU camera computation with GPU rendering
+    compute_states_and_cameras(0);
+    upload_cameras();
+
+    for (TI step_i = 0; step_i < STEPS; step_i++) {
+        if (no_probe) {
+#if RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE == RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH
+            rlt::render_launch(device, *env.renderer);
+#elif RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE == RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_RGBD
+            rlt::render_launch(device, *env.renderer);
+#else
+            rlt::render_launch(device, *env.renderer);
+#endif
+        }
+        else {
+            rlt::render_launch(device, *env.renderer);
+        }
+
+        // Overlap: compute next frame's cameras while GPU renders current frame
+        if (step_i + 1 < STEPS) {
+            compute_states_and_cameras(step_i + 1);
+        }
+
+        if (no_probe) {
+#if RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE == RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH
+            rlt::render_sync(device, *env.renderer);
+#elif RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE == RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_RGBD
+            rlt::render_sync(device, *env.renderer);
+#else
+            rlt::render_sync(device, *env.renderer);
+#endif
+        }
+        else {
+            rlt::render_sync(device, *env.renderer);
+        }
+
+        if (int err = do_video_output(step_i)) return err;
+
+        // Upload next frame's cameras (GPU is idle now, buffer is safe to overwrite)
+        if (step_i + 1 < STEPS) {
+            upload_cameras();
+        }
+    }
+    rlt::synchronize(device, *env.renderer);
+    auto t1 = std::chrono::steady_clock::now();
+    const double elapsed = std::chrono::duration<double>(t1 - t0).count();
+
+    const double fps = (static_cast<double>(NUM_ENVS) * STEPS) / elapsed;
+    std::cout << "raytracing_example: " << NUM_ENVS << " envs, " << STEPS << " batched observe steps" << std::endl;
+    std::cout << "elapsed: " << elapsed << " s, effective frame throughput: " << fps << " frames/s" << std::endl;
+
+#if RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE != RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH
+    if (mp4_pipe != nullptr) {
+        const int ffmpeg_status = pclose(mp4_pipe);
+        mp4_pipe = nullptr;
+        if (ffmpeg_status != 0) {
+            std::cerr << "ffmpeg exited with status " << ffmpeg_status << std::endl;
+            return 1;
+        }
+        std::cout << "video written: raytracing_example_megaframe.mp4" << std::endl;
+    }
+#endif
+
+#if RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE != RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_DEPTH
+    rlt::save_image(device, *env.renderer, "raytracing_example_grid.png");
+#endif
+#if RL_TOOLS_RENDERING_RAYTRACING_EXAMPLE_OUTPUT_MODE != RL_TOOLS_RENDERING_RAYTRACING_OUTPUT_RGB
+    rlt::save_depth_image(device, *env.renderer, "raytracing_example_depth.png");
+    rlt::save_depth(device, *env.renderer, "raytracing_example_depth.bin");
+#endif
+
+    rlt::free(device, camera_staging);
+    rlt::free(device, env);
+    rlt::free(device, pixels);
+    rlt::free(device, states);
+    rlt::free(device, parameters);
+    return 0;
+}