hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/src/rendering/raytracing/yaw_prediction/yaw_prediction_cuda.cu

@@ -0,0 +1,699 @@
+#define RL_TOOLS_OPERATIONS_CPU_MUX_INCLUDE_CUDA
+#include <rl_tools/operations/cpu_mux.h>
+
+// Adam instance operations must come before layer operations (see CLAUDE.md)
+#include <rl_tools/nn/optimizers/adam/instance/operations_generic.h>
+#include <rl_tools/nn/optimizers/adam/instance/operations_cuda.h>
+
+// Layer operations
+#include <rl_tools/nn/layers/dense/operations_generic.h>
+#include <rl_tools/nn/layers/conv2d/operations_generic.h>
+#include <rl_tools/nn/layers/avg_pool2d/operations_generic.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/operations_generic.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/operations_cuda.h>
+
+// Model operations
+#include <rl_tools/nn_models/sequential/operations_generic.h>
+
+// CUDA-specific operations
+#include <rl_tools/nn/operations_cuda.h>
+#include <rl_tools/nn_models/operations_generic.h>
+
+// Tensor operations
+#include <rl_tools/containers/tensor/operations_generic.h>
+#include <rl_tools/containers/tensor/operations_cpu.h>
+
+// Model definition + operations (must come before Adam optimizer operations)
+#include "model.h"
+#include "model_operations.h"
+#include "model_forward_cuda.h"
+
+// Adam optimizer operations (after model operations so ADL finds model's _reset_optimizer_state)
+#include <rl_tools/nn/optimizers/adam/operations_generic.h>
+#include <rl_tools/nn/optimizers/adam/operations_cuda.h>
+
+// Scene management (compiled separately as .cpp to avoid NVCC issues with environment code)
+
+#include "scene.h"
+
+#include "../example/environment/environment.h"
+#include "../example/environment/operations_cpu.h"
+
+// Experiment tracking (logging, tensorboard)
+#include <rl_tools/utils/extrack/extrack.h>
+#include <rl_tools/utils/extrack/operations_cpu.h>
+
+// Checkpointing (optional, requires HDF5)
+#ifdef RL_TOOLS_ENABLE_HDF5
+#include <rl_tools/persist/backends/hdf5/operations_cpu.h>
+#include <rl_tools/nn/layers/dense/persist.h>
+#include <rl_tools/nn/layers/conv2d/persist.h>
+#include <rl_tools/nn/layers/avg_pool2d/persist.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/persist.h>
+#include <rl_tools/nn_models/sequential/persist.h>
+#include <rl_tools/nn_models/parallel/persist.h>
+#include "model_persist.h"
+#endif
+
+#include <iostream>
+#include <vector>
+#include <cmath>
+#include <chrono>
+#include <cstring>
+#include <random>
+#include <algorithm>
+#include <filesystem>
+#include <csignal>
+#include <thread>
+#include <mutex>
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+
+namespace fs = std::filesystem;
+
+namespace rlt = rl_tools;
+
+static volatile std::sig_atomic_t signal_received = 0;
+static void signal_handler(int signal) {
+    signal_received = signal;
+}
+
+// ---- Configuration ----
+using T = __nv_bfloat16;
+using DEVICE_CPU = rlt::devices::DEVICE_FACTORY<>;
+using DEVICE_CUDA = rlt::devices::DEVICE_FACTORY_CUDA<>;
+using TI = DEVICE_CPU::index_t;
+using TI_CUDA = DEVICE_CUDA::index_t;
+
+namespace yp = rlt::rendering::raytracing::yaw_prediction;
+static constexpr TI_CUDA BATCH_SIZE = yp::SCENE_NUM_CAMERAS / 2;
+static constexpr TI_CUDA NUM_CAMERAS = yp::SCENE_NUM_CAMERAS;
+static constexpr TI CAM_WIDTH = yp::SCENE_CAM_WIDTH;
+static constexpr TI CAM_HEIGHT = yp::SCENE_CAM_HEIGHT;
+static constexpr TI NUM_ITERATIONS = 10000;
+
+struct TrainingConfig {
+    static constexpr float LEARNING_RATE = 1e-3f;
+    static constexpr float MAX_ANGLE = 3.14159265358979323846f / 6.0f; // 30 degrees
+    static constexpr float FOV_MIN = 17.188733853924695f;  // degrees, narrow
+    static constexpr float FOV_MAX = 68.75493541569878f;   // degrees, wide
+    static constexpr float HUBER_DELTA = 0.1f;
+    static constexpr float LOSS_WEIGHT_PX = 1.0f;
+    static constexpr float LOSS_WEIGHT_PY = 1.0f;
+    static constexpr float LOSS_WEIGHT_ROLL = 1.0f;
+    static constexpr TI LOG_INTERVAL = 100;
+    static constexpr TI VAL_INTERVAL = 1000;
+    static constexpr TI CHECKPOINT_INTERVAL = 100000;
+};
+
+using TYPE_POLICY = rlt::numeric_types::Policy<float,
+    rlt::numeric_types::UseCase<rlt::numeric_types::categories::Parameter, T>,
+    rlt::numeric_types::UseCase<rlt::numeric_types::categories::Activation, T>,
+    rlt::numeric_types::UseCase<rlt::numeric_types::categories::Gradient, T>,
+    rlt::numeric_types::UseCase<rlt::numeric_types::categories::MasterParameter, float>>;
+using T_ACTIVATION = TYPE_POLICY::GET<rlt::numeric_types::categories::Activation>;
+using T_GRADIENT = TYPE_POLICY::GET<rlt::numeric_types::categories::Gradient>;
+
+struct AdamParams : rlt::nn::optimizers::adam::DEFAULT_PARAMETERS_PYTORCH<TYPE_POLICY> {
+    static constexpr float ALPHA = TrainingConfig::LEARNING_RATE;
+};
+
+using OPTIMIZER_SPEC = rlt::nn::optimizers::adam::Specification<TYPE_POLICY, TI_CUDA, AdamParams>;
+using OPTIMIZER = rlt::nn::optimizers::Adam<OPTIMIZER_SPEC>;
+using GPU_CAPABILITY = rlt::nn::capability::Gradient<rlt::nn::parameters::Adam>;
+using GPU_MODEL = rlt::rendering::raytracing::yaw_prediction::MODEL<GPU_CAPABILITY, TYPE_POLICY, TI_CUDA, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH>;
+
+using CPU_TYPE_POLICY = rlt::numeric_types::Policy<float>;
+using CPU_CAPABILITY = rlt::nn::capability::Gradient<rlt::nn::parameters::Adam>;
+using CPU_MODEL = rlt::rendering::raytracing::yaw_prediction::MODEL<CPU_CAPABILITY, CPU_TYPE_POLICY, TI, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH>;
+using CPU_MODEL_INFERENCE = typename CPU_MODEL::template CHANGE_CAPABILITY<rlt::nn::capability::Forward<>>;
+
+using GPU_INPUT_SHAPE = rlt::tensor::Shape<TI_CUDA, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH, 3>;
+using GPU_INPUT_SPEC = rlt::tensor::Specification<T_ACTIVATION, TI_CUDA, GPU_INPUT_SHAPE>;
+using GPU_OUTPUT_SHAPE = typename GPU_MODEL::OUTPUT_SHAPE;
+using GPU_D_OUTPUT_SPEC = rlt::tensor::Specification<T_GRADIENT, TI_CUDA, GPU_OUTPUT_SHAPE>;
+
+// ---- CUDA kernels (loss only — model forward/backward moved to model_forward_cuda.h) ----
+
+// Element-wise Huber loss and gradient (operates on all BATCH_SIZE * OUTPUT_DIM elements)
+__global__ void huber_loss_gradient_kernel(
+    const T_ACTIVATION* __restrict__ predictions,
+    const float* __restrict__ targets,
+    T_GRADIENT* __restrict__ d_output,
+    float* __restrict__ losses,
+    int total_elements,
+    int output_dim,
+    float delta,
+    float weight_px,
+    float weight_py,
+    float weight_roll
+) {
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total_elements) return;
+
+    const int component = idx % output_dim;
+    const float weight = (component == 0) ? weight_px : (component == 1) ? weight_py : weight_roll;
+
+    const float pred = (float)predictions[idx];
+    const float tgt = targets[idx];
+    const float diff = pred - tgt;
+    const float abs_diff = abs(diff);
+
+    const float scale = weight / static_cast<float>(total_elements);
+    if (abs_diff <= delta) {
+        losses[idx] = weight * 0.5f * diff * diff;
+        d_output[idx] = (T_GRADIENT)(scale * diff);
+    } else {
+        losses[idx] = weight * delta * (abs_diff - 0.5f * delta);
+        d_output[idx] = (T_GRADIENT)(scale * delta * ((diff > 0.0f) - (diff < 0.0f)));
+    }
+}
+
+__global__ void reduce_loss_kernel(
+    const float* __restrict__ losses,
+    float* __restrict__ total_loss,
+    float* __restrict__ component_losses,
+    int batch_size,
+    int output_dim
+) {
+    float total = 0.0f;
+    for (int c = 0; c < output_dim; c++) {
+        float sum = 0.0f;
+        for (int b = 0; b < batch_size; b++) {
+            sum += losses[b * output_dim + c];
+        }
+        component_losses[c] = sum / static_cast<float>(batch_size);
+        total += sum;
+    }
+    *total_loss = total / static_cast<float>(batch_size * output_dim);
+}
+
+int main(int argc, char** argv) {
+    std::string scene_dir;
+    std::string single_scene;
+    TI num_scenes = 20;
+    TI num_scenes_per_batch = 8;
+    TI num_iterations = 100000;
+
+    TI num_val_scenes = 10;
+    TI num_load_threads = 8;
+
+    for (int i = 1; i < argc; i++) {
+        std::string arg = argv[i];
+        if (arg == "--scene-dir" && i + 1 < argc) {
+            scene_dir = argv[++i];
+        } else if (arg == "--scene" && i + 1 < argc) {
+            single_scene = argv[++i];
+        } else if (arg == "--num-scenes" && i + 1 < argc) {
+            num_scenes = std::atoi(argv[++i]);
+        } else if (arg == "--num-scenes-per-batch" && i + 1 < argc) {
+            num_scenes_per_batch = std::atoi(argv[++i]);
+        } else if (arg == "--num-iterations" && i + 1 < argc) {
+            num_iterations = std::atoi(argv[++i]);
+        } else if (arg == "--num-val-scenes" && i + 1 < argc) {
+            num_val_scenes = std::atoi(argv[++i]);
+        } else if (arg == "--num-load-threads" && i + 1 < argc) {
+            num_load_threads = std::atoi(argv[++i]);
+        } else {
+            std::cerr << "Usage: " << argv[0]
+                      << " [--scene-dir <dir> | --scene <path.glb>]"
+                      << " [--num-scenes N] [--num-scenes-per-batch M]"
+                      << " [--num-iterations N] [--num-val-scenes N]"
+                      << " [--num-load-threads N]" << std::endl;
+            return 1;
+        }
+    }
+
+    // Collect available scene paths (sorted, then split into train + val)
+    std::vector<std::string> all_scene_paths;
+    std::vector<std::string> val_scene_paths;
+    if (!scene_dir.empty()) {
+        std::vector<std::string> found_paths;
+        for (auto& entry : fs::directory_iterator(scene_dir)) {
+            if (entry.path().extension() == ".glb") {
+                found_paths.push_back(entry.path().string());
+            }
+        }
+        std::sort(found_paths.begin(), found_paths.end());
+        if (found_paths.empty()) {
+            std::cerr << "No .glb files found in " << scene_dir << std::endl;
+            return 1;
+        }
+        // First num_scenes for training, next num_val_scenes for validation
+        TI total_needed = num_scenes + num_val_scenes;
+        if (total_needed > found_paths.size()) {
+            std::cerr << "Need " << total_needed << " scenes (train=" << num_scenes << " + val=" << num_val_scenes << ") but only found " << found_paths.size() << std::endl;
+            return 1;
+        }
+        all_scene_paths.assign(found_paths.begin(), found_paths.begin() + num_scenes);
+        val_scene_paths.assign(found_paths.begin() + num_scenes, found_paths.begin() + total_needed);
+    } else if (!single_scene.empty()) {
+        all_scene_paths.push_back(single_scene);
+        num_scenes_per_batch = 1;
+        num_val_scenes = 0;
+    } else {
+        std::cerr << "Must specify --scene-dir or --scene" << std::endl;
+        return 1;
+    }
+
+    if (num_scenes_per_batch > all_scene_paths.size()) {
+        num_scenes_per_batch = all_scene_paths.size();
+    }
+
+    const TI samples_per_scene = BATCH_SIZE / num_scenes_per_batch;
+    if (samples_per_scene * num_scenes_per_batch != BATCH_SIZE) {
+        std::cerr << "BATCH_SIZE (" << BATCH_SIZE << ") must be divisible by num_scenes_per_batch (" << num_scenes_per_batch << ")" << std::endl;
+        return 1;
+    }
+
+    constexpr float PI = 3.14159265358979323846f;
+
+    std::cout << "Rotation prediction training (multi-scene)" << std::endl;
+    std::cout << "  Training scenes: " << all_scene_paths.size() << std::endl;
+    std::cout << "  Validation scenes: " << val_scene_paths.size() << std::endl;
+    std::cout << "  Scenes per batch: " << num_scenes_per_batch << std::endl;
+    std::cout << "  Samples per scene: " << samples_per_scene << std::endl;
+    std::cout << "  Batch size: " << BATCH_SIZE << std::endl;
+    std::cout << "  Image resolution: " << CAM_WIDTH << "x" << CAM_HEIGHT << std::endl;
+    std::cout << "  Max angle: " << TrainingConfig::MAX_ANGLE * 180.0f / PI << " degrees" << std::endl;
+    std::cout << "  Learning rate: " << TrainingConfig::LEARNING_RATE << std::endl;
+    std::cout << "  Num iterations: " << num_iterations << std::endl;
+
+    // ---- Load all scenes in parallel ----
+    struct LoadedScene {
+        yp::SceneHandle* handle = nullptr;
+        std::string path;
+    };
+    std::vector<LoadedScene> loaded_scenes(all_scene_paths.size());
+    std::vector<LoadedScene> val_scenes(val_scene_paths.size());
+
+    {
+        std::mutex cout_mutex;
+        TI total_scenes = all_scene_paths.size() + val_scene_paths.size();
+        TI actual_load_threads = std::min(num_load_threads, total_scenes);
+
+        // Build a flat list of load tasks (train first, then val)
+        struct LoadTask {
+            LoadedScene* target;
+            const std::string* path;
+            const char* label;
+            TI index;
+            TI count;
+        };
+        std::vector<LoadTask> tasks;
+        tasks.reserve(total_scenes);
+        for (TI s = 0; s < all_scene_paths.size(); s++) {
+            loaded_scenes[s].path = all_scene_paths[s];
+            tasks.push_back({&loaded_scenes[s], &loaded_scenes[s].path, "train", s, (TI)all_scene_paths.size()});
+        }
+        for (TI s = 0; s < val_scene_paths.size(); s++) {
+            val_scenes[s].path = val_scene_paths[s];
+            tasks.push_back({&val_scenes[s], &val_scenes[s].path, "val", s, (TI)val_scene_paths.size()});
+        }
+
+        std::atomic<TI> next_task{0};
+        std::vector<std::thread> load_threads;
+        load_threads.reserve(actual_load_threads);
+        for (TI t = 0; t < actual_load_threads; t++) {
+            load_threads.emplace_back([&]() {
+                while (true) {
+                    TI task_idx = next_task.fetch_add(1);
+                    if (task_idx >= tasks.size()) break;
+                    auto& task = tasks[task_idx];
+                    auto* handle = yp::create_scene(task.path->c_str());
+                    task.target->handle = handle;
+                    std::lock_guard<std::mutex> lock(cout_mutex);
+                    std::cout << "Loaded " << task.label << " scene " << task.index << "/" << task.count << ": "
+                              << fs::path(*task.path).filename().string()
+                              << " (" << yp::get_num_indoor_states(handle) << " indoor states)" << std::endl;
+                }
+            });
+        }
+        for (auto& t : load_threads) t.join();
+        std::cout << "All " << total_scenes << " scenes loaded (" << actual_load_threads << " threads)." << std::endl;
+    }
+
+    // RNG for selecting scenes per batch
+    std::mt19937 scene_rng(42);
+
+    // ---- Device setup ----
+    DEVICE_CPU device_cpu;
+    DEVICE_CUDA device_cuda;
+    rlt::init(device_cpu);
+    rlt::init(device_cuda);
+
+    DEVICE_CUDA::SPEC::RANDOM::ENGINE<> rng_cuda;
+    rlt::malloc(device_cuda, rng_cuda);
+    rlt::init(device_cuda, rng_cuda, 42);
+
+    // ---- Model setup ----
+    OPTIMIZER optimizer;
+    GPU_MODEL model;
+    typename GPU_MODEL::template Buffer<true> model_buffer;
+    rlt::malloc(device_cuda, optimizer);
+    rlt::malloc(device_cuda, model);
+    rlt::malloc(device_cuda, model_buffer);
+    rlt::init(device_cuda, optimizer);
+
+    // Init weights on CPU, then copy to GPU
+    CPU_MODEL model_cpu;
+    CPU_MODEL_INFERENCE model_cpu_inference;
+    rlt::malloc(device_cpu, model_cpu);
+    rlt::malloc(device_cpu, model_cpu_inference);
+    {
+        DEVICE_CPU::SPEC::RANDOM::ENGINE<> rng_cpu;
+        rlt::malloc(device_cpu, rng_cpu);
+        rlt::init(device_cpu, rng_cpu, 42);
+        rlt::init_weights(device_cpu, model_cpu, rng_cpu);
+        rlt::free(device_cpu, rng_cpu);
+    }
+    rlt::copy(device_cpu, device_cuda, model_cpu, model);
+    rlt::reset_optimizer_state(device_cuda, optimizer, model);
+    std::cout << "Model initialized" << std::endl;
+
+    // ---- GPU tensors ----
+    rlt::Tensor<GPU_INPUT_SPEC> gpu_input_a, gpu_input_b;
+    rlt::malloc(device_cuda, gpu_input_a);
+    rlt::malloc(device_cuda, gpu_input_b);
+    using GPU_D_INPUT_SPEC = rlt::tensor::Specification<T_GRADIENT, TI_CUDA, GPU_INPUT_SHAPE>;
+    rlt::Tensor<GPU_D_INPUT_SPEC> gpu_d_input_a, gpu_d_input_b;
+    rlt::malloc(device_cuda, gpu_d_input_a);
+    rlt::malloc(device_cuda, gpu_d_input_b);
+
+    rlt::Tensor<GPU_D_OUTPUT_SPEC> gpu_d_output;
+    rlt::malloc(device_cuda, gpu_d_output);
+    rlt::disable_dynamic_memory_allocation(device_cuda);
+
+    static constexpr TI OUTPUT_DIM = 3;
+    static constexpr TI TOTAL_OUTPUT_ELEMENTS = BATCH_SIZE * OUTPUT_DIM;
+
+    float* gpu_targets;
+    cudaMalloc(&gpu_targets, TOTAL_OUTPUT_ELEMENTS * sizeof(float));
+
+    float* gpu_losses;
+    cudaMalloc(&gpu_losses, TOTAL_OUTPUT_ELEMENTS * sizeof(float));
+
+    float* gpu_total_loss;
+    cudaMalloc(&gpu_total_loss, sizeof(float));
+
+    float* gpu_component_losses;
+    cudaMalloc(&gpu_component_losses, OUTPUT_DIM * sizeof(float));
+
+    // CPU-side buffers
+    std::vector<float> cpu_targets(TOTAL_OUTPUT_ELEMENTS);
+    std::vector<rlt::rendering::raytracing::Camera<float>> cameras(NUM_CAMERAS);
+
+    // ---- Extrack setup (experiment tracking, tensorboard logging) ----
+    rlt::utils::extrack::Config<TI> extrack_config;
+    rlt::utils::extrack::Paths extrack_paths;
+    extrack_config.name = "yaw-prediction";
+    extrack_config.population_variates = "cross-conv_channel-multiplier_resolution";
+    extrack_config.population_values = std::to_string(ABLATION_USE_CROSS_CONV) + "_" + std::to_string(ABLATION_CHANNEL_MULTIPLIER) + "_" + std::to_string(ABLATION_RESOLUTION);
+    rlt::init(device_cpu, extrack_config, extrack_paths, 0);
+
+    // ---- Signal handler for clean shutdown ----
+    std::signal(SIGINT, signal_handler);
+
+    // ---- Training loop ----
+    auto train_mode = rlt::Mode<rlt::mode::Default<>>{};
+    using CLOCK = std::chrono::high_resolution_clock;
+    CLOCK::time_point total_start;
+    bool throughput_timer_started = false;
+
+    // Indices for shuffled scene selection
+    std::vector<TI> scene_indices(loaded_scenes.size());
+    std::iota(scene_indices.begin(), scene_indices.end(), 0);
+
+    for (TI iteration = 0; iteration < num_iterations && !signal_received; iteration++) {
+        if (!throughput_timer_started && iteration == 1) {
+            total_start = CLOCK::now();
+            throughput_timer_started = true;
+        }
+        rlt::set_step(device_cpu, device_cpu.logger, iteration);
+
+        // ---- Select random scenes for this batch ----
+        std::shuffle(scene_indices.begin(), scene_indices.end(), scene_rng);
+
+        // ---- Sample and launch async renders for all scenes ----
+        for (TI s = 0; s < num_scenes_per_batch; s++) {
+            TI scene_idx = scene_indices[s];
+            const TI batch_offset = s * samples_per_scene;
+
+            yp::sample_camera_batch(
+                loaded_scenes[scene_idx].handle,
+                cameras.data(),
+                cpu_targets.data() + batch_offset * OUTPUT_DIM,
+                samples_per_scene,
+                TrainingConfig::MAX_ANGLE,
+                TrainingConfig::FOV_MIN,
+                TrainingConfig::FOV_MAX
+            );
+
+            yp::render_batch<true>(loaded_scenes[scene_idx].handle, cameras.data());
+        }
+
+        // ---- Sync all renders ----
+        for (TI s = 0; s < num_scenes_per_batch; s++) {
+            TI scene_idx = scene_indices[s];
+            auto& handle = loaded_scenes[scene_idx].handle;
+            rlt::render_sync(handle->device, *handle->env.renderer);
+        }
+
+        // ---- Convert rendered framebuffers to float tensors ----
+        for (TI s = 0; s < num_scenes_per_batch; s++) {
+            TI scene_idx = scene_indices[s];
+            const TI batch_offset = s * samples_per_scene;
+            uint32_t* device_fb = yp::get_framebuffer_device_ptr(loaded_scenes[scene_idx].handle);
+
+            yp::cuda_kernels::rgba_to_activation_kernel<<<samples_per_scene, 256, 0, device_cuda.stream>>>(
+                device_fb, gpu_input_a._data, CAM_WIDTH, CAM_HEIGHT, 0, batch_offset
+            );
+            yp::cuda_kernels::rgba_to_activation_kernel<<<samples_per_scene, 256, 0, device_cuda.stream>>>(
+                device_fb, gpu_input_b._data, CAM_WIDTH, CAM_HEIGHT, samples_per_scene, batch_offset
+            );
+        }
+        cudaMemcpyAsync(gpu_targets, cpu_targets.data(), TOTAL_OUTPUT_ELEMENTS * sizeof(float), cudaMemcpyHostToDevice, device_cuda.stream);
+
+        // ---- Forward ----
+        rlt::zero_gradient(device_cuda, model);
+        rlt::forward_cuda(device_cuda, model, gpu_input_a, gpu_input_b, model_buffer, rng_cuda, train_mode);
+        auto output_view = rlt::output(device_cuda, model);
+
+        // ---- Loss + gradient ----
+        {
+            constexpr int threads = 256;
+            constexpr int blocks = (TOTAL_OUTPUT_ELEMENTS + threads - 1) / threads;
+            huber_loss_gradient_kernel<<<blocks, threads, 0, device_cuda.stream>>>(
+                output_view._data, gpu_targets, gpu_d_output._data, gpu_losses, TOTAL_OUTPUT_ELEMENTS, OUTPUT_DIM,
+                TrainingConfig::HUBER_DELTA, TrainingConfig::LOSS_WEIGHT_PX, TrainingConfig::LOSS_WEIGHT_PY, TrainingConfig::LOSS_WEIGHT_ROLL
+            );
+        }
+
+        // ---- Backward ----
+        rlt::backward_full_cuda(device_cuda, model, gpu_input_a, gpu_input_b, gpu_d_output, gpu_d_input_a, gpu_d_input_b, model_buffer);
+
+        // ---- Optimizer step ----
+        rlt::step(device_cuda, optimizer, model);
+
+        // ---- Checkpointing ----
+#ifdef RL_TOOLS_ENABLE_HDF5
+        if (iteration % TrainingConfig::CHECKPOINT_INTERVAL == 0 || iteration == num_iterations - 1) {
+            cudaStreamSynchronize(device_cuda.stream);
+            rlt::copy(device_cuda, device_cpu, model, model_cpu);
+            rlt::copy(device_cpu, device_cpu, model_cpu, model_cpu_inference);
+            auto step_folder = rlt::get_step_folder(device_cpu, extrack_config, extrack_paths, iteration);
+            auto file = rl_tools::persist::backends::hdf5::File((step_folder / "checkpoint.h5").string(), rl_tools::persist::backends::hdf5::Mode::WRITE);
+            auto mg = rlt::create_group(device_cpu, file, "model");
+            rlt::save(device_cpu, model_cpu_inference, mg);
+            std::cout << "  Checkpoint: " << (step_folder / "checkpoint.h5").string() << std::endl;
+        }
+#endif
+
+        // ---- Logging ----
+        if (iteration % TrainingConfig::LOG_INTERVAL == 0 || iteration == num_iterations - 1) {
+            reduce_loss_kernel<<<1, 1, 0, device_cuda.stream>>>(
+                gpu_losses, gpu_total_loss, gpu_component_losses, BATCH_SIZE, OUTPUT_DIM
+            );
+            cudaStreamSynchronize(device_cuda.stream);
+
+            float loss_val;
+            cudaMemcpy(&loss_val, gpu_total_loss, sizeof(float), cudaMemcpyDeviceToHost);
+            float component_losses[OUTPUT_DIM];
+            cudaMemcpy(component_losses, gpu_component_losses, OUTPUT_DIM * sizeof(float), cudaMemcpyDeviceToHost);
+
+            std::vector<T_ACTIVATION> pred_buf_raw(TOTAL_OUTPUT_ELEMENTS);
+            cudaMemcpy(pred_buf_raw.data(), output_view._data, TOTAL_OUTPUT_ELEMENTS * sizeof(T_ACTIVATION), cudaMemcpyDeviceToHost);
+
+            float err_px = 0.0f, err_py = 0.0f, err_roll = 0.0f;
+            for (TI_CUDA i = 0; i < BATCH_SIZE; i++) {
+                err_px   += std::abs((float)pred_buf_raw[i * OUTPUT_DIM + 0] - cpu_targets[i * OUTPUT_DIM + 0]);
+                err_py   += std::abs((float)pred_buf_raw[i * OUTPUT_DIM + 1] - cpu_targets[i * OUTPUT_DIM + 1]);
+                err_roll += std::abs((float)pred_buf_raw[i * OUTPUT_DIM + 2] - cpu_targets[i * OUTPUT_DIM + 2]);
+            }
+            err_px /= BATCH_SIZE;
+            err_py /= BATCH_SIZE;
+            err_roll /= BATCH_SIZE;
+
+            auto now = CLOCK::now();
+            double elapsed_s = throughput_timer_started
+                ? std::chrono::duration<double>(now - total_start).count()
+                : 0.0;
+            double samples_per_s = elapsed_s > 0.0 && iteration > 0
+                ? static_cast<double>(iteration * BATCH_SIZE) / elapsed_s
+                : 0.0;
+
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/loss", loss_val);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/loss_px", component_losses[0]);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/loss_py", component_losses[1]);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/loss_roll", component_losses[2]);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/err_px", err_px);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/err_py", err_py);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/err_roll", err_roll);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "train/samples_per_s", samples_per_s);
+
+            std::cout << "[iter " << iteration << "/" << num_iterations << "]"
+                      << "  loss=" << loss_val
+                      << "  loss_px=" << component_losses[0]
+                      << "  loss_py=" << component_losses[1]
+                      << "  loss_roll=" << component_losses[2]
+                      << "  err_px=" << err_px
+                      << "  err_py=" << err_py
+                      << "  err_roll=" << err_roll
+                      << "  samples_per_s=" << samples_per_s
+                      << "  time=" << elapsed_s << "s"
+                      << std::endl;
+        }
+
+        // ---- Validation ----
+        if (val_scenes.size() > 0 && (iteration % TrainingConfig::VAL_INTERVAL == 0 || iteration == num_iterations - 1)) {
+            auto eval_mode = rlt::Mode<rlt::mode::Evaluation<>>{};
+            float val_loss_sum = 0.0f;
+            float val_loss_px_sum = 0.0f, val_loss_py_sum = 0.0f, val_loss_roll_sum = 0.0f;
+            float val_err_px_sum = 0.0f, val_err_py_sum = 0.0f, val_err_roll_sum = 0.0f;
+            TI val_total_samples = 0;
+
+            for (TI vs = 0; vs < val_scenes.size(); vs++) {
+                // Sample a full batch from this validation scene
+                yp::sample_camera_batch(
+                    val_scenes[vs].handle,
+                    cameras.data(),
+                    cpu_targets.data(),
+                    BATCH_SIZE,
+                    TrainingConfig::MAX_ANGLE,
+                    TrainingConfig::FOV_MIN,
+                    TrainingConfig::FOV_MAX
+                );
+                yp::render_batch<false>(val_scenes[vs].handle, cameras.data());
+                uint32_t* device_fb = yp::get_framebuffer_device_ptr(val_scenes[vs].handle);
+
+                yp::cuda_kernels::rgba_to_activation_kernel<<<BATCH_SIZE, 256, 0, device_cuda.stream>>>(
+                    device_fb, gpu_input_a._data, CAM_WIDTH, CAM_HEIGHT, 0, 0
+                );
+                yp::cuda_kernels::rgba_to_activation_kernel<<<BATCH_SIZE, 256, 0, device_cuda.stream>>>(
+                    device_fb, gpu_input_b._data, CAM_WIDTH, CAM_HEIGHT, BATCH_SIZE, 0
+                );
+
+                cudaMemcpyAsync(gpu_targets, cpu_targets.data(), TOTAL_OUTPUT_ELEMENTS * sizeof(float),
+                                cudaMemcpyHostToDevice, device_cuda.stream);
+
+                // Forward pass (evaluation mode)
+                rlt::forward_cuda(device_cuda, model, gpu_input_a, gpu_input_b, model_buffer, rng_cuda, eval_mode);
+                auto val_output_view = rlt::output(device_cuda, model);
+
+                // Compute loss
+                {
+                    constexpr int threads = 256;
+                    constexpr int blocks = (TOTAL_OUTPUT_ELEMENTS + threads - 1) / threads;
+                    huber_loss_gradient_kernel<<<blocks, threads, 0, device_cuda.stream>>>(
+                        val_output_view._data, gpu_targets, gpu_d_output._data, gpu_losses, TOTAL_OUTPUT_ELEMENTS, OUTPUT_DIM,
+                        TrainingConfig::HUBER_DELTA, TrainingConfig::LOSS_WEIGHT_PX, TrainingConfig::LOSS_WEIGHT_PY, TrainingConfig::LOSS_WEIGHT_ROLL
+                    );
+                }
+                reduce_loss_kernel<<<1, 1, 0, device_cuda.stream>>>(
+                    gpu_losses, gpu_total_loss, gpu_component_losses, BATCH_SIZE, OUTPUT_DIM
+                );
+                cudaStreamSynchronize(device_cuda.stream);
+
+                float scene_loss;
+                cudaMemcpy(&scene_loss, gpu_total_loss, sizeof(float), cudaMemcpyDeviceToHost);
+                float scene_component_losses[OUTPUT_DIM];
+                cudaMemcpy(scene_component_losses, gpu_component_losses, OUTPUT_DIM * sizeof(float), cudaMemcpyDeviceToHost);
+
+                std::vector<T_ACTIVATION> val_pred_buf_raw(TOTAL_OUTPUT_ELEMENTS);
+                cudaMemcpy(val_pred_buf_raw.data(), val_output_view._data, TOTAL_OUTPUT_ELEMENTS * sizeof(T_ACTIVATION), cudaMemcpyDeviceToHost);
+
+                float scene_err_px = 0.0f, scene_err_py = 0.0f, scene_err_roll = 0.0f;
+                for (TI_CUDA i = 0; i < BATCH_SIZE; i++) {
+                    scene_err_px   += std::abs((float)val_pred_buf_raw[i * OUTPUT_DIM + 0] - cpu_targets[i * OUTPUT_DIM + 0]);
+                    scene_err_py   += std::abs((float)val_pred_buf_raw[i * OUTPUT_DIM + 1] - cpu_targets[i * OUTPUT_DIM + 1]);
+                    scene_err_roll += std::abs((float)val_pred_buf_raw[i * OUTPUT_DIM + 2] - cpu_targets[i * OUTPUT_DIM + 2]);
+                }
+
+                val_loss_sum += scene_loss * TOTAL_OUTPUT_ELEMENTS;
+                val_loss_px_sum += scene_component_losses[0] * BATCH_SIZE;
+                val_loss_py_sum += scene_component_losses[1] * BATCH_SIZE;
+                val_loss_roll_sum += scene_component_losses[2] * BATCH_SIZE;
+                val_err_px_sum += scene_err_px;
+                val_err_py_sum += scene_err_py;
+                val_err_roll_sum += scene_err_roll;
+                val_total_samples += BATCH_SIZE;
+            }
+
+            float val_loss = val_loss_sum / (val_total_samples * OUTPUT_DIM);
+            float val_loss_px = val_loss_px_sum / val_total_samples;
+            float val_loss_py = val_loss_py_sum / val_total_samples;
+            float val_loss_roll = val_loss_roll_sum / val_total_samples;
+            float val_err_px = val_err_px_sum / val_total_samples;
+            float val_err_py = val_err_py_sum / val_total_samples;
+            float val_err_roll = val_err_roll_sum / val_total_samples;
+
+            rlt::add_scalar(device_cpu, device_cpu.logger, "val/loss", val_loss);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "val/loss_px", val_loss_px);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "val/loss_py", val_loss_py);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "val/loss_roll", val_loss_roll);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "val/err_px", val_err_px);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "val/err_py", val_err_py);
+            rlt::add_scalar(device_cpu, device_cpu.logger, "val/err_roll", val_err_roll);
+
+            std::cout << "[iter " << iteration << "] VAL loss=" << val_loss
+                      << "  loss_px=" << val_loss_px
+                      << "  loss_py=" << val_loss_py
+                      << "  loss_roll=" << val_loss_roll
+                      << "  err_px=" << val_err_px
+                      << "  err_py=" << val_err_py
+                      << "  err_roll=" << val_err_roll << std::endl;
+        }
+    }
+
+    if (signal_received) {
+        std::cout << "Training interrupted by signal " << signal_received << "." << std::endl;
+    } else {
+        std::cout << "Training complete." << std::endl;
+    }
+
+    // ---- Cleanup ----
+    cudaFree(gpu_targets);
+    cudaFree(gpu_losses);
+    cudaFree(gpu_total_loss);
+    cudaFree(gpu_component_losses);
+    rlt::free(device_cuda, gpu_input_a);
+    rlt::free(device_cuda, gpu_input_b);
+    rlt::free(device_cuda, gpu_d_input_a);
+    rlt::free(device_cuda, gpu_d_input_b);
+    rlt::free(device_cuda, gpu_d_output);
+    rlt::free(device_cuda, model_buffer);
+    rlt::free(device_cuda, model);
+    rlt::free(device_cuda, optimizer);
+    rlt::free(device_cuda, rng_cuda);
+    rlt::free(device_cpu, model_cpu);
+    rlt::free(device_cpu, model_cpu_inference);
+    for (auto& s : loaded_scenes) {
+        yp::destroy_scene(s.handle);
+    }
+    for (auto& s : val_scenes) {
+        yp::destroy_scene(s.handle);
+    }
+    rlt::free(device_cpu, device_cpu.logger);
+
+    return 0;
+}