diff --git a/config/craftax.ini b/config/craftax.ini index 22e0787987..fd21c161e4 100644 --- a/config/craftax.ini +++ b/config/craftax.ini @@ -2,8 +2,8 @@ env_name = craftax [vec] -total_agents = 16384 -num_buffers = 16 +total_agents = 4096 +num_buffers = 4 num_threads = 16 [env] @@ -13,18 +13,76 @@ seed_offset = 0 # Bounds world diversity: at most reset_pool_size unique maps are # ever seen per process. Set to 0 to disable (required for the # parity harness to maintain exact per-seed determinism). -reset_pool_size = 1024 +reset_pool_size = 32 [train] -total_timesteps = 200_000_000 -learning_rate = 0.008 -ent_coef = 0.02 -gamma = 0.997 -gae_lambda = 0.95 +total_timesteps = 1_000_000_000 +learning_rate = 0.011222594 +ent_coef = 0.005 +gamma = 0.98190003 +gae_lambda = 0.9217566 horizon = 128 minibatch_size = 32768 +replay_ratio = 1.7488248 +clip_coef = 0.41524854 +vf_coef = 2.7007978 +vf_clip_coef = 0.01 +max_grad_norm = 3.0366656 +beta1 = 0.93491626 +beta2 = 0.99692244 +eps = 1.2829034e-08 +vtrace_rho_clip = 0.79657767 +vtrace_c_clip = 1.0472921 +prio_alpha = 0.46023299 +prio_beta0 = 0.21990312 [policy] -hidden_size = 32 -num_layers = 1 +hidden_size = 128 +num_layers = 3 expansion_factor = 1 + +[sweep] +downsample = 5 + +[sweep.train.total_timesteps] +distribution = log_normal +min = 1e8 +max = 1e9 +mean = 5e8 +scale = time + +[sweep.train.ent_coef] +distribution = log_normal +min = 0.0001 +max = 0.005 +scale = auto + +[sweep.train.learning_rate] +distribution = log_normal +min = 0.0015 +max = 0.015 +scale = 0.5 + +[sweep.policy.num_layers] +distribution = uniform +min = 3 +max = 8 +scale = auto + +[sweep.train.horizon] +distribution = uniform_pow2 +min = 128 +max = 512 +scale = auto + +[sweep.policy.hidden_size] +distribution = uniform_pow2 +min = 128 +max = 512 +scale = auto + +[sweep.vec.num_buffers] +distribution = uniform +min = 1 +max = 4 +scale = auto diff --git a/src/craftax.cu b/src/craftax.cu new file mode 100644 index 0000000000..391e65bc38 --- /dev/null +++ b/src/craftax.cu @@ -0,0 +1,166 @@ +// Craftax CUDA encoder: bag-of-embeddings, concat, projection +// Included by ocean.cu — requires precision_t, PrecisionTensor, Allocator, puf_mm, etc. + +static constexpr int CX_OBS_ROWS = 9, CX_OBS_COLS = 11; +static constexpr int CX_NUM_CELLS = CX_OBS_ROWS * CX_OBS_COLS; // 99 +static constexpr int CX_CHANNELS = 8; // block,item,vis,5 mobs +static constexpr int CX_MAP_FLOATS = CX_NUM_CELLS * CX_CHANNELS; // 792 +static constexpr int CX_NUM_SCALARS = 51; // inventory/status tail +static constexpr int CX_EMB_DIM = 16; +static constexpr int CX_VOCAB_TOTAL = 154; // sum of per-channel vocabs +static constexpr int CX_MAP_EMB = CX_NUM_CELLS * CX_EMB_DIM; // 1584 +static constexpr int CX_CONCAT = CX_MAP_EMB + CX_NUM_SCALARS; // 1635 +// Per-channel base offset into the shared embedding table (vocab caps from +// worldgen.h: block<64, item+1<8, visible<2, mob+1<16). +__constant__ int CX_OFFSETS[CX_CHANNELS] = {0, 64, 72, 74, 90, 106, 122, 138}; + +// Cast float accumulation buffer to precision_t. +__global__ void craftax_float_to_precision_kernel( + precision_t* __restrict__ dst, const float* __restrict__ src, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) dst[idx] = from_float(src[idx]); +} + +// Bag-of-embeddings: sum channel embeddings per cell into concat's first +// CX_MAP_EMB columns. +__global__ void craftax_embed_bag_kernel( + precision_t* __restrict__ concat, const precision_t* __restrict__ obs, + const precision_t* __restrict__ embed_w, int B, int obs_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * CX_NUM_CELLS * CX_EMB_DIM) return; + int d = idx % CX_EMB_DIM; + int cell = (idx / CX_EMB_DIM) % CX_NUM_CELLS; + int b = idx / (CX_NUM_CELLS * CX_EMB_DIM); + const precision_t* ids = obs + (int64_t)b * obs_size + cell * CX_CHANNELS; + float sum = 0.0f; + #pragma unroll + for (int ch = 0; ch < CX_CHANNELS; ch++) { + int row = CX_OFFSETS[ch] + (int)to_float(ids[ch]); + sum += to_float(embed_w[row * CX_EMB_DIM + d]); + } + concat[(int64_t)b * CX_CONCAT + cell * CX_EMB_DIM + d] = from_float(sum); +} + +// Copy scalar tail into concat's trailing CX_NUM_SCALARS columns. +__global__ void craftax_copy_scalars_kernel( + precision_t* __restrict__ concat, const precision_t* __restrict__ obs, + int B, int obs_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * CX_NUM_SCALARS) return; + int j = idx % CX_NUM_SCALARS; + int b = idx / CX_NUM_SCALARS; + concat[(int64_t)b * CX_CONCAT + CX_MAP_EMB + j] = + obs[(int64_t)b * obs_size + CX_MAP_FLOATS + j]; +} + +// Embedding backward: scatter-add grad to every channel row that fed the sum +// (float buffer avoids bf16 atomics). +__global__ void craftax_embed_bag_backward_kernel( + float* __restrict__ embed_wgrad_f, const precision_t* __restrict__ grad_concat, + const precision_t* __restrict__ obs, int B, int obs_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * CX_NUM_CELLS * CX_EMB_DIM) return; + int d = idx % CX_EMB_DIM; + int cell = (idx / CX_EMB_DIM) % CX_NUM_CELLS; + int b = idx / (CX_NUM_CELLS * CX_EMB_DIM); + float g = to_float(grad_concat[(int64_t)b * CX_CONCAT + cell * CX_EMB_DIM + d]); + const precision_t* ids = obs + (int64_t)b * obs_size + cell * CX_CHANNELS; + #pragma unroll + for (int ch = 0; ch < CX_CHANNELS; ch++) { + int row = CX_OFFSETS[ch] + (int)to_float(ids[ch]); + atomicAdd(&embed_wgrad_f[row * CX_EMB_DIM + d], g); + } +} + +struct CraftaxEncoderWeights { + PrecisionTensor embed_w, proj_w; + int obs_size, hidden; +}; + +struct CraftaxEncoderActivations { + PrecisionTensor concat, out, saved_obs; // concat is reused as grad_concat in backward + PrecisionTensor embed_wgrad, proj_wgrad; + FloatTensor embed_wgrad_f; // float accumulation buffer for scatter-add +}; + +static PrecisionTensor craftax_encoder_forward(void* w, void* activations, PrecisionTensor input, cudaStream_t stream) { + CraftaxEncoderWeights* ew = (CraftaxEncoderWeights*)w; + CraftaxEncoderActivations* a = (CraftaxEncoderActivations*)activations; + int B = input.shape[0]; + if (a->saved_obs.data) puf_copy(&a->saved_obs, &input, stream); + + craftax_embed_bag_kernel<<>>( + a->concat.data, input.data, ew->embed_w.data, B, ew->obs_size); + craftax_copy_scalars_kernel<<>>( + a->concat.data, input.data, B, ew->obs_size); + + puf_mm(&a->concat, &ew->proj_w, &a->out, stream); // (B, CONCAT) @ (hidden, CONCAT)^T + return a->out; +} + +static void craftax_encoder_backward(void* w, void* activations, PrecisionTensor grad, cudaStream_t stream) { + CraftaxEncoderWeights* ew = (CraftaxEncoderWeights*)w; + CraftaxEncoderActivations* a = (CraftaxEncoderActivations*)activations; + int B = grad.shape[0]; + + // proj weight grad uses forward concat values -- must run before concat is reused. + puf_mm_tn(&grad, &a->concat, &a->proj_wgrad, stream); + // grad w.r.t. concat = grad @ proj_w; overwrite concat in place. + PrecisionTensor grad_concat = {.data = a->concat.data, .shape = {B, CX_CONCAT}}; + puf_mm_nn(&grad, &ew->proj_w, &grad_concat, stream); + + int embed_n = CX_VOCAB_TOTAL * CX_EMB_DIM; + cudaMemsetAsync(a->embed_wgrad_f.data, 0, embed_n * sizeof(float), stream); + craftax_embed_bag_backward_kernel<<>>( + a->embed_wgrad_f.data, grad_concat.data, a->saved_obs.data, B, ew->obs_size); + craftax_float_to_precision_kernel<<>>( + a->embed_wgrad.data, a->embed_wgrad_f.data, embed_n); +} + +static void craftax_encoder_init_weights(void* w, uint64_t* seed, cudaStream_t stream) { + CraftaxEncoderWeights* ew = (CraftaxEncoderWeights*)w; + puf_normal_init(&ew->embed_w, 1.0f, (*seed)++, stream); + PrecisionTensor pw = {.data = ew->proj_w.data, .shape = {ew->hidden, CX_CONCAT}}; + puf_kaiming_init(&pw, 1.0f, (*seed)++, stream); +} + +static void craftax_encoder_reg_params(void* w, Allocator* alloc) { + CraftaxEncoderWeights* ew = (CraftaxEncoderWeights*)w; + ew->embed_w = {.shape = {CX_VOCAB_TOTAL, CX_EMB_DIM}}; + ew->proj_w = {.shape = {ew->hidden, CX_CONCAT}}; + alloc_register(alloc, &ew->embed_w); + alloc_register(alloc, &ew->proj_w); +} + +static void craftax_encoder_reg_train(void* w, void* activations, Allocator* acts, Allocator* grads, int B_TT) { + CraftaxEncoderWeights* ew = (CraftaxEncoderWeights*)w; + CraftaxEncoderActivations* a = (CraftaxEncoderActivations*)activations; + *a = {}; + a->concat = {.shape = {B_TT, CX_CONCAT}}; + a->out = {.shape = {B_TT, ew->hidden}}; + a->saved_obs = {.shape = {B_TT, ew->obs_size}}; + alloc_register(acts, &a->concat); alloc_register(acts, &a->out); alloc_register(acts, &a->saved_obs); + a->embed_wgrad = {.shape = {CX_VOCAB_TOTAL, CX_EMB_DIM}}; + a->embed_wgrad_f = {.shape = {CX_VOCAB_TOTAL, CX_EMB_DIM}}; + a->proj_wgrad = {.shape = {ew->hidden, CX_CONCAT}}; + alloc_register(grads, &a->embed_wgrad); + alloc_register(acts, &a->embed_wgrad_f); + alloc_register(grads, &a->proj_wgrad); +} + +static void craftax_encoder_reg_rollout(void* w, void* activations, Allocator* alloc, int B) { + CraftaxEncoderWeights* ew = (CraftaxEncoderWeights*)w; + CraftaxEncoderActivations* a = (CraftaxEncoderActivations*)activations; + a->concat = {.shape = {B, CX_CONCAT}}; + a->out = {.shape = {B, ew->hidden}}; + alloc_register(alloc, &a->concat); alloc_register(alloc, &a->out); +} + +static void* craftax_encoder_create_weights(void* self) { + Encoder* e = (Encoder*)self; + CraftaxEncoderWeights* ew = (CraftaxEncoderWeights*)calloc(1, sizeof(CraftaxEncoderWeights)); + ew->obs_size = e->in_dim; ew->hidden = e->out_dim; + return ew; +} +static void craftax_encoder_free_weights(void* weights) { free(weights); } +static void craftax_encoder_free_activations(void* activations) { free(activations); } diff --git a/src/nmmo3.cu b/src/nmmo3.cu new file mode 100644 index 0000000000..4723021b43 --- /dev/null +++ b/src/nmmo3.cu @@ -0,0 +1,571 @@ +// NMMO3 CUDA encoder: multihot, cuDNN conv, embedding, concat, projection +// Included by pufferlib.cu — requires precision_t, PrecisionTensor, Allocator, puf_mm, etc. + +#include "cudnn_conv2d.cu" + +// ---- NMMO3 constants ---- + +static constexpr int N3_MAP_H = 11, N3_MAP_W = 15, N3_NFEAT = 10; +static constexpr int N3_MULTIHOT = 59; +static constexpr int N3_MAP_SIZE = N3_MAP_H * N3_MAP_W * N3_NFEAT; +static constexpr int N3_PLAYER = 47, N3_REWARD = 10; +static constexpr int N3_EMBED_DIM = 32, N3_EMBED_VOCAB = 128; +static constexpr int N3_PLAYER_EMBED = N3_PLAYER * N3_EMBED_DIM; +static constexpr int N3_C1_IC = 59, N3_C1_OC = 128, N3_C1_K = 5, N3_C1_S = 3; +static constexpr int N3_C1_OH = 3, N3_C1_OW = 4; +static constexpr int N3_C2_IC = 128, N3_C2_OC = 128, N3_C2_K = 3, N3_C2_S = 1; +static constexpr int N3_C2_OH = 1, N3_C2_OW = 2; +static constexpr int N3_CONV_FLAT = N3_C2_OC * N3_C2_OH * N3_C2_OW; +static constexpr int N3_CONCAT = N3_CONV_FLAT + N3_PLAYER_EMBED + N3_PLAYER + N3_REWARD; + +__constant__ int N3_OFFSETS[10] = {0, 4, 8, 25, 30, 33, 38, 43, 48, 55}; + +static cudnnDataType_t n3_cudnn_dtype() { + return (PRECISION_SIZE == 2) ? CUDNN_DATA_BFLOAT16 : CUDNN_DATA_FLOAT; +} + +// ---- NMMO3 kernels ---- + +__global__ void n3_multihot_kernel( + precision_t* __restrict__ out, const precision_t* __restrict__ obs, int B, int obs_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * N3_MAP_H * N3_MAP_W) return; + int b = idx / (N3_MAP_H * N3_MAP_W), rem = idx % (N3_MAP_H * N3_MAP_W); + int h = rem / N3_MAP_W, w = rem % N3_MAP_W; + const precision_t* src = obs + b * obs_size + (h * N3_MAP_W + w) * N3_NFEAT; + precision_t* dst = out + b * N3_MULTIHOT * N3_MAP_H * N3_MAP_W; + for (int f = 0; f < N3_NFEAT; f++) + dst[(N3_OFFSETS[f] + (int)to_float(src[f])) * N3_MAP_H * N3_MAP_W + h * N3_MAP_W + w] = from_float(1.0f); +} + +__global__ void n3_embedding_kernel( + precision_t* __restrict__ out, const precision_t* __restrict__ obs, + const precision_t* __restrict__ embed_w, int B, int obs_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * N3_PLAYER) return; + int b = idx / N3_PLAYER, f = idx % N3_PLAYER; + int val = (int)to_float(obs[b * obs_size + N3_MAP_SIZE + f]); + const precision_t* src = embed_w + val * N3_EMBED_DIM; + precision_t* dst = out + b * N3_PLAYER_EMBED + f * N3_EMBED_DIM; + for (int d = 0; d < N3_EMBED_DIM; d++) dst[d] = src[d]; +} + +__global__ void n3_concat_kernel( + precision_t* __restrict__ out, const precision_t* __restrict__ conv_flat, + const precision_t* __restrict__ embed, const precision_t* __restrict__ obs, + int B, int obs_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * N3_CONCAT) return; + int b = idx / N3_CONCAT, c = idx % N3_CONCAT; + precision_t val; + if (c < N3_CONV_FLAT) { + int oc = c / (N3_C2_OH * N3_C2_OW), r = c % (N3_C2_OH * N3_C2_OW); + int oh = r / N3_C2_OW, ow = r % N3_C2_OW; + val = conv_flat[b * N3_CONV_FLAT + oc * N3_C2_OH * N3_C2_OW + oh * N3_C2_OW + ow]; + } else if (c < N3_CONV_FLAT + N3_PLAYER_EMBED) + val = embed[b * N3_PLAYER_EMBED + (c - N3_CONV_FLAT)]; + else if (c < N3_CONV_FLAT + N3_PLAYER_EMBED + N3_PLAYER) + val = obs[b * obs_size + N3_MAP_SIZE + (c - N3_CONV_FLAT - N3_PLAYER_EMBED)]; + else + val = obs[b * obs_size + obs_size - N3_REWARD + (c - N3_CONV_FLAT - N3_PLAYER_EMBED - N3_PLAYER)]; + out[idx] = val; +} + +__global__ void n3_bias_relu_kernel( + precision_t* __restrict__ data, const precision_t* __restrict__ bias, int total, int dim) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) return; + data[idx] = from_float(fmaxf(0.0f, to_float(data[idx]) + to_float(bias[idx % dim]))); +} + +__global__ void n3_relu_backward_kernel( + precision_t* __restrict__ grad, const precision_t* __restrict__ out, int total) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) return; + if (to_float(out[idx]) <= 0.0f) grad[idx] = from_float(0.0f); +} + + +__global__ void bias_grad_kernel( + precision_t* __restrict__ bgrad, const precision_t* __restrict__ grad, int N, int dim) { + int d = blockIdx.x; + if (d >= dim) return; + float sum = 0.0f; + for (int i = threadIdx.x; i < N; i += blockDim.x) + sum += to_float(grad[i * dim + d]); + for (int offset = 16; offset > 0; offset >>= 1) + sum += __shfl_down_sync(0xffffffff, sum, offset); + __shared__ float sdata[32]; + int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + if (lane == 0) sdata[warp] = sum; + __syncthreads(); + if (warp == 0) { + sum = (lane < (blockDim.x + 31) / 32) ? sdata[lane] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) + sum += __shfl_down_sync(0xffffffff, sum, offset); + if (lane == 0) bgrad[d] = from_float(sum); + } +} + +// NCHW bias grad: sum over (B, OH, OW) for each OC channel +__global__ void n3_conv_bias_grad_nchw( + precision_t* __restrict__ bgrad, const precision_t* __restrict__ grad, + int B, int OC, int spatial) { + int oc = blockIdx.x; + if (oc >= OC) return; + float sum = 0.0f; + int total = B * spatial; + for (int i = threadIdx.x; i < total; i += blockDim.x) { + int b = i / spatial, s = i % spatial; + sum += to_float(grad[b * OC * spatial + oc * spatial + s]); + } + for (int offset = 16; offset > 0; offset >>= 1) + sum += __shfl_down_sync(0xffffffff, sum, offset); + __shared__ float sdata[32]; + int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + if (lane == 0) sdata[warp] = sum; + __syncthreads(); + if (warp == 0) { + sum = (lane < (blockDim.x + 31) / 32) ? sdata[lane] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) + sum += __shfl_down_sync(0xffffffff, sum, offset); + if (lane == 0) bgrad[oc] = from_float(sum); + } +} + +__global__ void n3_concat_backward_conv_kernel( + precision_t* __restrict__ conv_grad, const precision_t* __restrict__ concat_grad, int B) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * N3_CONV_FLAT) return; + int b = idx / N3_CONV_FLAT, c = idx % N3_CONV_FLAT; + conv_grad[b * N3_CONV_FLAT + c] = concat_grad[b * N3_CONCAT + c]; +} + +// Embedding backward: scatter-add grad from concat_grad's player_embed region +// into embed_wgrad (float accumulation buffer). +// Each (b, f) looked up row obs[b, MAP_SIZE+f] from the table. +__global__ void n3_embedding_backward_kernel( + float* __restrict__ embed_wgrad_f, + const precision_t* __restrict__ concat_grad, + const precision_t* __restrict__ obs, + int B, int obs_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * N3_PLAYER * N3_EMBED_DIM) return; + int b = idx / (N3_PLAYER * N3_EMBED_DIM); + int rem = idx % (N3_PLAYER * N3_EMBED_DIM); + int f = rem / N3_EMBED_DIM; + int d = rem % N3_EMBED_DIM; + int val = (int)to_float(obs[b * obs_size + N3_MAP_SIZE + f]); + float g = to_float(concat_grad[b * N3_CONCAT + N3_CONV_FLAT + f * N3_EMBED_DIM + d]); + atomicAdd(&embed_wgrad_f[val * N3_EMBED_DIM + d], g); +} + +// Cast float buffer to precision_t +__global__ void n3_float_to_precision_kernel( + precision_t* __restrict__ dst, const float* __restrict__ src, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) dst[idx] = from_float(src[idx]); +} + +// ---- atomicAdd for precision_t ---- +#ifdef PRECISION_FLOAT +__device__ __forceinline__ void atomicAdd_precision(precision_t* addr, precision_t val) { + atomicAdd(addr, val); +} +#else +__device__ __forceinline__ void atomicAdd_precision(precision_t* addr, precision_t val) { + // bf16 atomicAdd via CAS on enclosing 32-bit word + unsigned int* addr_u32 = (unsigned int*)((size_t)addr & ~2ULL); + bool is_high = ((size_t)addr & 2) != 0; + unsigned int old_u32 = *addr_u32, assumed; + do { + assumed = old_u32; + __nv_bfloat16* pair = (__nv_bfloat16*)&old_u32; + float sum = __bfloat162float(pair[is_high]) + __bfloat162float(val); + unsigned int new_u32 = assumed; + ((__nv_bfloat16*)&new_u32)[is_high] = __float2bfloat16(sum); + old_u32 = atomicCAS(addr_u32, assumed, new_u32); + } while (old_u32 != assumed); +} +#endif + +// ---- NCHW bias kernels for im2col conv path ---- + +__global__ void conv_bias_kernel(precision_t* __restrict__ data, + const precision_t* __restrict__ bias, int B, int OC, int spatial) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = B * OC * spatial; + if (idx >= total) return; + int oc = (idx / spatial) % OC; + data[idx] = from_float(to_float(data[idx]) + to_float(bias[oc])); +} + +__global__ void conv_bias_relu_kernel(precision_t* __restrict__ data, + const precision_t* __restrict__ bias, int B, int OC, int spatial) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = B * OC * spatial; + if (idx >= total) return; + int oc = (idx / spatial) % OC; + data[idx] = from_float(fmaxf(0.0f, to_float(data[idx]) + to_float(bias[oc]))); +} + +// ---- im2col + cuBLAS conv (no cuDNN) ---- +// NCHW layout throughout. Weight stored as (OC, IC*K*K). +// im2col produces (B*OH*OW, IC*K*K), matmul with W^T gives (B*OH*OW, OC), +// then reshape to NCHW (B, OC, OH, OW). + +__global__ void im2col_kernel( + const precision_t* __restrict__ input, precision_t* __restrict__ col, + int B, int IC, int IH, int IW, int K, int S, int OH, int OW +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = B * OH * OW * IC * K * K; + if (idx >= total) return; + int col_w = IC * K * K; + int row = idx / col_w; + int c = idx % col_w; + int b = row / (OH * OW); + int rem = row % (OH * OW); + int oh = rem / OW, ow = rem % OW; + int ic = c / (K * K), kk = c % (K * K); + int kh = kk / K, kw = kk % K; + int ih = oh * S + kh, iw = ow * S + kw; + col[idx] = input[b * IC * IH * IW + ic * IH * IW + ih * IW + iw]; +} + +// Backward: col2im — input-centric gather to avoid atomics. +// Each thread owns one (b, ic, ih, iw) element and sums contributions from all +// (oh, ow, kh, kw) patches that map to it. +__global__ void col2im_kernel( + const precision_t* __restrict__ col, precision_t* __restrict__ grad_input, + int B, int IC, int IH, int IW, int K, int S, int OH, int OW +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = B * IC * IH * IW; + if (idx >= total) return; + int iw = idx % IW; + int ih = (idx / IW) % IH; + int ic = (idx / (IW * IH)) % IC; + int b = idx / (IW * IH * IC); + float sum = 0.0f; + for (int kh = 0; kh < K; kh++) { + int ih_off = ih - kh; + if (ih_off < 0 || ih_off % S != 0) continue; + int oh = ih_off / S; + if (oh >= OH) continue; + for (int kw = 0; kw < K; kw++) { + int iw_off = iw - kw; + if (iw_off < 0 || iw_off % S != 0) continue; + int ow = iw_off / S; + if (ow >= OW) continue; + int col_idx = (b * OH * OW + oh * OW + ow) * (IC * K * K) + ic * K * K + kh * K + kw; + sum += to_float(col[col_idx]); + } + } + grad_input[idx] = from_float(sum); +} + +// Transpose (B, OC, OH, OW) -> (B*OH*OW, OC) [NCHW to row-major spatial-first] +__global__ void nchw_to_rows_kernel( + const precision_t* __restrict__ src, precision_t* __restrict__ dst, + int B, int OC, int spatial +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = B * OC * spatial; + if (idx >= total) return; + int b = idx / (OC * spatial); + int oc = (idx / spatial) % OC; + int s = idx % spatial; + dst[(b * spatial + s) * OC + oc] = src[idx]; +} + +// Transpose (B*OH*OW, OC) -> (B, OC, OH, OW) [row-major spatial-first to NCHW] +__global__ void rows_to_nchw_kernel( + const precision_t* __restrict__ src, precision_t* __restrict__ dst, + int B, int OC, int spatial +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = B * OC * spatial; + if (idx >= total) return; + int b = idx / (OC * spatial); + int oc = (idx / spatial) % OC; + int s = idx % spatial; + dst[idx] = src[(b * spatial + s) * OC + oc]; +} + +// Forward: im2col conv + bias + optional relu. All NCHW. +// col_buf: pre-allocated (max_B * OH * OW, IC * K * K) +// mm_buf: pre-allocated (max_B * OH * OW, OC) — row-major (spatial-first) +static void gemm_conv_forward( + PrecisionTensor* weight, PrecisionTensor* bias, + precision_t* input, precision_t* output, + precision_t* col_buf, precision_t* mm_buf, + int B, int IC, int IH, int IW, int OC, int K, int S, int OH, int OW, + bool relu, cudaStream_t stream +) { + int col_rows = B * OH * OW; + int col_cols = IC * K * K; + int total_col = col_rows * col_cols; + int total_out = B * OC * OH * OW; + + // im2col: input NCHW -> col (B*OH*OW, IC*K*K) + im2col_kernel<<>>( + input, col_buf, B, IC, IH, IW, K, S, OH, OW); + + // matmul: col (B*OH*OW, IC*K*K) @ W^T (IC*K*K, OC) = mm_buf (B*OH*OW, OC) + PrecisionTensor col_t = {.data = col_buf, .shape = {col_rows, col_cols}}; + PrecisionTensor mm_t = {.data = mm_buf, .shape = {col_rows, OC}}; + puf_mm(&col_t, weight, &mm_t, stream); + + // transpose (B*OH*OW, OC) -> (B, OC, OH, OW) NCHW + bias + relu + int spatial = OH * OW; + rows_to_nchw_kernel<<>>( + mm_buf, output, B, OC, spatial); + if (relu) { + conv_bias_relu_kernel<<>>( + output, bias->data, B, OC, spatial); + } else { + conv_bias_kernel<<>>( + output, bias->data, B, OC, spatial); + } +} + +// Backward: weight grad + optional input grad via im2col/col2im + cuBLAS. +// grad_output is NCHW (B, OC, OH, OW). saved_input is NCHW. +// Caller handles relu backward and bias grad (same as cuDNN path). +static void gemm_conv_backward( + PrecisionTensor* weight, + precision_t* saved_input, precision_t* grad_output, + precision_t* wgrad, precision_t* input_grad, + precision_t* col_buf, precision_t* mm_buf, + int B, int IC, int IH, int IW, int OC, int K, int S, int OH, int OW, + cudaStream_t stream +) { + int col_rows = B * OH * OW; + int col_cols = IC * K * K; + int total_col = col_rows * col_cols; + int total_out = B * OC * OH * OW; + int spatial = OH * OW; + + // Transpose grad_output NCHW -> (B*OH*OW, OC) + nchw_to_rows_kernel<<>>( + grad_output, mm_buf, B, OC, spatial); + + // im2col of saved_input + im2col_kernel<<>>( + saved_input, col_buf, B, IC, IH, IW, K, S, OH, OW); + + // Weight grad: mm_buf^T (OC, B*OH*OW) @ col_buf (B*OH*OW, IC*K*K) = wgrad (OC, IC*K*K) + PrecisionTensor mm_t = {.data = mm_buf, .shape = {col_rows, OC}}; + PrecisionTensor col_t = {.data = col_buf, .shape = {col_rows, col_cols}}; + PrecisionTensor wg_t = {.data = wgrad, .shape = {OC, col_cols}}; + puf_mm_tn(&mm_t, &col_t, &wg_t, stream); + + // Input grad (optional): mm_buf (B*OH*OW, OC) @ weight (OC, IC*K*K) = col_grad (B*OH*OW, IC*K*K) + if (input_grad) { + puf_mm_nn(&mm_t, weight, &col_t, stream); // reuse col_buf as col_grad + col2im_kernel<<>>( + col_buf, input_grad, B, IC, IH, IW, K, S, OH, OW); + } +} + +// ---- NMMO3 encoder structs ---- + +struct NMMO3EncoderWeights { + ConvWeights conv1, conv2; + PrecisionTensor embed_w, proj_w, proj_b; + int obs_size, hidden; +}; + +struct NMMO3EncoderActivations { + ConvActivations conv1, conv2; + PrecisionTensor col1, mm1, col2, mm2; // im2col + matmul scratch buffers + PrecisionTensor multihot, embed_out, concat, out, saved_obs; + PrecisionTensor embed_wgrad, proj_wgrad, proj_bgrad; + FloatTensor embed_wgrad_f; // float accumulation buffer for scatter-add +}; + +static NMMO3EncoderWeights* nmmo3_encoder_create(int obs_size, int hidden) { + NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)calloc(1, sizeof(NMMO3EncoderWeights)); + ew->obs_size = obs_size; ew->hidden = hidden; + conv_init(&ew->conv1, N3_C1_IC, N3_C1_OC, N3_C1_K, N3_C1_S, N3_MAP_H, N3_MAP_W, true); + conv_init(&ew->conv2, N3_C2_IC, N3_C2_OC, N3_C2_K, N3_C2_S, N3_C1_OH, N3_C1_OW, false); + return ew; +} + +// ---- NMMO3 encoder interface ---- + +static PrecisionTensor nmmo3_encoder_forward(void* w, void* activations, PrecisionTensor input, cudaStream_t stream) { + NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; + NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; + int B = input.shape[0]; + + if (a->saved_obs.data) puf_copy(&a->saved_obs, &input, stream); + + cudaMemsetAsync(a->multihot.data, 0, (int64_t)B * N3_MULTIHOT * N3_MAP_H * N3_MAP_W * sizeof(precision_t), stream); + n3_multihot_kernel<<>>( + a->multihot.data, input.data, B, ew->obs_size); + + gemm_conv_forward(&ew->conv1.w, &ew->conv1.b, a->multihot.data, a->conv1.out.data, + a->col1.data, a->mm1.data, B, N3_C1_IC, N3_MAP_H, N3_MAP_W, + N3_C1_OC, N3_C1_K, N3_C1_S, N3_C1_OH, N3_C1_OW, true, stream); + if (a->conv1.saved_input.data) + cudaMemcpyAsync(a->conv1.saved_input.data, a->multihot.data, + (int64_t)B * N3_C1_IC * N3_MAP_H * N3_MAP_W * sizeof(precision_t), cudaMemcpyDeviceToDevice, stream); + gemm_conv_forward(&ew->conv2.w, &ew->conv2.b, a->conv1.out.data, a->conv2.out.data, + a->col2.data, a->mm2.data, B, N3_C2_IC, N3_C1_OH, N3_C1_OW, + N3_C2_OC, N3_C2_K, N3_C2_S, N3_C2_OH, N3_C2_OW, false, stream); + if (a->conv2.saved_input.data) + cudaMemcpyAsync(a->conv2.saved_input.data, a->conv1.out.data, + (int64_t)B * N3_C2_IC * N3_C1_OH * N3_C1_OW * sizeof(precision_t), cudaMemcpyDeviceToDevice, stream); + + n3_embedding_kernel<<>>( + a->embed_out.data, input.data, ew->embed_w.data, B, ew->obs_size); + n3_concat_kernel<<>>( + a->concat.data, a->conv2.out.data, a->embed_out.data, input.data, B, ew->obs_size); + + puf_mm(&a->concat, &ew->proj_w, &a->out, stream); + n3_bias_relu_kernel<<hidden), BLOCK_SIZE, 0, stream>>>( + a->out.data, ew->proj_b.data, B * ew->hidden, ew->hidden); + return a->out; +} + +static void nmmo3_encoder_backward(void* w, void* activations, PrecisionTensor grad, cudaStream_t stream) { + NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; + NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; + int B = grad.shape[0], H = ew->hidden; + + n3_relu_backward_kernel<<>>( + grad.data, a->out.data, B * H); + bias_grad_kernel<<>>( + a->proj_bgrad.data, grad.data, B, H); + puf_mm_tn(&grad, &a->concat, &a->proj_wgrad, stream); + + PrecisionTensor grad_concat = {.data = a->concat.data, .shape = {B, N3_CONCAT}}; + puf_mm_nn(&grad, &ew->proj_w, &grad_concat, stream); + + n3_concat_backward_conv_kernel<<>>( + a->conv2.grad.data, grad_concat.data, B); + + n3_conv_bias_grad_nchw<<conv2.OC, 256, 0, stream>>>( + a->conv2.bgrad.data, a->conv2.grad.data, + B, ew->conv2.OC, ew->conv2.OH * ew->conv2.OW); + gemm_conv_backward(&ew->conv2.w, a->conv2.saved_input.data, a->conv2.grad.data, + a->conv2.wgrad.data, a->conv1.grad.data, + a->col2.data, a->mm2.data, B, N3_C2_IC, N3_C1_OH, N3_C1_OW, + N3_C2_OC, N3_C2_K, N3_C2_S, N3_C2_OH, N3_C2_OW, stream); + + n3_relu_backward_kernel<<conv1.OC * ew->conv1.OH * ew->conv1.OW), BLOCK_SIZE, 0, stream>>>( + a->conv1.grad.data, a->conv1.out.data, + B * ew->conv1.OC * ew->conv1.OH * ew->conv1.OW); + n3_conv_bias_grad_nchw<<conv1.OC, 256, 0, stream>>>( + a->conv1.bgrad.data, a->conv1.grad.data, + B, ew->conv1.OC, ew->conv1.OH * ew->conv1.OW); + gemm_conv_backward(&ew->conv1.w, a->conv1.saved_input.data, a->conv1.grad.data, + a->conv1.wgrad.data, NULL, + a->col1.data, a->mm1.data, B, N3_C1_IC, N3_MAP_H, N3_MAP_W, + N3_C1_OC, N3_C1_K, N3_C1_S, N3_C1_OH, N3_C1_OW, stream); + + // Embedding backward: scatter-add from concat gradient into float buffer, then cast + int embed_n = N3_EMBED_VOCAB * N3_EMBED_DIM; + cudaMemsetAsync(a->embed_wgrad_f.data, 0, embed_n * sizeof(float), stream); + n3_embedding_backward_kernel<<>>( + a->embed_wgrad_f.data, grad_concat.data, a->saved_obs.data, B, ew->obs_size); + n3_float_to_precision_kernel<<>>( + a->embed_wgrad.data, a->embed_wgrad_f.data, embed_n); +} + +static void nmmo3_encoder_init_weights(void* w, uint64_t* seed, cudaStream_t stream) { + NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; + conv_init_weights(&ew->conv1, seed, stream); + conv_init_weights(&ew->conv2, seed, stream); + auto init2d = [&](PrecisionTensor& t, int rows, int cols, float gain) { + PrecisionTensor wt = {.data = t.data, .shape = {rows, cols}}; + puf_kaiming_init(&wt, gain, (*seed)++, stream); + }; + puf_normal_init(&ew->embed_w, 1.0f, (*seed)++, stream); + init2d(ew->proj_w, ew->hidden, N3_CONCAT, 1.0f); + cudaMemsetAsync(ew->proj_b.data, 0, numel(ew->proj_b.shape) * sizeof(precision_t), stream); +} + +static void nmmo3_encoder_reg_params(void* w, Allocator* alloc) { + NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; + conv_reg_params(&ew->conv1, alloc); + conv_reg_params(&ew->conv2, alloc); + ew->embed_w = {.shape = {N3_EMBED_VOCAB, N3_EMBED_DIM}}; + ew->proj_w = {.shape = {ew->hidden, N3_CONCAT}}; + ew->proj_b = {.shape = {ew->hidden}}; + alloc_register(alloc,&ew->embed_w); + alloc_register(alloc,&ew->proj_w); alloc_register(alloc,&ew->proj_b); +} + +static void nmmo3_encoder_reg_train(void* w, void* activations, Allocator* acts, Allocator* grads, int B_TT) { + NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; + NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; + *a = {}; + a->multihot = {.shape = {B_TT, N3_MULTIHOT * N3_MAP_H * N3_MAP_W}}; + alloc_register(acts,&a->multihot); + // Conv1 buffers + a->conv1.out = {.shape = {B_TT * N3_C1_OC * N3_C1_OH * N3_C1_OW}}; + a->conv1.grad = {.shape = {B_TT * N3_C1_OC * N3_C1_OH * N3_C1_OW}}; + a->conv1.saved_input = {.shape = {B_TT * N3_C1_IC * N3_MAP_H * N3_MAP_W}}; + a->conv1.wgrad = {.shape = {N3_C1_OC, N3_C1_IC * N3_C1_K * N3_C1_K}}; + a->conv1.bgrad = {.shape = {N3_C1_OC}}; + alloc_register(acts,&a->conv1.out); alloc_register(acts,&a->conv1.grad); alloc_register(acts,&a->conv1.saved_input); + alloc_register(grads,&a->conv1.wgrad); alloc_register(grads,&a->conv1.bgrad); + a->col1 = {.shape = {B_TT * N3_C1_OH * N3_C1_OW, N3_C1_IC * N3_C1_K * N3_C1_K}}; + a->mm1 = {.shape = {B_TT * N3_C1_OH * N3_C1_OW, N3_C1_OC}}; + alloc_register(acts,&a->col1); alloc_register(acts,&a->mm1); + // Conv2 buffers + a->conv2.out = {.shape = {B_TT * N3_C2_OC * N3_C2_OH * N3_C2_OW}}; + a->conv2.grad = {.shape = {B_TT * N3_C2_OC * N3_C2_OH * N3_C2_OW}}; + a->conv2.saved_input = {.shape = {B_TT * N3_C2_IC * N3_C1_OH * N3_C1_OW}}; + a->conv2.wgrad = {.shape = {N3_C2_OC, N3_C2_IC * N3_C2_K * N3_C2_K}}; + a->conv2.bgrad = {.shape = {N3_C2_OC}}; + alloc_register(acts,&a->conv2.out); alloc_register(acts,&a->conv2.grad); alloc_register(acts,&a->conv2.saved_input); + alloc_register(grads,&a->conv2.wgrad); alloc_register(grads,&a->conv2.bgrad); + a->col2 = {.shape = {B_TT * N3_C2_OH * N3_C2_OW, N3_C2_IC * N3_C2_K * N3_C2_K}}; + a->mm2 = {.shape = {B_TT * N3_C2_OH * N3_C2_OW, N3_C2_OC}}; + alloc_register(acts,&a->col2); alloc_register(acts,&a->mm2); + a->embed_out = {.shape = {B_TT, N3_PLAYER_EMBED}}; + a->concat = {.shape = {B_TT, N3_CONCAT}}; + a->out = {.shape = {B_TT, ew->hidden}}; + a->saved_obs = {.shape = {B_TT, ew->obs_size}}; + alloc_register(acts,&a->embed_out); alloc_register(acts,&a->concat); + alloc_register(acts,&a->out); alloc_register(acts,&a->saved_obs); + a->embed_wgrad = {.shape = {N3_EMBED_VOCAB, N3_EMBED_DIM}}; + a->embed_wgrad_f = {.shape = {N3_EMBED_VOCAB, N3_EMBED_DIM}}; + a->proj_wgrad = {.shape = {ew->hidden, N3_CONCAT}}; + a->proj_bgrad = {.shape = {ew->hidden}}; + alloc_register(grads,&a->embed_wgrad); + alloc_register(acts,&a->embed_wgrad_f); + alloc_register(grads,&a->proj_wgrad); alloc_register(grads,&a->proj_bgrad); +} + +static void nmmo3_encoder_reg_rollout(void* w, void* activations, Allocator* alloc, int B) { + NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; + NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; + a->multihot = {.shape = {B, N3_MULTIHOT * N3_MAP_H * N3_MAP_W}}; + alloc_register(alloc,&a->multihot); + a->conv1.out = {.shape = {B * N3_C1_OC * N3_C1_OH * N3_C1_OW}}; + alloc_register(alloc,&a->conv1.out); + a->col1 = {.shape = {B * N3_C1_OH * N3_C1_OW, N3_C1_IC * N3_C1_K * N3_C1_K}}; + a->mm1 = {.shape = {B * N3_C1_OH * N3_C1_OW, N3_C1_OC}}; + alloc_register(alloc,&a->col1); alloc_register(alloc,&a->mm1); + a->conv2.out = {.shape = {B * N3_C2_OC * N3_C2_OH * N3_C2_OW}}; + alloc_register(alloc,&a->conv2.out); + a->col2 = {.shape = {B * N3_C2_OH * N3_C2_OW, N3_C2_IC * N3_C2_K * N3_C2_K}}; + a->mm2 = {.shape = {B * N3_C2_OH * N3_C2_OW, N3_C2_OC}}; + alloc_register(alloc,&a->col2); alloc_register(alloc,&a->mm2); + a->embed_out = {.shape = {B, N3_PLAYER_EMBED}}; + a->concat = {.shape = {B, N3_CONCAT}}; + a->out = {.shape = {B, ew->hidden}}; + alloc_register(alloc,&a->embed_out); alloc_register(alloc,&a->concat); alloc_register(alloc,&a->out); +} + +static void* nmmo3_encoder_create_weights(void* self) { + Encoder* e = (Encoder*)self; + return nmmo3_encoder_create(e->in_dim, e->out_dim); +} +static void nmmo3_encoder_free_weights(void* weights) { free(weights); } +static void nmmo3_encoder_free_activations(void* activations) { free(activations); } diff --git a/src/ocean.cu b/src/ocean.cu index baaa9b7be6..e75ae8c013 100644 --- a/src/ocean.cu +++ b/src/ocean.cu @@ -1,577 +1,29 @@ -// NMMO3 CUDA encoder: multihot, cuDNN conv, embedding, concat, projection +// Custom CUDA encoder registry for ocean environments. +// Each env's encoder lives in its own translation-unit include below; this file +// only wires them into the Encoder vtable via create_custom_encoder(). // Included by pufferlib.cu — requires precision_t, PrecisionTensor, Allocator, puf_mm, etc. -#include "cudnn_conv2d.cu" - -// ---- NMMO3 constants ---- - -static constexpr int N3_MAP_H = 11, N3_MAP_W = 15, N3_NFEAT = 10; -static constexpr int N3_MULTIHOT = 59; -static constexpr int N3_MAP_SIZE = N3_MAP_H * N3_MAP_W * N3_NFEAT; -static constexpr int N3_PLAYER = 47, N3_REWARD = 10; -static constexpr int N3_EMBED_DIM = 32, N3_EMBED_VOCAB = 128; -static constexpr int N3_PLAYER_EMBED = N3_PLAYER * N3_EMBED_DIM; -static constexpr int N3_C1_IC = 59, N3_C1_OC = 128, N3_C1_K = 5, N3_C1_S = 3; -static constexpr int N3_C1_OH = 3, N3_C1_OW = 4; -static constexpr int N3_C2_IC = 128, N3_C2_OC = 128, N3_C2_K = 3, N3_C2_S = 1; -static constexpr int N3_C2_OH = 1, N3_C2_OW = 2; -static constexpr int N3_CONV_FLAT = N3_C2_OC * N3_C2_OH * N3_C2_OW; -static constexpr int N3_CONCAT = N3_CONV_FLAT + N3_PLAYER_EMBED + N3_PLAYER + N3_REWARD; - -__constant__ int N3_OFFSETS[10] = {0, 4, 8, 25, 30, 33, 38, 43, 48, 55}; - -static cudnnDataType_t n3_cudnn_dtype() { - return (PRECISION_SIZE == 2) ? CUDNN_DATA_BFLOAT16 : CUDNN_DATA_FLOAT; -} - -// ---- NMMO3 kernels ---- - -__global__ void n3_multihot_kernel( - precision_t* __restrict__ out, const precision_t* __restrict__ obs, int B, int obs_size) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= B * N3_MAP_H * N3_MAP_W) return; - int b = idx / (N3_MAP_H * N3_MAP_W), rem = idx % (N3_MAP_H * N3_MAP_W); - int h = rem / N3_MAP_W, w = rem % N3_MAP_W; - const precision_t* src = obs + b * obs_size + (h * N3_MAP_W + w) * N3_NFEAT; - precision_t* dst = out + b * N3_MULTIHOT * N3_MAP_H * N3_MAP_W; - for (int f = 0; f < N3_NFEAT; f++) - dst[(N3_OFFSETS[f] + (int)to_float(src[f])) * N3_MAP_H * N3_MAP_W + h * N3_MAP_W + w] = from_float(1.0f); -} - -__global__ void n3_embedding_kernel( - precision_t* __restrict__ out, const precision_t* __restrict__ obs, - const precision_t* __restrict__ embed_w, int B, int obs_size) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= B * N3_PLAYER) return; - int b = idx / N3_PLAYER, f = idx % N3_PLAYER; - int val = (int)to_float(obs[b * obs_size + N3_MAP_SIZE + f]); - const precision_t* src = embed_w + val * N3_EMBED_DIM; - precision_t* dst = out + b * N3_PLAYER_EMBED + f * N3_EMBED_DIM; - for (int d = 0; d < N3_EMBED_DIM; d++) dst[d] = src[d]; -} - -__global__ void n3_concat_kernel( - precision_t* __restrict__ out, const precision_t* __restrict__ conv_flat, - const precision_t* __restrict__ embed, const precision_t* __restrict__ obs, - int B, int obs_size) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= B * N3_CONCAT) return; - int b = idx / N3_CONCAT, c = idx % N3_CONCAT; - precision_t val; - if (c < N3_CONV_FLAT) { - int oc = c / (N3_C2_OH * N3_C2_OW), r = c % (N3_C2_OH * N3_C2_OW); - int oh = r / N3_C2_OW, ow = r % N3_C2_OW; - val = conv_flat[b * N3_CONV_FLAT + oc * N3_C2_OH * N3_C2_OW + oh * N3_C2_OW + ow]; - } else if (c < N3_CONV_FLAT + N3_PLAYER_EMBED) - val = embed[b * N3_PLAYER_EMBED + (c - N3_CONV_FLAT)]; - else if (c < N3_CONV_FLAT + N3_PLAYER_EMBED + N3_PLAYER) - val = obs[b * obs_size + N3_MAP_SIZE + (c - N3_CONV_FLAT - N3_PLAYER_EMBED)]; - else - val = obs[b * obs_size + obs_size - N3_REWARD + (c - N3_CONV_FLAT - N3_PLAYER_EMBED - N3_PLAYER)]; - out[idx] = val; -} - -__global__ void n3_bias_relu_kernel( - precision_t* __restrict__ data, const precision_t* __restrict__ bias, int total, int dim) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) return; - data[idx] = from_float(fmaxf(0.0f, to_float(data[idx]) + to_float(bias[idx % dim]))); -} - -__global__ void n3_relu_backward_kernel( - precision_t* __restrict__ grad, const precision_t* __restrict__ out, int total) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) return; - if (to_float(out[idx]) <= 0.0f) grad[idx] = from_float(0.0f); -} - - -__global__ void bias_grad_kernel( - precision_t* __restrict__ bgrad, const precision_t* __restrict__ grad, int N, int dim) { - int d = blockIdx.x; - if (d >= dim) return; - float sum = 0.0f; - for (int i = threadIdx.x; i < N; i += blockDim.x) - sum += to_float(grad[i * dim + d]); - for (int offset = 16; offset > 0; offset >>= 1) - sum += __shfl_down_sync(0xffffffff, sum, offset); - __shared__ float sdata[32]; - int lane = threadIdx.x % 32, warp = threadIdx.x / 32; - if (lane == 0) sdata[warp] = sum; - __syncthreads(); - if (warp == 0) { - sum = (lane < (blockDim.x + 31) / 32) ? sdata[lane] : 0.0f; - for (int offset = 16; offset > 0; offset >>= 1) - sum += __shfl_down_sync(0xffffffff, sum, offset); - if (lane == 0) bgrad[d] = from_float(sum); - } -} - -// NCHW bias grad: sum over (B, OH, OW) for each OC channel -__global__ void n3_conv_bias_grad_nchw( - precision_t* __restrict__ bgrad, const precision_t* __restrict__ grad, - int B, int OC, int spatial) { - int oc = blockIdx.x; - if (oc >= OC) return; - float sum = 0.0f; - int total = B * spatial; - for (int i = threadIdx.x; i < total; i += blockDim.x) { - int b = i / spatial, s = i % spatial; - sum += to_float(grad[b * OC * spatial + oc * spatial + s]); - } - for (int offset = 16; offset > 0; offset >>= 1) - sum += __shfl_down_sync(0xffffffff, sum, offset); - __shared__ float sdata[32]; - int lane = threadIdx.x % 32, warp = threadIdx.x / 32; - if (lane == 0) sdata[warp] = sum; - __syncthreads(); - if (warp == 0) { - sum = (lane < (blockDim.x + 31) / 32) ? sdata[lane] : 0.0f; - for (int offset = 16; offset > 0; offset >>= 1) - sum += __shfl_down_sync(0xffffffff, sum, offset); - if (lane == 0) bgrad[oc] = from_float(sum); - } -} - -__global__ void n3_concat_backward_conv_kernel( - precision_t* __restrict__ conv_grad, const precision_t* __restrict__ concat_grad, int B) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= B * N3_CONV_FLAT) return; - int b = idx / N3_CONV_FLAT, c = idx % N3_CONV_FLAT; - conv_grad[b * N3_CONV_FLAT + c] = concat_grad[b * N3_CONCAT + c]; -} - -// Embedding backward: scatter-add grad from concat_grad's player_embed region -// into embed_wgrad (float accumulation buffer). -// Each (b, f) looked up row obs[b, MAP_SIZE+f] from the table. -__global__ void n3_embedding_backward_kernel( - float* __restrict__ embed_wgrad_f, - const precision_t* __restrict__ concat_grad, - const precision_t* __restrict__ obs, - int B, int obs_size) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= B * N3_PLAYER * N3_EMBED_DIM) return; - int b = idx / (N3_PLAYER * N3_EMBED_DIM); - int rem = idx % (N3_PLAYER * N3_EMBED_DIM); - int f = rem / N3_EMBED_DIM; - int d = rem % N3_EMBED_DIM; - int val = (int)to_float(obs[b * obs_size + N3_MAP_SIZE + f]); - float g = to_float(concat_grad[b * N3_CONCAT + N3_CONV_FLAT + f * N3_EMBED_DIM + d]); - atomicAdd(&embed_wgrad_f[val * N3_EMBED_DIM + d], g); -} - -// Cast float buffer to precision_t -__global__ void n3_float_to_precision_kernel( - precision_t* __restrict__ dst, const float* __restrict__ src, int n) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < n) dst[idx] = from_float(src[idx]); -} - -// ---- atomicAdd for precision_t ---- -#ifdef PRECISION_FLOAT -__device__ __forceinline__ void atomicAdd_precision(precision_t* addr, precision_t val) { - atomicAdd(addr, val); -} -#else -__device__ __forceinline__ void atomicAdd_precision(precision_t* addr, precision_t val) { - // bf16 atomicAdd via CAS on enclosing 32-bit word - unsigned int* addr_u32 = (unsigned int*)((size_t)addr & ~2ULL); - bool is_high = ((size_t)addr & 2) != 0; - unsigned int old_u32 = *addr_u32, assumed; - do { - assumed = old_u32; - __nv_bfloat16* pair = (__nv_bfloat16*)&old_u32; - float sum = __bfloat162float(pair[is_high]) + __bfloat162float(val); - unsigned int new_u32 = assumed; - ((__nv_bfloat16*)&new_u32)[is_high] = __float2bfloat16(sum); - old_u32 = atomicCAS(addr_u32, assumed, new_u32); - } while (old_u32 != assumed); -} -#endif - -// ---- NCHW bias kernels for im2col conv path ---- - -__global__ void conv_bias_kernel(precision_t* __restrict__ data, - const precision_t* __restrict__ bias, int B, int OC, int spatial) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = B * OC * spatial; - if (idx >= total) return; - int oc = (idx / spatial) % OC; - data[idx] = from_float(to_float(data[idx]) + to_float(bias[oc])); -} - -__global__ void conv_bias_relu_kernel(precision_t* __restrict__ data, - const precision_t* __restrict__ bias, int B, int OC, int spatial) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = B * OC * spatial; - if (idx >= total) return; - int oc = (idx / spatial) % OC; - data[idx] = from_float(fmaxf(0.0f, to_float(data[idx]) + to_float(bias[oc]))); -} - -// ---- im2col + cuBLAS conv (no cuDNN) ---- -// NCHW layout throughout. Weight stored as (OC, IC*K*K). -// im2col produces (B*OH*OW, IC*K*K), matmul with W^T gives (B*OH*OW, OC), -// then reshape to NCHW (B, OC, OH, OW). - -__global__ void im2col_kernel( - const precision_t* __restrict__ input, precision_t* __restrict__ col, - int B, int IC, int IH, int IW, int K, int S, int OH, int OW -) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = B * OH * OW * IC * K * K; - if (idx >= total) return; - int col_w = IC * K * K; - int row = idx / col_w; - int c = idx % col_w; - int b = row / (OH * OW); - int rem = row % (OH * OW); - int oh = rem / OW, ow = rem % OW; - int ic = c / (K * K), kk = c % (K * K); - int kh = kk / K, kw = kk % K; - int ih = oh * S + kh, iw = ow * S + kw; - col[idx] = input[b * IC * IH * IW + ic * IH * IW + ih * IW + iw]; -} - -// Backward: col2im — input-centric gather to avoid atomics. -// Each thread owns one (b, ic, ih, iw) element and sums contributions from all -// (oh, ow, kh, kw) patches that map to it. -__global__ void col2im_kernel( - const precision_t* __restrict__ col, precision_t* __restrict__ grad_input, - int B, int IC, int IH, int IW, int K, int S, int OH, int OW -) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = B * IC * IH * IW; - if (idx >= total) return; - int iw = idx % IW; - int ih = (idx / IW) % IH; - int ic = (idx / (IW * IH)) % IC; - int b = idx / (IW * IH * IC); - float sum = 0.0f; - for (int kh = 0; kh < K; kh++) { - int ih_off = ih - kh; - if (ih_off < 0 || ih_off % S != 0) continue; - int oh = ih_off / S; - if (oh >= OH) continue; - for (int kw = 0; kw < K; kw++) { - int iw_off = iw - kw; - if (iw_off < 0 || iw_off % S != 0) continue; - int ow = iw_off / S; - if (ow >= OW) continue; - int col_idx = (b * OH * OW + oh * OW + ow) * (IC * K * K) + ic * K * K + kh * K + kw; - sum += to_float(col[col_idx]); - } - } - grad_input[idx] = from_float(sum); -} - -// Transpose (B, OC, OH, OW) -> (B*OH*OW, OC) [NCHW to row-major spatial-first] -__global__ void nchw_to_rows_kernel( - const precision_t* __restrict__ src, precision_t* __restrict__ dst, - int B, int OC, int spatial -) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = B * OC * spatial; - if (idx >= total) return; - int b = idx / (OC * spatial); - int oc = (idx / spatial) % OC; - int s = idx % spatial; - dst[(b * spatial + s) * OC + oc] = src[idx]; -} - -// Transpose (B*OH*OW, OC) -> (B, OC, OH, OW) [row-major spatial-first to NCHW] -__global__ void rows_to_nchw_kernel( - const precision_t* __restrict__ src, precision_t* __restrict__ dst, - int B, int OC, int spatial -) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = B * OC * spatial; - if (idx >= total) return; - int b = idx / (OC * spatial); - int oc = (idx / spatial) % OC; - int s = idx % spatial; - dst[idx] = src[(b * spatial + s) * OC + oc]; -} - -// Forward: im2col conv + bias + optional relu. All NCHW. -// col_buf: pre-allocated (max_B * OH * OW, IC * K * K) -// mm_buf: pre-allocated (max_B * OH * OW, OC) — row-major (spatial-first) -static void gemm_conv_forward( - PrecisionTensor* weight, PrecisionTensor* bias, - precision_t* input, precision_t* output, - precision_t* col_buf, precision_t* mm_buf, - int B, int IC, int IH, int IW, int OC, int K, int S, int OH, int OW, - bool relu, cudaStream_t stream -) { - int col_rows = B * OH * OW; - int col_cols = IC * K * K; - int total_col = col_rows * col_cols; - int total_out = B * OC * OH * OW; - - // im2col: input NCHW -> col (B*OH*OW, IC*K*K) - im2col_kernel<<>>( - input, col_buf, B, IC, IH, IW, K, S, OH, OW); - - // matmul: col (B*OH*OW, IC*K*K) @ W^T (IC*K*K, OC) = mm_buf (B*OH*OW, OC) - PrecisionTensor col_t = {.data = col_buf, .shape = {col_rows, col_cols}}; - PrecisionTensor mm_t = {.data = mm_buf, .shape = {col_rows, OC}}; - puf_mm(&col_t, weight, &mm_t, stream); - - // transpose (B*OH*OW, OC) -> (B, OC, OH, OW) NCHW + bias + relu - int spatial = OH * OW; - rows_to_nchw_kernel<<>>( - mm_buf, output, B, OC, spatial); - if (relu) { - conv_bias_relu_kernel<<>>( - output, bias->data, B, OC, spatial); - } else { - conv_bias_kernel<<>>( - output, bias->data, B, OC, spatial); - } -} - -// Backward: weight grad + optional input grad via im2col/col2im + cuBLAS. -// grad_output is NCHW (B, OC, OH, OW). saved_input is NCHW. -// Caller handles relu backward and bias grad (same as cuDNN path). -static void gemm_conv_backward( - PrecisionTensor* weight, - precision_t* saved_input, precision_t* grad_output, - precision_t* wgrad, precision_t* input_grad, - precision_t* col_buf, precision_t* mm_buf, - int B, int IC, int IH, int IW, int OC, int K, int S, int OH, int OW, - cudaStream_t stream -) { - int col_rows = B * OH * OW; - int col_cols = IC * K * K; - int total_col = col_rows * col_cols; - int total_out = B * OC * OH * OW; - int spatial = OH * OW; - - // Transpose grad_output NCHW -> (B*OH*OW, OC) - nchw_to_rows_kernel<<>>( - grad_output, mm_buf, B, OC, spatial); - - // im2col of saved_input - im2col_kernel<<>>( - saved_input, col_buf, B, IC, IH, IW, K, S, OH, OW); - - // Weight grad: mm_buf^T (OC, B*OH*OW) @ col_buf (B*OH*OW, IC*K*K) = wgrad (OC, IC*K*K) - PrecisionTensor mm_t = {.data = mm_buf, .shape = {col_rows, OC}}; - PrecisionTensor col_t = {.data = col_buf, .shape = {col_rows, col_cols}}; - PrecisionTensor wg_t = {.data = wgrad, .shape = {OC, col_cols}}; - puf_mm_tn(&mm_t, &col_t, &wg_t, stream); - - // Input grad (optional): mm_buf (B*OH*OW, OC) @ weight (OC, IC*K*K) = col_grad (B*OH*OW, IC*K*K) - if (input_grad) { - puf_mm_nn(&mm_t, weight, &col_t, stream); // reuse col_buf as col_grad - col2im_kernel<<>>( - col_buf, input_grad, B, IC, IH, IW, K, S, OH, OW); - } -} - -// ---- NMMO3 encoder structs ---- - -struct NMMO3EncoderWeights { - ConvWeights conv1, conv2; - PrecisionTensor embed_w, proj_w, proj_b; - int obs_size, hidden; -}; - -struct NMMO3EncoderActivations { - ConvActivations conv1, conv2; - PrecisionTensor col1, mm1, col2, mm2; // im2col + matmul scratch buffers - PrecisionTensor multihot, embed_out, concat, out, saved_obs; - PrecisionTensor embed_wgrad, proj_wgrad, proj_bgrad; - FloatTensor embed_wgrad_f; // float accumulation buffer for scatter-add -}; - -static NMMO3EncoderWeights* nmmo3_encoder_create(int obs_size, int hidden) { - NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)calloc(1, sizeof(NMMO3EncoderWeights)); - ew->obs_size = obs_size; ew->hidden = hidden; - conv_init(&ew->conv1, N3_C1_IC, N3_C1_OC, N3_C1_K, N3_C1_S, N3_MAP_H, N3_MAP_W, true); - conv_init(&ew->conv2, N3_C2_IC, N3_C2_OC, N3_C2_K, N3_C2_S, N3_C1_OH, N3_C1_OW, false); - return ew; -} - -// ---- NMMO3 encoder interface ---- - -static PrecisionTensor nmmo3_encoder_forward(void* w, void* activations, PrecisionTensor input, cudaStream_t stream) { - NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; - NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; - int B = input.shape[0]; - - if (a->saved_obs.data) puf_copy(&a->saved_obs, &input, stream); - - cudaMemsetAsync(a->multihot.data, 0, (int64_t)B * N3_MULTIHOT * N3_MAP_H * N3_MAP_W * sizeof(precision_t), stream); - n3_multihot_kernel<<>>( - a->multihot.data, input.data, B, ew->obs_size); - - gemm_conv_forward(&ew->conv1.w, &ew->conv1.b, a->multihot.data, a->conv1.out.data, - a->col1.data, a->mm1.data, B, N3_C1_IC, N3_MAP_H, N3_MAP_W, - N3_C1_OC, N3_C1_K, N3_C1_S, N3_C1_OH, N3_C1_OW, true, stream); - if (a->conv1.saved_input.data) - cudaMemcpyAsync(a->conv1.saved_input.data, a->multihot.data, - (int64_t)B * N3_C1_IC * N3_MAP_H * N3_MAP_W * sizeof(precision_t), cudaMemcpyDeviceToDevice, stream); - gemm_conv_forward(&ew->conv2.w, &ew->conv2.b, a->conv1.out.data, a->conv2.out.data, - a->col2.data, a->mm2.data, B, N3_C2_IC, N3_C1_OH, N3_C1_OW, - N3_C2_OC, N3_C2_K, N3_C2_S, N3_C2_OH, N3_C2_OW, false, stream); - if (a->conv2.saved_input.data) - cudaMemcpyAsync(a->conv2.saved_input.data, a->conv1.out.data, - (int64_t)B * N3_C2_IC * N3_C1_OH * N3_C1_OW * sizeof(precision_t), cudaMemcpyDeviceToDevice, stream); - - n3_embedding_kernel<<>>( - a->embed_out.data, input.data, ew->embed_w.data, B, ew->obs_size); - n3_concat_kernel<<>>( - a->concat.data, a->conv2.out.data, a->embed_out.data, input.data, B, ew->obs_size); - - puf_mm(&a->concat, &ew->proj_w, &a->out, stream); - n3_bias_relu_kernel<<hidden), BLOCK_SIZE, 0, stream>>>( - a->out.data, ew->proj_b.data, B * ew->hidden, ew->hidden); - return a->out; -} - -static void nmmo3_encoder_backward(void* w, void* activations, PrecisionTensor grad, cudaStream_t stream) { - NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; - NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; - int B = grad.shape[0], H = ew->hidden; - - n3_relu_backward_kernel<<>>( - grad.data, a->out.data, B * H); - bias_grad_kernel<<>>( - a->proj_bgrad.data, grad.data, B, H); - puf_mm_tn(&grad, &a->concat, &a->proj_wgrad, stream); - - PrecisionTensor grad_concat = {.data = a->concat.data, .shape = {B, N3_CONCAT}}; - puf_mm_nn(&grad, &ew->proj_w, &grad_concat, stream); - - n3_concat_backward_conv_kernel<<>>( - a->conv2.grad.data, grad_concat.data, B); - - n3_conv_bias_grad_nchw<<conv2.OC, 256, 0, stream>>>( - a->conv2.bgrad.data, a->conv2.grad.data, - B, ew->conv2.OC, ew->conv2.OH * ew->conv2.OW); - gemm_conv_backward(&ew->conv2.w, a->conv2.saved_input.data, a->conv2.grad.data, - a->conv2.wgrad.data, a->conv1.grad.data, - a->col2.data, a->mm2.data, B, N3_C2_IC, N3_C1_OH, N3_C1_OW, - N3_C2_OC, N3_C2_K, N3_C2_S, N3_C2_OH, N3_C2_OW, stream); - - n3_relu_backward_kernel<<conv1.OC * ew->conv1.OH * ew->conv1.OW), BLOCK_SIZE, 0, stream>>>( - a->conv1.grad.data, a->conv1.out.data, - B * ew->conv1.OC * ew->conv1.OH * ew->conv1.OW); - n3_conv_bias_grad_nchw<<conv1.OC, 256, 0, stream>>>( - a->conv1.bgrad.data, a->conv1.grad.data, - B, ew->conv1.OC, ew->conv1.OH * ew->conv1.OW); - gemm_conv_backward(&ew->conv1.w, a->conv1.saved_input.data, a->conv1.grad.data, - a->conv1.wgrad.data, NULL, - a->col1.data, a->mm1.data, B, N3_C1_IC, N3_MAP_H, N3_MAP_W, - N3_C1_OC, N3_C1_K, N3_C1_S, N3_C1_OH, N3_C1_OW, stream); - - // Embedding backward: scatter-add from concat gradient into float buffer, then cast - int embed_n = N3_EMBED_VOCAB * N3_EMBED_DIM; - cudaMemsetAsync(a->embed_wgrad_f.data, 0, embed_n * sizeof(float), stream); - n3_embedding_backward_kernel<<>>( - a->embed_wgrad_f.data, grad_concat.data, a->saved_obs.data, B, ew->obs_size); - n3_float_to_precision_kernel<<>>( - a->embed_wgrad.data, a->embed_wgrad_f.data, embed_n); -} - -static void nmmo3_encoder_init_weights(void* w, uint64_t* seed, cudaStream_t stream) { - NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; - conv_init_weights(&ew->conv1, seed, stream); - conv_init_weights(&ew->conv2, seed, stream); - auto init2d = [&](PrecisionTensor& t, int rows, int cols, float gain) { - PrecisionTensor wt = {.data = t.data, .shape = {rows, cols}}; - puf_kaiming_init(&wt, gain, (*seed)++, stream); - }; - puf_normal_init(&ew->embed_w, 1.0f, (*seed)++, stream); - init2d(ew->proj_w, ew->hidden, N3_CONCAT, 1.0f); - cudaMemsetAsync(ew->proj_b.data, 0, numel(ew->proj_b.shape) * sizeof(precision_t), stream); -} - -static void nmmo3_encoder_reg_params(void* w, Allocator* alloc) { - NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; - conv_reg_params(&ew->conv1, alloc); - conv_reg_params(&ew->conv2, alloc); - ew->embed_w = {.shape = {N3_EMBED_VOCAB, N3_EMBED_DIM}}; - ew->proj_w = {.shape = {ew->hidden, N3_CONCAT}}; - ew->proj_b = {.shape = {ew->hidden}}; - alloc_register(alloc,&ew->embed_w); - alloc_register(alloc,&ew->proj_w); alloc_register(alloc,&ew->proj_b); -} - -static void nmmo3_encoder_reg_train(void* w, void* activations, Allocator* acts, Allocator* grads, int B_TT) { - NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; - NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; - *a = {}; - a->multihot = {.shape = {B_TT, N3_MULTIHOT * N3_MAP_H * N3_MAP_W}}; - alloc_register(acts,&a->multihot); - // Conv1 buffers - a->conv1.out = {.shape = {B_TT * N3_C1_OC * N3_C1_OH * N3_C1_OW}}; - a->conv1.grad = {.shape = {B_TT * N3_C1_OC * N3_C1_OH * N3_C1_OW}}; - a->conv1.saved_input = {.shape = {B_TT * N3_C1_IC * N3_MAP_H * N3_MAP_W}}; - a->conv1.wgrad = {.shape = {N3_C1_OC, N3_C1_IC * N3_C1_K * N3_C1_K}}; - a->conv1.bgrad = {.shape = {N3_C1_OC}}; - alloc_register(acts,&a->conv1.out); alloc_register(acts,&a->conv1.grad); alloc_register(acts,&a->conv1.saved_input); - alloc_register(grads,&a->conv1.wgrad); alloc_register(grads,&a->conv1.bgrad); - a->col1 = {.shape = {B_TT * N3_C1_OH * N3_C1_OW, N3_C1_IC * N3_C1_K * N3_C1_K}}; - a->mm1 = {.shape = {B_TT * N3_C1_OH * N3_C1_OW, N3_C1_OC}}; - alloc_register(acts,&a->col1); alloc_register(acts,&a->mm1); - // Conv2 buffers - a->conv2.out = {.shape = {B_TT * N3_C2_OC * N3_C2_OH * N3_C2_OW}}; - a->conv2.grad = {.shape = {B_TT * N3_C2_OC * N3_C2_OH * N3_C2_OW}}; - a->conv2.saved_input = {.shape = {B_TT * N3_C2_IC * N3_C1_OH * N3_C1_OW}}; - a->conv2.wgrad = {.shape = {N3_C2_OC, N3_C2_IC * N3_C2_K * N3_C2_K}}; - a->conv2.bgrad = {.shape = {N3_C2_OC}}; - alloc_register(acts,&a->conv2.out); alloc_register(acts,&a->conv2.grad); alloc_register(acts,&a->conv2.saved_input); - alloc_register(grads,&a->conv2.wgrad); alloc_register(grads,&a->conv2.bgrad); - a->col2 = {.shape = {B_TT * N3_C2_OH * N3_C2_OW, N3_C2_IC * N3_C2_K * N3_C2_K}}; - a->mm2 = {.shape = {B_TT * N3_C2_OH * N3_C2_OW, N3_C2_OC}}; - alloc_register(acts,&a->col2); alloc_register(acts,&a->mm2); - a->embed_out = {.shape = {B_TT, N3_PLAYER_EMBED}}; - a->concat = {.shape = {B_TT, N3_CONCAT}}; - a->out = {.shape = {B_TT, ew->hidden}}; - a->saved_obs = {.shape = {B_TT, ew->obs_size}}; - alloc_register(acts,&a->embed_out); alloc_register(acts,&a->concat); - alloc_register(acts,&a->out); alloc_register(acts,&a->saved_obs); - a->embed_wgrad = {.shape = {N3_EMBED_VOCAB, N3_EMBED_DIM}}; - a->embed_wgrad_f = {.shape = {N3_EMBED_VOCAB, N3_EMBED_DIM}}; - a->proj_wgrad = {.shape = {ew->hidden, N3_CONCAT}}; - a->proj_bgrad = {.shape = {ew->hidden}}; - alloc_register(grads,&a->embed_wgrad); - alloc_register(acts,&a->embed_wgrad_f); - alloc_register(grads,&a->proj_wgrad); alloc_register(grads,&a->proj_bgrad); -} - -static void nmmo3_encoder_reg_rollout(void* w, void* activations, Allocator* alloc, int B) { - NMMO3EncoderWeights* ew = (NMMO3EncoderWeights*)w; - NMMO3EncoderActivations* a = (NMMO3EncoderActivations*)activations; - a->multihot = {.shape = {B, N3_MULTIHOT * N3_MAP_H * N3_MAP_W}}; - alloc_register(alloc,&a->multihot); - a->conv1.out = {.shape = {B * N3_C1_OC * N3_C1_OH * N3_C1_OW}}; - alloc_register(alloc,&a->conv1.out); - a->col1 = {.shape = {B * N3_C1_OH * N3_C1_OW, N3_C1_IC * N3_C1_K * N3_C1_K}}; - a->mm1 = {.shape = {B * N3_C1_OH * N3_C1_OW, N3_C1_OC}}; - alloc_register(alloc,&a->col1); alloc_register(alloc,&a->mm1); - a->conv2.out = {.shape = {B * N3_C2_OC * N3_C2_OH * N3_C2_OW}}; - alloc_register(alloc,&a->conv2.out); - a->col2 = {.shape = {B * N3_C2_OH * N3_C2_OW, N3_C2_IC * N3_C2_K * N3_C2_K}}; - a->mm2 = {.shape = {B * N3_C2_OH * N3_C2_OW, N3_C2_OC}}; - alloc_register(alloc,&a->col2); alloc_register(alloc,&a->mm2); - a->embed_out = {.shape = {B, N3_PLAYER_EMBED}}; - a->concat = {.shape = {B, N3_CONCAT}}; - a->out = {.shape = {B, ew->hidden}}; - alloc_register(alloc,&a->embed_out); alloc_register(alloc,&a->concat); alloc_register(alloc,&a->out); -} - -static void* nmmo3_encoder_create_weights(void* self) { - Encoder* e = (Encoder*)self; - return nmmo3_encoder_create(e->in_dim, e->out_dim); -} -static void nmmo3_encoder_free_weights(void* weights) { free(weights); } -static void nmmo3_encoder_free_activations(void* activations) { free(activations); } +#include "nmmo3.cu" +#include "craftax.cu" // Override encoder vtable for known ocean environments. No-op for unknown envs. static void create_custom_encoder(const std::string& env_name, Encoder* enc) { + if (env_name == "craftax") { + *enc = Encoder{ + .forward = craftax_encoder_forward, + .backward = craftax_encoder_backward, + .init_weights = craftax_encoder_init_weights, + .reg_params = craftax_encoder_reg_params, + .reg_train = craftax_encoder_reg_train, + .reg_rollout = craftax_encoder_reg_rollout, + .create_weights = craftax_encoder_create_weights, + .free_weights = craftax_encoder_free_weights, + .free_activations = craftax_encoder_free_activations, + .in_dim = enc->in_dim, .out_dim = enc->out_dim, + .activation_size = sizeof(CraftaxEncoderActivations), + }; + return; + } if (env_name == "nmmo3") { *enc = Encoder{ .forward = nmmo3_encoder_forward,