"""Two-launch-per-iteration CUDA implementation of grid-foraging PPO.

The rollout kernel is a cooperative persistent kernel: each CUDA thread owns
one environment for all 32 steps, including policy inference, sampling, the
environment transition, GAE, normalization, and episode-return reduction.
The update kernel is another cooperative persistent kernel.  It performs all
four epochs by four minibatches internally, with grid barriers between the
forward/backward reduction and each Adam update.  Consequently launch count is
two per PPO iteration and never scales with rollout steps or minibatches.
"""
from __future__ import annotations

import os

os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")

from torch.utils.cpp_extension import load_inline

_CPP = r"""
#include <torch/extension.h>

torch::Tensor ppo_train_cuda(int64_t iterations, int64_t seed);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("train", &ppo_train_cuda, "Persistent CUDA PPO trainer");
}
"""


_CUDA = r"""
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <stdexcept>

namespace cg = cooperative_groups;

constexpr int N = 4096;
constexpr int T = 32;
constexpr int B = N * T;
constexpr int MB = B / 4;
constexpr int H = 64;
constexpr int P = 645;

constexpr int W1 = 0;       // 64 x 4
constexpr int B1 = 256;     // 64
constexpr int WPI = 320;    // 4 x 64
constexpr int BPI = 576;    // 4
constexpr int WV = 580;     // 64
constexpr int BV = 644;     // 1

constexpr float GAMMA = 0.99f;
constexpr float GAMLAM = 0.9405f;
constexpr float CLIP_LO = 0.8f;
constexpr float CLIP_HI = 1.2f;
constexpr float ENT = 0.01f;
constexpr float LR = 0.003f;

static inline void cuda_check(cudaError_t err, const char* where) {
    if (err != cudaSuccess) {
        throw std::runtime_error(std::string(where) + ": " + cudaGetErrorString(err));
    }
}

__host__ __device__ __forceinline__ uint64_t mix64(uint64_t x) {
    x += 0x9e3779b97f4a7c15ULL;
    x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ULL;
    x = (x ^ (x >> 27)) * 0x94d049bb133111ebULL;
    return x ^ (x >> 31);
}

__device__ __forceinline__ uint32_t next_u32(uint64_t& s) {
    // xorshift64*: independent stream per environment.
    s ^= s >> 12;
    s ^= s << 25;
    s ^= s >> 27;
    return static_cast<uint32_t>((s * 0x2545f4914f6cdd1dULL) >> 32);
}

__device__ __forceinline__ float uniform01(uint64_t& s) {
    return (static_cast<float>(next_u32(s)) + 0.5f) * 0x1.0p-32f;
}

__device__ __forceinline__ int rand_cell(uint64_t& s) {
    // Multiply-high maps all 2^32 RNG values to [0, 11) without modulo bias.
    return static_cast<int>((static_cast<uint64_t>(next_u32(s)) * 11ULL) >> 32);
}

__device__ __forceinline__ float fast_tanh(float x) {
    return tanhf(x);
}

__global__ void init_kernel(float* w, float* m, float* v, uint64_t* rng,
                            uint64_t seed) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < P) {
        m[i] = 0.0f;
        v[i] = 0.0f;
        bool bias = (i >= B1 && i < WPI) || (i >= BPI && i < WV) || i == BV;
        if (bias) {
            w[i] = 0.0f;
        } else {
            float bound = i < B1 ? 0.5f : 0.125f;
            uint64_t z = mix64(seed ^ (0xd1b54a32d192ed03ULL * (uint64_t)(i + 1)));
            float u = (static_cast<float>(static_cast<uint32_t>(z >> 32)) + 0.5f)
                    * 0x1.0p-32f;
            w[i] = (2.0f * u - 1.0f) * bound;
        }
    }
    if (i < N) {
        uint64_t s = mix64(seed ^ (0x9e3779b97f4a7c15ULL * (uint64_t)(i + 1)));
        rng[i] = s ? s : 0x6a09e667f3bcc909ULL;
    }
}

__device__ __forceinline__ void policy_forward(
    const float* __restrict__ w, const float* __restrict__ obs,
    float* __restrict__ hidden, float logits[4], float& value) {
    #pragma unroll
    for (int j = 0; j < H; ++j) {
        float z = w[B1 + j];
        z = fmaf(w[W1 + j * 4 + 0], obs[0], z);
        z = fmaf(w[W1 + j * 4 + 1], obs[1], z);
        z = fmaf(w[W1 + j * 4 + 2], obs[2], z);
        z = fmaf(w[W1 + j * 4 + 3], obs[3], z);
        hidden[j] = fast_tanh(z);
    }
    #pragma unroll
    for (int a = 0; a < 4; ++a) {
        float z = w[BPI + a];
        #pragma unroll 4
        for (int j = 0; j < H; ++j) z = fmaf(w[WPI + a * H + j], hidden[j], z);
        logits[a] = z;
    }
    value = w[BV];
    #pragma unroll 4
    for (int j = 0; j < H; ++j) value = fmaf(w[WV + j], hidden[j], value);
}

__global__ void rollout_kernel(
    const float* __restrict__ w,
    uint64_t* __restrict__ rng,
    float* __restrict__ obs_buf,
    uint8_t* __restrict__ act_buf,
    float* __restrict__ old_logp,
    float* __restrict__ val_buf,
    uint8_t* __restrict__ rew_buf,
    float* __restrict__ adv_buf,
    float* __restrict__ ret_buf,
    float* __restrict__ rollout_hidden,
    float* __restrict__ stats,
    float* __restrict__ curve,
    int iteration) {
    cg::grid_group grid = cg::this_grid();
    int e = blockIdx.x * blockDim.x + threadIdx.x;
    extern __shared__ float reduce[];
    float* rollout_w = reduce + 3 * blockDim.x;

    if (e == 0) {
        stats[0] = 0.0f; // sum advantage
        stats[1] = 0.0f; // sum advantage squared
        stats[2] = 0.0f; // sum episode return
    }
    grid.sync();
    for (int p = threadIdx.x; p < P; p += blockDim.x) rollout_w[p] = w[p];
    __syncthreads();

    float adv_sum = 0.0f;
    float adv_sq = 0.0f;
    float episode_return = 0.0f;
    if (e < N) {
        uint64_t state = rng[e];
        int ax = rand_cell(state);
        int ay = rand_cell(state);
        int fx = rand_cell(state);
        int fy = rand_cell(state);
        float* hidden = rollout_hidden + e * H;

        #pragma unroll 1
        for (int t = 0; t < T; ++t) {
            int s = t * N + e;
            float o[4] = {
                static_cast<float>(fx - ax) * (1.0f / 11.0f),
                static_cast<float>(fy - ay) * (1.0f / 11.0f),
                static_cast<float>(ax) * 0.1f,
                static_cast<float>(ay) * 0.1f
            };
            #pragma unroll
            for (int k = 0; k < 4; ++k) obs_buf[s * 4 + k] = o[k];

            float logits[4], value;
            policy_forward(rollout_w, o, hidden, logits, value);
            float mx = fmaxf(fmaxf(logits[0], logits[1]), fmaxf(logits[2], logits[3]));
            float p0 = __expf(logits[0] - mx);
            float p1 = __expf(logits[1] - mx);
            float p2 = __expf(logits[2] - mx);
            float p3 = __expf(logits[3] - mx);
            float inv = 1.0f / (p0 + p1 + p2 + p3);
            p0 *= inv; p1 *= inv; p2 *= inv; p3 *= inv;
            float u = uniform01(state);
            int action = u < p0 ? 0 : (u < p0 + p1 ? 1 : (u < p0 + p1 + p2 ? 2 : 3));
            float pa = action == 0 ? p0 : (action == 1 ? p1 : (action == 2 ? p2 : p3));
            act_buf[s] = static_cast<uint8_t>(action);
            old_logp[s] = __logf(fmaxf(pa, 1.0e-20f));
            val_buf[s] = value;

            if (action == 0) ay = max(0, ay - 1);
            else if (action == 1) ay = min(10, ay + 1);
            else if (action == 2) ax = max(0, ax - 1);
            else ax = min(10, ax + 1);
            bool caught = ax == fx && ay == fy;
            rew_buf[s] = static_cast<uint8_t>(caught);
            episode_return += caught ? 1.0f : 0.0f;
            if (caught) {
                fx = rand_cell(state);
                fy = rand_cell(state);
            }
        }

        float gae = 0.0f;
        #pragma unroll
        for (int t = T - 1; t >= 0; --t) {
            int s = t * N + e;
            float delta = static_cast<float>(rew_buf[s]) - val_buf[s];
            if (t != T - 1) delta += GAMMA * val_buf[s + N];
            gae = delta + (t == T - 1 ? 0.0f : GAMLAM * gae);
            adv_buf[s] = gae;
            ret_buf[s] = gae + val_buf[s];
            adv_sum += gae;
            adv_sq = fmaf(gae, gae, adv_sq);
        }
        rng[e] = state;
    }

    float* r0 = reduce;
    float* r1 = reduce + blockDim.x;
    float* r2 = reduce + 2 * blockDim.x;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    #pragma unroll
    for (int d = 16; d; d >>= 1) {
        adv_sum += __shfl_down_sync(0xffffffffU, adv_sum, d);
        adv_sq += __shfl_down_sync(0xffffffffU, adv_sq, d);
        episode_return += __shfl_down_sync(0xffffffffU, episode_return, d);
    }
    if (lane == 0) {
        r0[warp] = adv_sum;
        r1[warp] = adv_sq;
        r2[warp] = episode_return;
    }
    __syncthreads();
    if (warp == 0) {
        int nw = blockDim.x >> 5;
        float s0 = lane < nw ? r0[lane] : 0.0f;
        float s1 = lane < nw ? r1[lane] : 0.0f;
        float s2 = lane < nw ? r2[lane] : 0.0f;
        #pragma unroll
        for (int d = 16; d; d >>= 1) {
            s0 += __shfl_down_sync(0xffffffffU, s0, d);
            s1 += __shfl_down_sync(0xffffffffU, s1, d);
            s2 += __shfl_down_sync(0xffffffffU, s2, d);
        }
        if (lane == 0) {
            stats[blockIdx.x] = s0;
            stats[32 + blockIdx.x] = s1;
            stats[64 + blockIdx.x] = s2;
        }
    }
    grid.sync();
    if (e == 0) {
        float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f;
        #pragma unroll
        for (int b = 0; b < 32; ++b) {
            s0 += stats[b];
            s1 += stats[32 + b];
            s2 += stats[64 + b];
        }
        stats[0] = s0;
        stats[1] = s1;
        stats[2] = s2;
    }
    grid.sync();

    float mean = stats[0] / static_cast<float>(B);
    float centered_ss = fmaxf(0.0f, stats[1] - stats[0] * mean);
    float inv_std = rsqrtf(centered_ss / static_cast<float>(B - 1) + 1.0e-8f);
    if (e < N) {
        #pragma unroll
        for (int t = 0; t < T; ++t) {
            int s = t * N + e;
            adv_buf[s] = (adv_buf[s] - mean) * inv_std;
        }
    }
    if (e == 0) curve[iteration] = stats[2] / static_cast<float>(N);
}

__device__ __forceinline__ uint32_t hash32(uint32_t x) {
    x ^= x >> 16;
    x *= 0x7feb352dU;
    x ^= x >> 15;
    x *= 0x846ca68bU;
    return x ^ (x >> 16);
}

__device__ __forceinline__ int permuted_index(int pos, uint32_t a, uint32_t b) {
    // Shuffle the 4096 warp-sized sample groups while retaining coalescing
    // inside each group.  Odd ``a`` makes this a bijection over the batch.
    uint32_t group = static_cast<uint32_t>(pos) >> 5;
    uint32_t lane = static_cast<uint32_t>(pos) & 31U;
    uint32_t shuffled = (a * group + b) & ((B >> 5) - 1);
    return static_cast<int>((shuffled << 5) | lane);
}

__global__ void update_kernel(
    float* __restrict__ w,
    float* __restrict__ adam_m,
    float* __restrict__ adam_v,
    const float* __restrict__ obs,
    const uint8_t* __restrict__ actions,
    const float* __restrict__ old_logp,
    const float* __restrict__ advantages,
    const float* __restrict__ returns,
    float* __restrict__ hidden,
    float* __restrict__ dpre,
    float* __restrict__ dlogits,
    float* __restrict__ dvalue,
    float* __restrict__ grad,
    float* __restrict__ stats,
    uint64_t seed,
    int iteration) {
    cg::grid_group grid = cg::this_grid();
    int gt = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    extern __shared__ float red[];

    #pragma unroll 1
    for (int epoch = 0; epoch < 4; ++epoch) {
        uint32_t key = hash32(static_cast<uint32_t>(seed) ^
                              (0x9e3779b9U * static_cast<uint32_t>(iteration * 4 + epoch + 1)));
        uint32_t pa = (hash32(key ^ 0xa511e9b3U) | 1U);
        uint32_t pb = hash32(key ^ 0x63d83595U) & (B - 1);

        #pragma unroll 1
        for (int mini = 0; mini < 4; ++mini) {
            float* sw = red + blockDim.x;
            for (int p = threadIdx.x; p < P; p += blockDim.x) sw[p] = w[p];
            __syncthreads();
            // Per-sample forward and loss derivatives for this minibatch.
            for (int q = gt; q < MB; q += stride) {
                int idx = permuted_index(mini * MB + q, pa, pb);
                const float* o = obs + idx * 4;
                float* hh = hidden + q * H;
                #pragma unroll
                for (int j = 0; j < H; ++j) {
                    float z = sw[B1 + j];
                    z = fmaf(sw[W1 + j * 4 + 0], o[0], z);
                    z = fmaf(sw[W1 + j * 4 + 1], o[1], z);
                    z = fmaf(sw[W1 + j * 4 + 2], o[2], z);
                    z = fmaf(sw[W1 + j * 4 + 3], o[3], z);
                    hh[j] = fast_tanh(z);
                }

                float logits[4];
                #pragma unroll
                for (int a = 0; a < 4; ++a) {
                    float z = sw[BPI + a];
                    #pragma unroll 4
                    for (int j = 0; j < H; ++j) z = fmaf(sw[WPI + a * H + j], hh[j], z);
                    logits[a] = z;
                }
                float value = sw[BV];
                #pragma unroll 4
                for (int j = 0; j < H; ++j) value = fmaf(sw[WV + j], hh[j], value);

                float mx = fmaxf(fmaxf(logits[0], logits[1]), fmaxf(logits[2], logits[3]));
                float ex[4];
                ex[0] = __expf(logits[0] - mx);
                ex[1] = __expf(logits[1] - mx);
                ex[2] = __expf(logits[2] - mx);
                ex[3] = __expf(logits[3] - mx);
                float invz = 1.0f / (ex[0] + ex[1] + ex[2] + ex[3]);
                float prob[4], logp[4];
                float sum_p_logp = 0.0f;
                #pragma unroll
                for (int a = 0; a < 4; ++a) {
                    prob[a] = ex[a] * invz;
                    logp[a] = __logf(fmaxf(prob[a], 1.0e-20f));
                    sum_p_logp = fmaf(prob[a], logp[a], sum_p_logp);
                }
                int action = static_cast<int>(actions[idx]);
                float ratio = __expf(logp[action] - old_logp[idx]);
                float av = advantages[idx];
                bool active = av >= 0.0f ? ratio <= CLIP_HI : ratio >= CLIP_LO;
                float dlp = active ? -av * ratio : 0.0f;
                #pragma unroll
                for (int a = 0; a < 4; ++a) {
                    float g = dlp * ((a == action ? 1.0f : 0.0f) - prob[a]);
                    g += ENT * prob[a] * (logp[a] - sum_p_logp);
                    dlogits[q * 4 + a] = g;
                }
                float dv = value - returns[idx];
                dvalue[q] = dv;
                float* dp = dpre + q * H;
                #pragma unroll
                for (int j = 0; j < H; ++j) {
                    float dh = dv * sw[WV + j];
                    #pragma unroll
                    for (int a = 0; a < 4; ++a)
                        dh = fmaf(dlogits[q * 4 + a], sw[WPI + a * H + j], dh);
                    dp[j] = dh * (1.0f - hh[j] * hh[j]);
                }
            }
            grid.sync();

            // One CTA reduces each parameter, cycling over all 645 parameters.
            for (int p = blockIdx.x; p < P; p += gridDim.x) {
                float sum = 0.0f;
                for (int q = threadIdx.x; q < MB; q += blockDim.x) {
                    float term;
                    if (p < B1) {
                        int j = p >> 2;
                        int k = p & 3;
                        int idx = permuted_index(mini * MB + q, pa, pb);
                        term = dpre[q * H + j] * obs[idx * 4 + k];
                    } else if (p < WPI) {
                        term = dpre[q * H + (p - B1)];
                    } else if (p < BPI) {
                        int x = p - WPI;
                        int a = x / H;
                        int j = x - a * H;
                        term = dlogits[q * 4 + a] * hidden[q * H + j];
                    } else if (p < WV) {
                        term = dlogits[q * 4 + (p - BPI)];
                    } else if (p < BV) {
                        term = dvalue[q] * hidden[q * H + (p - WV)];
                    } else {
                        term = dvalue[q];
                    }
                    sum += term;
                }
                int lane = threadIdx.x & 31;
                int warp = threadIdx.x >> 5;
                #pragma unroll
                for (int d = 16; d; d >>= 1)
                    sum += __shfl_down_sync(0xffffffffU, sum, d);
                if (lane == 0) red[warp] = sum;
                __syncthreads();
                if (warp == 0) {
                    float total = lane < (blockDim.x >> 5) ? red[lane] : 0.0f;
                    #pragma unroll
                    for (int d = 16; d; d >>= 1)
                        total += __shfl_down_sync(0xffffffffU, total, d);
                    if (lane == 0) grad[p] = total * (1.0f / static_cast<float>(MB));
                }
                __syncthreads();
            }
            grid.sync();

            // Global L2 gradient clipping, matching clip_grad_norm_(..., 0.5).
            if (blockIdx.x == 0) {
                float ss = 0.0f;
                for (int p = threadIdx.x; p < P; p += blockDim.x) ss = fmaf(grad[p], grad[p], ss);
                int lane = threadIdx.x & 31;
                int warp = threadIdx.x >> 5;
                #pragma unroll
                for (int d = 16; d; d >>= 1)
                    ss += __shfl_down_sync(0xffffffffU, ss, d);
                if (lane == 0) red[warp] = ss;
                __syncthreads();
                if (warp == 0) {
                    float total = lane < (blockDim.x >> 5) ? red[lane] : 0.0f;
                    #pragma unroll
                    for (int d = 16; d; d >>= 1)
                        total += __shfl_down_sync(0xffffffffU, total, d);
                    if (lane == 0) {
                        float norm = sqrtf(total);
                        stats[0] = fminf(1.0f, 0.5f / (norm + 1.0e-6f));
                        int step = (iteration * 16) + epoch * 4 + mini + 1;
                        stats[1] = 1.0f / (1.0f - powf(0.9f, static_cast<float>(step)));
                        stats[2] = 1.0f / (1.0f - powf(0.999f, static_cast<float>(step)));
                    }
                }
            }
            grid.sync();

            for (int p = gt; p < P; p += stride) {
                float g = grad[p] * stats[0];
                float mm = fmaf(0.9f, adam_m[p], 0.1f * g);
                float vv = fmaf(0.999f, adam_v[p], 0.001f * g * g);
                adam_m[p] = mm;
                adam_v[p] = vv;
                float mhat = mm * stats[1];
                float vhat = vv * stats[2];
                w[p] -= LR * mhat / (sqrtf(vhat) + 1.0e-8f);
            }
            grid.sync();
        }
    }
}

struct Workspace {
    bool ready = false;
    float *w, *m, *v, *obs, *old_logp, *values, *adv, *ret;
    float *roll_hidden, *upd_hidden, *dpre, *dlogits, *dvalue, *grad, *stats, *curve;
    uint8_t *actions, *rewards;
    uint64_t *rng;

    void allocate() {
        if (ready) return;
        cuda_check(cudaMalloc(&w, P * sizeof(float)), "cudaMalloc w");
        cuda_check(cudaMalloc(&m, P * sizeof(float)), "cudaMalloc m");
        cuda_check(cudaMalloc(&v, P * sizeof(float)), "cudaMalloc v");
        cuda_check(cudaMalloc(&rng, N * sizeof(uint64_t)), "cudaMalloc rng");
        cuda_check(cudaMalloc(&obs, B * 4 * sizeof(float)), "cudaMalloc obs");
        cuda_check(cudaMalloc(&actions, B * sizeof(uint8_t)), "cudaMalloc actions");
        cuda_check(cudaMalloc(&old_logp, B * sizeof(float)), "cudaMalloc old_logp");
        cuda_check(cudaMalloc(&values, B * sizeof(float)), "cudaMalloc values");
        cuda_check(cudaMalloc(&rewards, B * sizeof(uint8_t)), "cudaMalloc rewards");
        cuda_check(cudaMalloc(&adv, B * sizeof(float)), "cudaMalloc adv");
        cuda_check(cudaMalloc(&ret, B * sizeof(float)), "cudaMalloc ret");
        cuda_check(cudaMalloc(&roll_hidden, N * H * sizeof(float)), "cudaMalloc roll_hidden");
        cuda_check(cudaMalloc(&upd_hidden, MB * H * sizeof(float)), "cudaMalloc upd_hidden");
        cuda_check(cudaMalloc(&dpre, MB * H * sizeof(float)), "cudaMalloc dpre");
        cuda_check(cudaMalloc(&dlogits, MB * 4 * sizeof(float)), "cudaMalloc dlogits");
        cuda_check(cudaMalloc(&dvalue, MB * sizeof(float)), "cudaMalloc dvalue");
        cuda_check(cudaMalloc(&grad, P * sizeof(float)), "cudaMalloc grad");
        cuda_check(cudaMalloc(&stats, 3 * 32 * sizeof(float)), "cudaMalloc stats");
        cuda_check(cudaMalloc(&curve, 40 * sizeof(float)), "cudaMalloc curve");
        ready = true;
    }
};

static Workspace ws;

torch::Tensor ppo_train_cuda(int64_t iterations, int64_t seed) {
    cuda_check(cudaSetDevice(0), "cudaSetDevice");
    if (iterations < 1) iterations = 1;
    if (iterations > 40) throw std::runtime_error("at most 40 PPO iterations are supported");
    ws.allocate();

    init_kernel<<<16, 256>>>(ws.w, ws.m, ws.v, ws.rng, static_cast<uint64_t>(seed));
    cuda_check(cudaGetLastError(), "init_kernel");

    int sm_count = 0;
    cuda_check(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0), "SM count");
    int active = 0;
    cuda_check(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &active, update_kernel, 128, (128 + P) * sizeof(float)), "update occupancy");
    int update_blocks = active * sm_count;
    if (update_blocks < 1) throw std::runtime_error("update kernel has zero occupancy");

    for (int it = 0; it < iterations; ++it) {
        void* rargs[] = {
            &ws.w, &ws.rng, &ws.obs, &ws.actions, &ws.old_logp, &ws.values,
            &ws.rewards, &ws.adv, &ws.ret, &ws.roll_hidden, &ws.stats, &ws.curve, &it
        };
        cuda_check(cudaLaunchCooperativeKernel(
            reinterpret_cast<void*>(rollout_kernel), dim3(32), dim3(128), rargs,
            (3 * 128 + P) * sizeof(float), 0), "rollout launch");

        void* uargs[] = {
            &ws.w, &ws.m, &ws.v, &ws.obs, &ws.actions, &ws.old_logp, &ws.adv,
            &ws.ret, &ws.upd_hidden, &ws.dpre, &ws.dlogits, &ws.dvalue, &ws.grad,
            &ws.stats, &seed, &it
        };
        cuda_check(cudaLaunchCooperativeKernel(
            reinterpret_cast<void*>(update_kernel), dim3(update_blocks), dim3(128), uargs,
            (128 + P) * sizeof(float), 0), "update launch");
    }

    auto out = torch::empty({iterations}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU));
    cuda_check(cudaMemcpy(out.data_ptr<float>(), ws.curve, iterations * sizeof(float),
                          cudaMemcpyDeviceToHost), "curve copy");
    return out;
}
"""


_ext = load_inline(
    name="rl_grid_ppo_sm120_v13_deterministic",
    cpp_sources=_CPP,
    cuda_sources=_CUDA,
    extra_cflags=["-O3"],
    extra_cuda_cflags=[
        "-O3",
        "--use_fast_math",
        "-gencode=arch=compute_120,code=sm_120",
    ],
    with_cuda=True,
    verbose=False,
)


def train(total_env_steps: int, seed: int) -> list[float]:
    iterations = max(1, int(total_env_steps) // (32 * 4096))
    return _ext.train(iterations, int(seed)).tolist()