diff --git a/config/nethack.ini b/config/nethack.ini index 5df02d918f..cfaea2c7b9 100644 --- a/config/nethack.ini +++ b/config/nethack.ini @@ -7,21 +7,16 @@ num_buffers = 4 num_threads = 32 [env] -gold_coef = 0.00267383433 -exp_coef = 0.112984791 -descent_coef = 0.762886882 -# pays once per new unique (dnum, dlevel) floor entered (branches count) -floor_coef = 0.0 -xp_coef = 0.2399306 -scout_coef = 0.0104040429 -ac_coef = 0.0 -death_penalty = -0.324047506 +gold_coef = 0.096887961 +score_coef = 0 +exp_coef = 0.041234985 +descent_coef = 1.0 +floor_coef = 0.38962388 +xp_coef = 0.741106272 +scout_coef = 0.000813971157 +death_penalty = -0.286484599 # flags -# scout reward scales with xp level relative to depth; <= 0 disables -scout_ready = 0.0 -# fraction of protection-spell AC excluded from the ac reward; 1 = armor only -ac_nospell = 1.0 # 1 removes the verb from the action space mask_search20 = 0 mask_run = 0 @@ -35,20 +30,20 @@ num_layers = 3 [train] total_timesteps = 1_000_000_000 minibatch_size = 4096 -horizon = 128 -learning_rate = 0.000765701407 -momentum = 0.987343371 +horizon = 256 +learning_rate = 0.00154131022 +momentum = 0.956484258 verb_eps = 0.005 verb_eps_anneal_start = 0.4 verb_eps_anneal_end = 1.0 -gamma = 0.998305142 -gae_lambda = 0.265489638 -replay_ratio = 1.68762803 -clip_coef = 0.100000001 -vf_coef = 4 -vf_clip_coef = 0.00999999978 -max_grad_norm = 0.961730838 -ent_coef = 0.0229146201 +gamma = 0.999329984 +gae_lambda = 0.199999988 +replay_ratio = 1.20105004 +clip_coef = 0.36599344 +vf_coef = 0.882273436 +vf_clip_coef = 1.30713189 +max_grad_norm = 0.335048616 +ent_coef = 0.00555205066 anneal_lr = 1 anneal_ent_coef = 0 min_lr_ratio = 0.0 @@ -67,8 +62,14 @@ scale = time [sweep.env.exp_coef] distribution = log_normal -min = 0.01 -max = 0.2 +min = 0.001 +max = 0.1 +scale = auto + +[sweep.env.score_coef] +distribution = log_normal +min = 0.001 +max = 0.1 scale = auto [sweep.env.gold_coef] @@ -91,7 +92,7 @@ scale = auto [sweep.env.scout_coef] distribution = log_normal -min = 0.001 +min = 0.0001 max = 0.1 scale = auto @@ -113,12 +114,6 @@ min = 0.2 max = 0.8 scale = auto -[sweep.env.ac_coef] -distribution = log_normal -min = 0.01 -max = 2.0 -scale = auto - [sweep.env.floor_coef] distribution = uniform min = 0.0 diff --git a/ocean/nethack/depth.ini b/ocean/nethack/depth.ini index 03cba7becc..3e752367e4 100644 --- a/ocean/nethack/depth.ini +++ b/ocean/nethack/depth.ini @@ -1,7 +1,8 @@ -# async depth config: promoted winner of the max_depth sweep (run 0306, -# tail depth 11.0, durable AC best-in-cohort). Rewards are descent-dominated -# with a per-floor bonus; SEARCH20 and RUN are masked. -# To use: copy to config/depth.ini, then ./puffer train depth +# Depth recipe: sweep-0306 winner (tail depth 11.0, async, RNN carry) with +# challenge multi-role. ~366M steps. SEARCH20 and RUN are masked. +# Use: copy over config/nethack.ini, then ./build.sh nethack && ./puffer train +# NOTE horizon*num_buffers > total_agents: runs are single-draw (the async +# multi-stream cuBLAS nondeterminism); use num_buffers=1 for bit-repro. [base] env_name = nethack reset_every_horizon = 0 @@ -11,29 +12,18 @@ async = 1 total_agents = 512 num_buffers = 4 num_threads = 32 -action_mask_size = 734 [env] gold_coef = 0.00100000005 exp_coef = 0.00999999978 descent_coef = 1.0 -# pays once per new unique (dnum, dlevel) floor entered (branches count) floor_coef = 0.727649987 xp_coef = 0.33030045 scout_coef = 0.0165148061 -ac_coef = 0.0 death_penalty = -0.0327779055 - -# flags -# scout reward scales with xp level relative to depth; <= 0 disables -scout_ready = 0.0 -# fraction of protection-spell AC excluded from the ac reward; 1 = armor only -ac_nospell = 1.0 -# 1 removes the verb from the action space mask_search20 = 1 mask_run = 1 -# 1 = random role/race/gender/align per reset (challenge protocol) -multi_role = 0 +multi_role = 1 [policy] hidden_size = 512 @@ -62,60 +52,3 @@ min_lr_ratio = 0.0 min_ent_coef_ratio = 0.1 vtrace_rho_clip = 1.0 vtrace_c_clip = 1.0 - -[sweep] -metric = max_depth - -[sweep.train.total_timesteps] -distribution = log_normal -min = 2e8 -max = 1e9 -scale = time - -[sweep.env.exp_coef] -distribution = log_normal -min = 0.01 -max = 1.0 -scale = auto - -[sweep.env.gold_coef] -distribution = log_normal -min = 0.001 -max = 1.0 -scale = auto - -[sweep.env.descent_coef] -distribution = uniform -min = 0.0 -max = 1.0 -scale = auto - -[sweep.env.floor_coef] -distribution = uniform -min = 0.0 -max = 1.0 -scale = auto - -[sweep.env.xp_coef] -distribution = uniform -min = 0.0 -max = 1.0 -scale = auto - -[sweep.env.scout_coef] -distribution = log_normal -min = 0.001 -max = 1.0 -scale = auto - -[sweep.env.death_penalty] -distribution = uniform -min = -1.0 -max = 0.0 -scale = auto - -[sweep.env.ac_coef] -distribution = log_normal -min = 0.01 -max = 2.0 -scale = auto diff --git a/ocean/nethack/depth_multirole.ini b/ocean/nethack/depth_multirole.ini deleted file mode 100644 index 974508e5a2..0000000000 --- a/ocean/nethack/depth_multirole.ini +++ /dev/null @@ -1,56 +0,0 @@ -# depth.ini (sweep 0306, tail depth 11.0, async, RNN carry) with challenge -# multi-role. ~366M steps; Fin's single-role 400M depth runs were 33-43 min. -# Use: copy over config/nethack.ini then ./build.sh nethack && ./puffer train -[base] -env_name = nethack -reset_every_horizon = 0 -async = 1 - -[vec] -total_agents = 512 -num_buffers = 4 -num_threads = 32 -action_mask_size = 734 - -[env] -gold_coef = 0.00100000005 -exp_coef = 0.00999999978 -descent_coef = 1.0 -floor_coef = 0.727649987 -xp_coef = 0.33030045 -scout_coef = 0.0165148061 -ac_coef = 0.0 -death_penalty = -0.0327779055 -scout_ready = 0.0 -ac_nospell = 1.0 -mask_search20 = 1 -mask_run = 1 -multi_role = 1 - -[policy] -hidden_size = 512 -num_layers = 3 - -[train] -total_timesteps = 366129632 -minibatch_size = 8192 -horizon = 512 -learning_rate = 0.00112607062 -momentum = 0.98707521 -verb_eps = 0.005 -verb_eps_anneal_start = 0.4 -verb_eps_anneal_end = 1.0 -gamma = 0.9966259 -gae_lambda = 0.83733207 -replay_ratio = 1.82700419 -clip_coef = 0.245328903 -vf_coef = 3.35665774 -vf_clip_coef = 2.18958735 -max_grad_norm = 1.28095329 -ent_coef = 0.00987035036 -anneal_lr = 1 -anneal_ent_coef = 0 -min_lr_ratio = 0.0 -min_ent_coef_ratio = 0.1 -vtrace_rho_clip = 1.0 -vtrace_c_clip = 1.0 diff --git a/ocean/nethack/nethack.c b/ocean/nethack/nethack.c index 8b4d97b359..b2592cbf43 100644 --- a/ocean/nethack/nethack.c +++ b/ocean/nethack/nethack.c @@ -1,6 +1,9 @@ #include #include #include +#ifdef __EMSCRIPTEN__ +#include +#endif #include #include #include @@ -40,6 +43,59 @@ static void env_open(Nethack* env) { + NETHACK_SPELL_SLOTS, 1); a->rewards = (float*)calloc(1, sizeof(float)); a->terminals = (float*)calloc(1, sizeof(float)); + // NH_ROLE=arc (demo only): pin that role; race/gender/align stay random. + { + const char* s = getenv("NH_ROLE"); + if (s && s[0]) { + char buf[32]; + size_t n = 0; + for (const char* p = s; *p && n + 1 < sizeof(buf); p++) { + char c = *p; + if (c >= 'A' && c <= 'Z') { + c = (char)(c - 'A' + 'a'); + } + if (c == ' ' || c == '_' || c == '-') { + continue; + } + buf[n++] = c; + } + buf[n] = 0; + static const struct { const char* key; const char* role; } map[] = { + {"0", "archeologist"}, {"arc", "archeologist"}, + {"arch", "archeologist"}, {"archeologist", "archeologist"}, + {"archaeologist", "archeologist"}, {"archologist", "archeologist"}, + {"1", "barbarian"}, {"bar", "barbarian"}, {"barbarian", "barbarian"}, + {"2", "caveman"}, {"cav", "caveman"}, {"caveman", "caveman"}, + {"3", "healer"}, {"hea", "healer"}, {"healer", "healer"}, + {"4", "knight"}, {"kni", "knight"}, {"knight", "knight"}, + {"5", "monk"}, {"mon", "monk"}, {"monk", "monk"}, + {"6", "priest"}, {"pri", "priest"}, {"priest", "priest"}, + {"7", "rogue"}, {"rog", "rogue"}, {"rogue", "rogue"}, + {"8", "ranger"}, {"ran", "ranger"}, {"ranger", "ranger"}, + {"9", "samurai"}, {"sam", "samurai"}, {"samurai", "samurai"}, + {"10", "tourist"}, {"tou", "tourist"}, {"tourist", "tourist"}, + {"11", "valkyrie"}, {"val", "valkyrie"}, {"valkyrie", "valkyrie"}, + {"12", "wizard"}, {"wiz", "wizard"}, {"wizard", "wizard"}, + }; + const char* role = NULL; + for (size_t i = 0; i < sizeof(map) / sizeof(map[0]); i++) { + if (strcmp(buf, map[i].key) == 0) { + role = map[i].role; + break; + } + } + if (!role) { + fprintf(stderr, "nethack: unknown NH_ROLE=%s (try arc, bar, cav, ...)\n", s); + exit(1); + } + static char opts[512]; + snprintf(opts, sizeof(opts), + "name:Agent,role:%s,race:random,gender:random,align:random," + NETHACK_OPTIONS_TAIL "!status_updates", role); + nethack_options_override = opts; + fprintf(stderr, "nethack demo: NH_ROLE=%s\n", role); + } + } init(env); // NH_MULTI=1: random role/race/gender/align per reset (challenge protocol) const char* mr = getenv("NH_MULTI"); @@ -67,7 +123,7 @@ static void env_close(Nethack* env) { #define DEMO_VOCAB 5977 #define DEMO_EMBED 32 #define DEMO_BL_FEAT (25 + 7 + 13 + NETHACK_NUM_ACTIONS + NETHACK_NUM_OCLASSES + 2 + 8 + 2 + 2 \ - + 1 + 2 + 20) + + 1 + 2 + 20 + 8) #define DEMO_SPKEY 16 #define DEMO_SPIN (DEMO_EMBED + 4) #define DEMO_INV_HID 16 // 16-dim slot rep: pool bottleneck + decoder key (unified) @@ -81,24 +137,24 @@ static void env_close(Nethack* env) { #define DEMO_QDIM ((DEMO_PTR_HEADS + 1) * DEMO_INV_HID) #define DEMO_DEC_LIN (NETHACK_NUM_ACTIONS + NETHACK_DIR_HEADS * NETHACK_NUM_DIRS + 1) #define DEMO_DEC_PAD ((DEMO_DEC_LIN + 7) / 8 * 8) -#define DEMO_LOC_IN (NETHACK_CROP_GRID * DEMO_EMBED) // 9x9 crop, per-cell embeds +#define DEMO_LOCC_DIM 8 +#define DEMO_LOCC_CLASSES 9 +#define DEMO_LOC_IN (NETHACK_CROP_GRID * DEMO_LOCC_DIM) // 9x9 semantic-class crop #define DEMO_LOC_HID 256 -#define DEMO_PW 5 -#define DEMO_PH 5 -#define DEMO_PX 16 -#define DEMO_PY 5 -#define DEMO_TOK (DEMO_PX * DEMO_PY) // 5x5 patches over 79x21 -#define DEMO_PCELLS (DEMO_PW * DEMO_PH) // off-map cells read the pad glyph -#define DEMO_P1 16 -#define DEMO_GLB_IN (DEMO_PCELLS * DEMO_EMBED) // per-patch flatten (glyph slice) -#define DEMO_GLB_HID 128 +#define DEMO_LOC_H1 256 +#define DEMO_TERRF 592 // 12 landmarks x 4 + 8 sectors x 4 bands x 17 classes +#define DEMO_TERR_H1 256 +#define DEMO_GLB_HID 128 // terrain branch output width // trigram message branch, mirroring NH_MSG_* in ocean/nethack/nethack.cu #define DEMO_MSG_LEN NETHACK_MSG_LEN -#define DEMO_MSG_VOCAB 4096 -#define DEMO_MSG_LOG2V 12 -#define DEMO_MSG_HID 32 -#define DEMO_MSG_CONCAT_OFF (DEMO_LOC_HID + DEMO_GLB_HID + DEMO_INV_POOL + 64 + DEMO_BL_FEAT) +#define DEMO_MSG_VOCAB 1024 +#define DEMO_MSG_LOG2V 10 +#define DEMO_MSG_HID 256 +#define DEMO_BLH 64 +#define DEMO_MSG_CONCAT_OFF (DEMO_LOC_HID + DEMO_GLB_HID + DEMO_BLH + DEMO_BL_FEAT) #define DEMO_SPELL_CONCAT_OFF (DEMO_MSG_CONCAT_OFF + DEMO_MSG_HID) +#define DEMO_SPM 64 +#define DEMO_SPELL_SLICE (2*DEMO_SPM + 4) // sum | max entity pools + doorstep scalars // identity-table channel (NH_ID_EMBED in nethack.cu); presence is inferred // per-checkpoint so one binary loads both eras #define DEMO_IDE_ROLE 16 @@ -107,8 +163,17 @@ static void env_close(Nethack* env) { #define DEMO_IDE_ALGN 8 #define DEMO_IDE_DIM (DEMO_IDE_ROLE + DEMO_IDE_RACE + DEMO_IDE_GEND + DEMO_IDE_ALGN) #define DEMO_IDE_NUMEL (13*DEMO_IDE_ROLE + 5*DEMO_IDE_RACE + 2*DEMO_IDE_GEND + 3*DEMO_IDE_ALGN) -#define DEMO_IDE_CONCAT_OFF (DEMO_SPELL_CONCAT_OFF + DEMO_SPKEY) -#define DEMO_CONCAT (DEMO_IDE_CONCAT_OFF + DEMO_IDE_DIM) // buffer max; live dim = net->concat_dim +#define DEMO_IDE_CONCAT_OFF (DEMO_SPELL_CONCAT_OFF + DEMO_SPELL_SLICE) +// V6-min streams (nethack.cu): tok48 -> rep16 -> MLP 64 -> sum|max + gates +#define DEMO_LABK 16 +#define DEMO_LAB_IN 48 +#define DEMO_MV 64 +#define DEMO_MINV_OFF (DEMO_IDE_CONCAT_OFF + DEMO_IDE_DIM) +#define DEMO_PASS_OFF (DEMO_MINV_OFF + 2*DEMO_MV) +#define DEMO_PACC 160 // accessory panel: [amulet | ring A | ring B | eyewear] x [r16|sfeat24] +#define DEMO_MLM_OFF (DEMO_PASS_OFF + 40 + 40 + DEMO_INV_HID + DEMO_PACC) +#define DEMO_MLI_OFF (DEMO_MLM_OFF + 2*DEMO_MV + DEMO_INV_HID) +#define DEMO_CONCAT (DEMO_MLI_OFF + 2*DEMO_MV + 2*DEMO_INV_HID) // per-blstat normalization, mirroring NH_BL_SCALE / NH_BL_ISLOG in ocean/nethack/nethack.cu static const float DEMO_BL_SCALE[27] = { @@ -125,18 +190,18 @@ typedef struct { float *embed; // (5977, 32) E_res float *ekind_w, *esub_w; // (14, 32), (944, 32) factor tables float *e_eff; // materialized E_res + E_kind + E_sub - float *loc_w, *loc_b; // (256, 2592), (256) - float *g1_w, *g1_xy, *g1_b; // (16, 800), (16, 2), (16): per-patch embed+flatten + hero dx,dy -> 16 - float *g2_w, *g2_b; // (128, 16), (128): 16 -> 128, maxed over tokens + float *loc_w, *loc_b; // (256, 648) over class-crop embeds + float *terr1_w, *terr1_b, *terr2_w, *terr2_b; // 592 -> 256 -> 128 terrain MLP + float *locc_w; // (9, 8) semantic-class embed table float *inv1_w, *inv1_b; // (16, 32), (16): per-slot features (pointer keys) float *inv1s_w; // (16, 24): gated item-state path into the slot MLP float *invt_w; // (16, 32): discovered-type channel (zero-init grown) - float *inv2_w, *inv2_b; // (128, 16), (128): pooled trunk summary (max over slots) + float *mv1_w, *mv1_b, *mv2_w, *mv2_b; // inventory deep pool 16->64->64 float *bl_w, *bl_b; // (64, DEMO_BL_FEAT), (64) float *proj_w, *proj_b; // (H, DEMO_CONCAT), (H) float *msg_w; // (4096, 32) trigram embedding table float *spk_w; // (16, 36) spell slot-rep projection - float *spk2_w, *spk2_b; // (16, 16), (16) spell pool (inv2 idiom) + float *spm1_w, *spm1_b, *spm2_w, *spm2_b; // spell per-slot MLP 16->64->64 float *ide_role_w, *ide_race_w, *ide_gend_w, *ide_algn_w; // identity tables (ide era) float *dec_lin; // (DEMO_DEC_PAD, H) bias-free; rows [26 verb | 48 dir | value], 75 used float *dec_q; // (DEMO_QDIM, H): thirteen stacked 16-dim queries (12 item + spell) @@ -145,11 +210,15 @@ typedef struct { MinGRU* mingru; Multidiscrete* md; int hidden_size, num_layers, num_actions; - int ide, concat_dim; // per-checkpoint layout (identity-table era or not) - float x[DEMO_LOC_IN]; // crop cell embeds, flattened - float px[DEMO_GLB_IN]; // one patch's cell embeds, flattened - float t16[DEMO_P1]; - float t128[DEMO_GLB_HID]; + float *loc2_w, *loc2_b; float loc_h1[DEMO_LOC_H1]; + float terr_tf[DEMO_TERRF]; float terr_h[DEMO_TERR_H1]; + float *mr_w, *mr_b, *mm1_w, *mm1_b, *mm2_w, *mm2_b; // monster stream + float *ir_w, *ir_b, *im1_w, *im1_b, *im2_w, *im2_b; // item stream + float x[DEMO_LOC_IN]; // crop class embeds, flattened + float wld[NETHACK_INV_SLOTS]; // per-slot wielded bit + float sfeat[NETHACK_INV_SLOTS * DEMO_SFEAT]; // per-slot state features + int otyp[NETHACK_INV_SLOTS]; // slot glyph -> otyp (-1 if not an object glyph) + int occ[NETHACK_INV_SLOTS]; // per-slot occupancy float slots[DEMO_INV_FLAT]; // per-slot post-relu features (decoder keys) float spkeys[NETHACK_SPELL_SLOTS * DEMO_SPKEY]; // relu'd spell slot reps float concat[DEMO_CONCAT]; // [local hid | global hid | inv pool | bl hidden | bl feats | msg] @@ -162,27 +231,34 @@ typedef struct { // All tensors land on 8-float boundaries; only tau (12) needs padding (+4). #define DEMO_ENC_FIXED (DEMO_VOCAB*DEMO_EMBED \ + NH_GM_NKIND*DEMO_EMBED + NH_GM_NSUB*DEMO_EMBED \ - + DEMO_LOC_HID*DEMO_LOC_IN + DEMO_LOC_HID \ - + DEMO_P1*DEMO_GLB_IN + DEMO_P1*2 + DEMO_P1 \ - + DEMO_GLB_HID*DEMO_P1 + DEMO_GLB_HID \ + + DEMO_LOC_H1*DEMO_LOC_IN + DEMO_LOC_H1 \ + + DEMO_LOC_HID*DEMO_LOC_H1 + DEMO_LOC_HID \ + + DEMO_TERR_H1*DEMO_TERRF + DEMO_TERR_H1 \ + + DEMO_GLB_HID*DEMO_TERR_H1 + DEMO_GLB_HID \ + + DEMO_LOCC_CLASSES*DEMO_LOCC_DIM \ + DEMO_INV_HID*DEMO_EMBED + DEMO_INV_HID \ + DEMO_INV_HID*DEMO_SFEAT \ + DEMO_INV_HID*DEMO_EMBED \ - + DEMO_INV_POOL*DEMO_INV_HID + DEMO_INV_POOL \ - + 64*DEMO_BL_FEAT + 64 \ + + DEMO_BLH*DEMO_BL_FEAT + DEMO_BLH \ + DEMO_MSG_VOCAB*DEMO_MSG_HID \ - + DEMO_SPKEY*DEMO_SPIN + DEMO_SPKEY*DEMO_SPKEY + DEMO_SPKEY) + + DEMO_SPKEY*DEMO_SPIN \ + + DEMO_SPM*DEMO_SPKEY + DEMO_SPM \ + + DEMO_SPM*DEMO_SPM + DEMO_SPM \ + + DEMO_IDE_NUMEL \ + + DEMO_MV*DEMO_INV_HID + DEMO_MV \ + + DEMO_MV*DEMO_MV + DEMO_MV \ + + 2*(DEMO_INV_HID*DEMO_LAB_IN + DEMO_INV_HID \ + + DEMO_MV*DEMO_INV_HID + DEMO_MV \ + + DEMO_MV*DEMO_MV + DEMO_MV)) #define DEMO_DEC_FIXED (DEMO_INV_HID*DEMO_INV_HID + 16) // k_w + tau padded 12->16 // ambiguities are possible; prefer the fewest layers (real configs have <= 8) -static int demo_infer_arch(int total, int* hidden, int* layers, int* actions, int* ide) { +static int demo_infer_arch(int total, int* hidden, int* layers, int* actions) { int best_l = 1 << 30; - for (int e = 0; e <= 1; e++) for (int H = 8; H <= 4096; H += 8) { - long rem = (long)total - DEMO_ENC_FIXED - DEMO_DEC_FIXED - (e ? DEMO_IDE_NUMEL : 0) - - (long)H * (DEMO_IDE_CONCAT_OFF + (e ? DEMO_IDE_DIM : 0) - + 1 + DEMO_DEC_PAD + DEMO_QDIM); + long rem = (long)total - DEMO_ENC_FIXED - DEMO_DEC_FIXED + - (long)H * (DEMO_CONCAT + 1 + DEMO_DEC_PAD + DEMO_QDIM); long per_layer = 3L * H * H; - if (rem <= 0) break; + if (rem <= 0) continue; if (rem % per_layer) continue; long L = rem / per_layer; if (L >= 1 && L < best_l) { @@ -190,7 +266,6 @@ static int demo_infer_arch(int total, int* hidden, int* layers, int* actions, in *hidden = H; *layers = (int)L; *actions = NETHACK_NUM_ACTIONS; - *ide = e; } } return best_l == 1 << 30 ? -1 : 0; @@ -199,46 +274,61 @@ static int demo_infer_arch(int total, int* hidden, int* layers, int* actions, in static NethackNet* make_nethack_net(Weights* w) { NethackNet* net = (NethackNet*)calloc(1, sizeof(NethackNet)); if (demo_infer_arch(w->size - 7, &net->hidden_size, &net->num_layers, - &net->num_actions, &net->ide) != 0) { + &net->num_actions) != 0) { fprintf(stderr, "nethack demo: cannot infer arch from %d floats — " "checkpoint is not a nethack policy with %d actions?\n", w->size - 7, NETHACK_NUM_ACTIONS); exit(1); } - net->concat_dim = DEMO_IDE_CONCAT_OFF + (net->ide ? DEMO_IDE_DIM : 0); - fprintf(stderr, "nethack demo: hidden=%d layers=%d actions=%d ide=%d (%d floats)\n", - net->hidden_size, net->num_layers, net->num_actions, net->ide, w->size - 7); + fprintf(stderr, "nethack demo: hidden=%d layers=%d actions=%d (%d floats)\n", + net->hidden_size, net->num_layers, net->num_actions, w->size - 7); net->hidden = (float*)calloc(net->hidden_size, sizeof(float)); net->embed = get_weights_aligned(w, DEMO_VOCAB * DEMO_EMBED); net->ekind_w = get_weights_aligned(w, NH_GM_NKIND * DEMO_EMBED); net->esub_w = get_weights_aligned(w, NH_GM_NSUB * DEMO_EMBED); - net->loc_w = get_weights_aligned(w, DEMO_LOC_HID * DEMO_LOC_IN); - net->loc_b = get_weights_aligned(w, DEMO_LOC_HID); - net->g1_w = get_weights_aligned(w, DEMO_P1 * DEMO_GLB_IN); - net->g1_xy = get_weights_aligned(w, DEMO_P1 * 2); - net->g1_b = get_weights_aligned(w, DEMO_P1); - net->g2_w = get_weights_aligned(w, DEMO_GLB_HID * DEMO_P1); - net->g2_b = get_weights_aligned(w, DEMO_GLB_HID); + net->loc_w = get_weights_aligned(w, DEMO_LOC_H1 * DEMO_LOC_IN); + net->loc_b = get_weights_aligned(w, DEMO_LOC_H1); + net->loc2_w = get_weights_aligned(w, DEMO_LOC_HID * DEMO_LOC_H1); + net->loc2_b = get_weights_aligned(w, DEMO_LOC_HID); + net->terr1_w = get_weights_aligned(w, DEMO_TERR_H1 * DEMO_TERRF); + net->terr1_b = get_weights_aligned(w, DEMO_TERR_H1); + net->terr2_w = get_weights_aligned(w, DEMO_GLB_HID * DEMO_TERR_H1); + net->terr2_b = get_weights_aligned(w, DEMO_GLB_HID); + net->locc_w = get_weights_aligned(w, DEMO_LOCC_CLASSES * DEMO_LOCC_DIM); net->inv1_w = get_weights_aligned(w, DEMO_INV_HID * DEMO_EMBED); net->inv1_b = get_weights_aligned(w, DEMO_INV_HID); net->inv1s_w = get_weights_aligned(w, DEMO_INV_HID * DEMO_SFEAT); net->invt_w = get_weights_aligned(w, DEMO_INV_HID * DEMO_EMBED); - net->inv2_w = get_weights_aligned(w, DEMO_INV_POOL * DEMO_INV_HID); - net->inv2_b = get_weights_aligned(w, DEMO_INV_POOL); - net->bl_w = get_weights_aligned(w, 64 * DEMO_BL_FEAT); - net->bl_b = get_weights_aligned(w, 64); - net->proj_w = get_weights_aligned(w, net->hidden_size * net->concat_dim); + net->bl_w = get_weights_aligned(w, DEMO_BLH * DEMO_BL_FEAT); + net->bl_b = get_weights_aligned(w, DEMO_BLH); + net->proj_w = get_weights_aligned(w, net->hidden_size * DEMO_CONCAT); net->proj_b = get_weights_aligned(w, net->hidden_size); net->msg_w = get_weights_aligned(w, DEMO_MSG_VOCAB * DEMO_MSG_HID); net->spk_w = get_weights_aligned(w, DEMO_SPKEY * DEMO_SPIN); - net->spk2_w = get_weights_aligned(w, DEMO_SPKEY * DEMO_SPKEY); - net->spk2_b = get_weights_aligned(w, DEMO_SPKEY); - if (net->ide) { - net->ide_role_w = get_weights_aligned(w, 13 * DEMO_IDE_ROLE); - net->ide_race_w = get_weights_aligned(w, 5 * DEMO_IDE_RACE); - net->ide_gend_w = get_weights_aligned(w, 2 * DEMO_IDE_GEND); - net->ide_algn_w = get_weights_aligned(w, 3 * DEMO_IDE_ALGN); - } + net->spm1_w = get_weights_aligned(w, DEMO_SPM * DEMO_SPKEY); + net->spm1_b = get_weights_aligned(w, DEMO_SPM); + net->spm2_w = get_weights_aligned(w, DEMO_SPM * DEMO_SPM); + net->spm2_b = get_weights_aligned(w, DEMO_SPM); + net->ide_role_w = get_weights_aligned(w, 13 * DEMO_IDE_ROLE); + net->ide_race_w = get_weights_aligned(w, 5 * DEMO_IDE_RACE); + net->ide_gend_w = get_weights_aligned(w, 2 * DEMO_IDE_GEND); + net->ide_algn_w = get_weights_aligned(w, 3 * DEMO_IDE_ALGN); + net->mv1_w = get_weights_aligned(w, DEMO_MV * DEMO_INV_HID); + net->mv1_b = get_weights_aligned(w, DEMO_MV); + net->mv2_w = get_weights_aligned(w, DEMO_MV * DEMO_MV); + net->mv2_b = get_weights_aligned(w, DEMO_MV); + net->mr_w = get_weights_aligned(w, DEMO_INV_HID * DEMO_LAB_IN); + net->mr_b = get_weights_aligned(w, DEMO_INV_HID); + net->mm1_w = get_weights_aligned(w, DEMO_MV * DEMO_INV_HID); + net->mm1_b = get_weights_aligned(w, DEMO_MV); + net->mm2_w = get_weights_aligned(w, DEMO_MV * DEMO_MV); + net->mm2_b = get_weights_aligned(w, DEMO_MV); + net->ir_w = get_weights_aligned(w, DEMO_INV_HID * DEMO_LAB_IN); + net->ir_b = get_weights_aligned(w, DEMO_INV_HID); + net->im1_w = get_weights_aligned(w, DEMO_MV * DEMO_INV_HID); + net->im1_b = get_weights_aligned(w, DEMO_MV); + net->im2_w = get_weights_aligned(w, DEMO_MV * DEMO_MV); + net->im2_b = get_weights_aligned(w, DEMO_MV); net->dec_lin = get_weights_aligned(w, DEMO_DEC_PAD * net->hidden_size); net->dec_q = get_weights_aligned(w, DEMO_QDIM * net->hidden_size); net->dec_k = get_weights_aligned(w, DEMO_INV_HID * DEMO_INV_HID); @@ -303,47 +393,104 @@ static int demo_glyph_at(const int16_t* glyphs, int r, int c) { return g; } +// glyph -> 9-class local id and glyph -> 17-class terrain id (mirrors nh_v5_luts_init) +static unsigned char demo_locc_lut[DEMO_VOCAB]; +static unsigned char demo_terrc_lut[DEMO_VOCAB]; +static void demo_v5_luts_init(void) { + static int done = 0; + if (done) return; + done = 1; + for (int g = 0; g < DEMO_VOCAB; g++) { + int ci = (g >= 2359 && g < 2359 + 96) ? g - 2359 : -1; + unsigned char lc = 7; // other + if ((ci >= 1 && ci <= 11) || ci == 17 || ci == 18) lc = 0; // wall/bars/tree + else if (ci >= 19 && ci <= 22) lc = 1; // floor/corridor + else if (ci >= 12 && ci <= 16) lc = 2; // door + else if (ci == 32 || ci == 34 || ci == 41 + || (ci >= 42 && ci <= 64)) lc = 3; // hazard terrain + else if (g >= 0 && g < 1144) lc = 4; // monster + else if ((g >= 1906 && g < 2359) || (g >= 1144 && g < 1525)) lc = 5; // item + else if (ci == 0) lc = 6; // unexplored + if (g == NETHACK_PAD_GLYPH) lc = 8; // off-map + demo_locc_lut[g] = lc; + unsigned char tc = 255; // not counted + if (ci == 23 || ci == 25) tc = 0; // upstairs + else if (ci == 24 || ci == 26) tc = 1; // downstairs + else if (ci == 27) tc = 2; // altar + else if (ci == 31) tc = 3; // fountain + else if (ci == 30) tc = 4; // sink + else if (ci == 29) tc = 5; // throne + else if (ci >= 12 && ci <= 14) tc = 6; // open door + else if (ci == 15 || ci == 16) tc = 7; // closed door + else if (ci >= 42 && ci <= 64) tc = 8; // trap + else if (ci == 28) tc = 9; // grave + else if (ci >= 1 && ci <= 11) tc = 10; // wall + else if (ci == 17 || ci == 18) tc = 11; // bars/tree + else if (ci == 0) tc = 12; // stone + else if (ci >= 19 && ci <= 22) tc = 13; // floor + else if ((g >= 1906 && g < 2359) || (g >= 1144 && g < 1525)) tc = 14; // item + else if (g >= 0 && g < 1144) tc = (g % 381 == 267) ? 16 : 15; // shk / monster + demo_terrc_lut[g] = tc; + } +} + static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // fills decoder->output + demo_v5_luts_init(); const int16_t* glyphs = (const int16_t*)(obs + NETHACK_OFF_GLYPHS); const unsigned char* bl = obs + NETHACK_OFF_BLSTATS; - // local view: per-cell embeds of the egocentric crop, flattened + // local view: 9x9 semantic-class crop (LUT), center forced to class 7 int hx = demo_i32(bl), hy = demo_i32(bl + 4); int half = NETHACK_CROP / 2; for (int p = 0; p < NETHACK_CROP_GRID; p++) { int g = demo_glyph_at(glyphs, hy - half + p / NETHACK_CROP, hx - half + p % NETHACK_CROP); - memcpy(net->x + p * DEMO_EMBED, net->e_eff + g * DEMO_EMBED, - DEMO_EMBED * sizeof(float)); + int cls = p == NETHACK_CROP_GRID / 2 ? 7 : demo_locc_lut[g]; + memcpy(net->x + p * DEMO_LOCC_DIM, net->locc_w + cls * DEMO_LOCC_DIM, + DEMO_LOCC_DIM * sizeof(float)); } - _linear(net->x, net->loc_w, net->loc_b, net->concat, 1, DEMO_LOC_IN, DEMO_LOC_HID); + _linear(net->x, net->loc_w, net->loc_b, net->loc_h1, 1, DEMO_LOC_IN, DEMO_LOC_H1); + _relu(net->loc_h1, net->loc_h1, DEMO_LOC_H1); + _linear(net->loc_h1, net->loc2_w, net->loc2_b, net->concat, 1, DEMO_LOC_H1, DEMO_LOC_HID); _relu(net->concat, net->concat, DEMO_LOC_HID); - // global view: per patch embed+flatten + normalized hero (dx,dy) -> 16 -> - // 128, elementwise max over the 80 tokens (off-map cells of ragged edge - // patches read the pad glyph) - float* glb = net->concat + DEMO_LOC_HID; - for (int o = 0; o < DEMO_GLB_HID; o++) glb[o] = -1e30f; - for (int tk = 0; tk < DEMO_TOK; tk++) { - int r0 = (tk / DEMO_PX) * DEMO_PH, c0 = (tk % DEMO_PX) * DEMO_PW; - for (int pos = 0; pos < DEMO_PCELLS; pos++) { - int g = demo_glyph_at(glyphs, r0 + pos / DEMO_PW, c0 + pos % DEMO_PW); - memcpy(net->px + pos * DEMO_EMBED, net->e_eff + g * DEMO_EMBED, - DEMO_EMBED * sizeof(float)); - } - float dx = (c0 + 0.5f * (DEMO_PW - 1) - hx) / (float)NH_COLS; - float dy = (r0 + 0.5f * (DEMO_PH - 1) - hy) / (float)NH_ROWS; - _linear(net->px, net->g1_w, net->g1_b, net->t16, 1, DEMO_GLB_IN, DEMO_P1); - for (int k = 0; k < DEMO_P1; k++) { - net->t16[k] += dx * net->g1_xy[k * 2] + dy * net->g1_xy[k * 2 + 1]; - if (net->t16[k] < 0.f) net->t16[k] = 0.f; + // terrain branch: landmark table + sector radar (mirrors nh_terr_feat_kernel) + { + float* lm = net->terr_tf; // 12 landmarks x [seen, dx, dy, dist] + float* sec = net->terr_tf + 48; // 8 sectors x 4 bands x 17 classes + for (int i = 0; i < DEMO_TERRF; i++) net->terr_tf[i] = 0.0f; + int lmd[12]; + for (int i = 0; i < 12; i++) lmd[i] = 1 << 30; + int hcell = hy * NH_COLS + hx; + for (int cell = 0; cell < NH_GRID; cell++) { + int g = glyphs[cell]; + if (g < 0) g = 0; + if (g >= DEMO_VOCAB) g = DEMO_VOCAB - 1; + int tc = cell == hcell ? 13 : demo_terrc_lut[g]; + if (tc == 255) continue; + int dy = cell / NH_COLS - hy, dx = cell % NH_COLS - hx; + int ady = dy < 0 ? -dy : dy, adx = dx < 0 ? -dx : dx; + int cheb = adx > ady ? adx : ady; + if (tc < 12 && cheb < lmd[tc]) { + lmd[tc] = cheb; + lm[tc * 4 + 0] = 1.0f; + lm[tc * 4 + 1] = (float)dx * (1.0f / 78.0f); + lm[tc * 4 + 2] = (float)dy * (1.0f / 20.0f); + lm[tc * 4 + 3] = (float)(cheb < 30 ? cheb : 30) * (1.0f / 30.0f); + } + float ang = atan2f((float)dy, (float)dx) + 3.14159265358979f; + int sct = ((int)(ang / 0.78539816339745f)) & 7; + int band = cheb < 3 ? 0 : cheb < 7 ? 1 : cheb < 15 ? 2 : 3; + sec[(sct * 4 + band) * 17 + tc] += 1.0f; } - _linear(net->t16, net->g2_w, net->g2_b, net->t128, 1, DEMO_P1, DEMO_GLB_HID); - for (int o = 0; o < DEMO_GLB_HID; o++) - if (net->t128[o] > glb[o]) glb[o] = net->t128[o]; + float inv_log = 1.0f / logf(1660.0f); + for (int i = 0; i < 8 * 4 * 17; i++) sec[i] = log1pf(sec[i]) * inv_log; + _linear(net->terr_tf, net->terr1_w, net->terr1_b, net->terr_h, 1, DEMO_TERRF, DEMO_TERR_H1); + _relu(net->terr_h, net->terr_h, DEMO_TERR_H1); + float* glb = net->concat + DEMO_LOC_HID; + _linear(net->terr_h, net->terr2_w, net->terr2_b, glb, 1, DEMO_TERR_H1, DEMO_GLB_HID); + _relu(glb, glb, DEMO_GLB_HID); } - _relu(glb, glb, DEMO_GLB_HID); - // inventory entities: per-slot embed -> shared 32->32 linear+relu (kept // as the pointer decoder's keys), then 32 -> 128 with max over slots for // the trunk (matches the CUDA fused pool) @@ -355,12 +502,14 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f if (g < 0) g = 0; if (g >= DEMO_VOCAB) g = DEMO_VOCAB - 1; const signed char* st = invst + slot * NLE_INV_STATE_FIELDS; - float sf[DEMO_SFEAT]; + float* sf = net->sfeat + slot * DEMO_SFEAT; for (int c = 0; c < 4; c++) sf[c] = st[0] == c ? 1.0f : 0.0f; int spe_known = st[1] != -128; sf[4] = (float)spe_known; - sf[5] = spe_known ? (float)st[1] * 0.1f : 0.0f; - sf[6] = log1pf(fmaxf((float)st[2], 0.0f)) * 0.5f; + int spe = st[1] < -7 ? -7 : (st[1] > 7 ? 7 : st[1]); + sf[5] = spe_known ? (float)spe * (1.0f / 7.0f) : 0.0f; + int quan = st[2] < 0 ? 0 : (st[2] > 30 ? 30 : st[2]); + sf[6] = (float)quan * (1.0f / 30.0f); sf[7] = (float)st[3] * (1.0f / 3.0f); sf[8] = (float)st[4] * (1.0f / 3.0f); for (int c = 0; c < 7; c++) sf[9 + c] = (float)((st[5] >> c) & 1); @@ -368,6 +517,9 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f int ot = inv[slot] - NH_GLYPH_OBJ_OFF; // armor slot category one-hot int cat = (ot >= 0 && ot < NH_NUM_OBJECTS) ? nh_obj_armcat[ot] : -1; for (int c = 0; c < 7; c++) sf[17 + c] = cat == c ? 1.0f : 0.0f; + net->otyp[slot] = (ot >= 0 && ot < NH_NUM_OBJECTS) ? ot : -1; + net->wld[slot] = sf[10]; + net->occ[slot] = inv[slot] != NETHACK_PAD_GLYPH; float* h32 = net->slots + slot * DEMO_INV_HID; _linear(net->e_eff + g * DEMO_EMBED, net->inv1_w, net->inv1_b, h32, 1, DEMO_EMBED, DEMO_INV_HID); @@ -382,21 +534,160 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f * net->e_eff[gt * DEMO_EMBED + d]; _relu(h32, h32, DEMO_INV_HID); } - float* invp = net->concat + DEMO_LOC_HID + DEMO_GLB_HID; - for (int o = 0; o < DEMO_INV_POOL; o++) { - float best = -1e30f; + { // inventory deep pool + pass-throughs (mirrors mv MLP + nh_min_summax + nh_pass) + float* dst = net->concat + DEMO_MINV_OFF; + float mx[DEMO_MV]; + int any = 0; + for (int d = 0; d < DEMO_MV; d++) { dst[d] = 0.0f; mx[d] = -1e30f; } for (int slot = 0; slot < NETHACK_INV_SLOTS; slot++) { - float v = 0.0f; - for (int k = 0; k < DEMO_INV_HID; k++) - v += net->inv2_w[o * DEMO_INV_HID + k] * net->slots[slot * DEMO_INV_HID + k]; - if (v > best) best = v; + if (!net->occ[slot]) continue; + any = 1; + const float* r = net->slots + slot * DEMO_INV_HID; + float h1[DEMO_MV]; + for (int d = 0; d < DEMO_MV; d++) { + float acc = net->mv1_b[d]; + for (int c = 0; c < DEMO_INV_HID; c++) + acc += net->mv1_w[d * DEMO_INV_HID + c] * r[c]; + h1[d] = acc > 0.f ? acc : 0.f; + } + for (int d = 0; d < DEMO_MV; d++) { + float acc = net->mv2_b[d]; + for (int c = 0; c < DEMO_MV; c++) + acc += net->mv2_w[d * DEMO_MV + c] * h1[c]; + float v = acc > 0.f ? acc : 0.f; + dst[d] += 0.2f * v; + if (v > mx[d]) mx[d] = v; + } + } + for (int d = 0; d < DEMO_MV; d++) dst[DEMO_MV + d] = any ? mx[d] : 0.f; + float* ps = net->concat + DEMO_PASS_OFF; // wld [r|sf] | quiver [r|sf] | worn mean + for (int d = 0; d < 40 + 40 + DEMO_INV_HID + DEMO_PACC; d++) ps[d] = 0.0f; + int nworn = 0; + for (int slot = 0; slot < NETHACK_INV_SLOTS; slot++) { + if (!net->occ[slot]) continue; + const float* r = net->slots + slot * DEMO_INV_HID; + const float* sf = net->sfeat + slot * DEMO_SFEAT; + if (sf[10] > 0.5f) + for (int e = 0; e < 40; e++) + ps[e] += e < DEMO_INV_HID ? r[e] : sf[e - DEMO_INV_HID]; + if (sf[12] > 0.5f) + for (int e = 0; e < 40; e++) + ps[40 + e] += e < DEMO_INV_HID ? r[e] : sf[e - DEMO_INV_HID]; + if (sf[9] > 0.5f) { + nworn++; + for (int e = 0; e < DEMO_INV_HID; e++) ps[80 + e] += r[e]; + } + } + if (nworn > 0) + for (int e = 0; e < DEMO_INV_HID; e++) ps[80 + e] /= (float)nworn; + { // accessory panel (mirrors nh_acc_slot): first worn owner of each slot wins + int nring = 0, have[4] = {0, 0, 0, 0}; + for (int slot = 0; slot < NETHACK_INV_SLOTS; slot++) { + if (!net->occ[slot]) continue; + const float* sf = net->sfeat + slot * DEMO_SFEAT; + if (sf[9] <= 0.5f) continue; + int ot = net->otyp[slot], a = -1; + if (ot >= 178 && ot <= 188) a = 0; + else if (ot >= 207 && ot <= 209) a = 3; + else if (ot >= 150 && ot <= 177) { if (nring < 2) a = 1 + nring; nring++; } + if (a < 0 || have[a]) continue; + have[a] = 1; + const float* r = net->slots + slot * DEMO_INV_HID; + float* dst = ps + 80 + DEMO_INV_HID + a * (DEMO_INV_HID + DEMO_SFEAT); + for (int e = 0; e < DEMO_INV_HID + DEMO_SFEAT; e++) + dst[e] = e < DEMO_INV_HID ? r[e] : sf[e - DEMO_INV_HID]; + } + } + } + { // typed token streams (mirrors nh_lab_tok_kernel + mr/mm MLPs + nh_min_summax + sgate) + for (int is_mon = 1; is_mon >= 0; is_mon--) { + int list_off = is_mon ? NETHACK_OFF_TOKM : NETHACK_OFF_TOKI; + const float *rw = is_mon ? net->mr_w : net->ir_w; + const float *rb = is_mon ? net->mr_b : net->ir_b; + const float *w1 = is_mon ? net->mm1_w : net->im1_w; + const float *b1 = is_mon ? net->mm1_b : net->im1_b; + const float *w2 = is_mon ? net->mm2_w : net->im2_w; + const float *b2 = is_mon ? net->mm2_b : net->im2_b; + float* dst = net->concat + (is_mon ? DEMO_MLM_OFF : DEMO_MLI_OFF); + float tok[DEMO_LABK][DEMO_LAB_IN], rep[DEMO_LABK][DEMO_INV_HID]; + int valid[DEMO_LABK]; + for (int k = 0; k < DEMO_LABK; k++) { + const unsigned char* e = obs + list_off + k * NETHACK_V3_MONF; + int row = e[0] | (e[1] << 8); + int dx = (signed char)e[2], dy = (signed char)e[3]; + int f4 = e[4], f5 = e[5], f6 = e[6]; + int g = row <= 0 ? -1 : is_mon ? row - 1 + : (row < 454 ? 1906 + row - 1 : 1144 + row - 454); + valid[k] = g >= 0; + float* o = tok[k]; + for (int d = 0; d < DEMO_LAB_IN; d++) o[d] = 0.f; + if (g >= 0) { + memcpy(o, net->e_eff + (size_t)g * DEMO_EMBED, DEMO_EMBED * sizeof(float)); + int adx = dx < 0 ? -dx : dx, ady = dy < 0 ? -dy : dy; + int cheb = adx > ady ? adx : ady; + float fdx = dx * (1.0f/40.0f), fdy = dy * (1.0f/11.0f); + o[32] = fdx < -1.f ? -1.f : (fdx > 1.f ? 1.f : fdx); + o[33] = fdy < -1.f ? -1.f : (fdy > 1.f ? 1.f : fdy); + o[34] = (cheb < 15 ? cheb : 15) * (1.0f/15.0f); + o[35] = k * (1.0f/15.0f); + if (is_mon) { + o[36] = (f4 & 1) ? 1.f : 0.f; + o[37] = (f4 & 8) ? 1.f : 0.f; + o[38] = (f4 & 4) ? 1.f : 0.f; + o[39] = cheb <= 1 ? 1.f : 0.f; + float df = f5 * 0.04f; o[40] = df > 1.f ? 1.f : df; + float sp = f6 * (1.0f/24.0f); o[41] = sp > 1.f ? 1.f : sp; + int hb = NH_MON_HAZ[(row - 1) % NH_MONS_STATIC_N]; + for (int hz = 0; hz < 4; hz++) o[42 + hz] = (float)((hb >> hz) & 1); + } else { + o[36] = (f5 & 1) ? 1.f : 0.f; + o[37] = (f5 & 2) ? 1.f : 0.f; + } + } + for (int j = 0; j < DEMO_INV_HID; j++) { // rep16 (decoder-free stream key) + float acc = rb[j]; + for (int c = 0; c < DEMO_LAB_IN; c++) acc += rw[j * DEMO_LAB_IN + c] * o[c]; + rep[k][j] = acc > 0.f ? acc : 0.f; + } + } + float mx[DEMO_MV]; + int any = 0; + for (int d = 0; d < DEMO_MV; d++) { dst[d] = 0.0f; mx[d] = -1e30f; } + for (int k = 0; k < DEMO_LABK; k++) { + if (!valid[k]) continue; + any = 1; + float h1[DEMO_MV]; + for (int j = 0; j < DEMO_MV; j++) { + float acc = b1[j]; + for (int c = 0; c < DEMO_INV_HID; c++) + acc += w1[j * DEMO_INV_HID + c] * rep[k][c]; + h1[j] = acc > 0.f ? acc : 0.f; + } + for (int j = 0; j < DEMO_MV; j++) { + float acc = b2[j]; + for (int c = 0; c < DEMO_MV; c++) acc += w2[j * DEMO_MV + c] * h1[c]; + float v = acc > 0.f ? acc : 0.f; + dst[j] += 0.25f * v; + if (v > mx[j]) mx[j] = v; + } + } + for (int d = 0; d < DEMO_MV; d++) dst[DEMO_MV + d] = any ? mx[d] : 0.f; + float* g0 = dst + 2 * DEMO_MV; // token-0 gate rep + for (int d = 0; d < DEMO_INV_HID; d++) g0[d] = valid[0] ? rep[0][d] : 0.f; + if (!is_mon) { // underfoot sum over rep16 of underfoot items (tok[36]) + float* uf = g0 + DEMO_INV_HID; + for (int d = 0; d < DEMO_INV_HID; d++) uf[d] = 0.f; + for (int k = 0; k < DEMO_LABK; k++) { + if (!valid[k] || tok[k][36] <= 0.5f) continue; + for (int d = 0; d < DEMO_INV_HID; d++) uf[d] += rep[k][d]; + } + } } - invp[o] = fmaxf(best + net->inv2_b[o], 0.0f); } // blstats+extra features (25 scalars, hunger 7, cond bits 13, prev verb // one-hot, inv class counts, hp/ene frac, dnum one-hot, engraving bits) - float* f = net->concat + DEMO_LOC_HID + DEMO_GLB_HID + DEMO_INV_POOL + 64; + float* f = net->concat + DEMO_LOC_HID + DEMO_GLB_HID + DEMO_BLH; int j = 0; for (int i = 0; i < 27; i++) { if (i == 21 || i == 25) continue; // hunger, condition: expanded below @@ -429,13 +720,14 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f float d = (float)demo_i32(ex + 4*(NETHACK_EXTRA_WEIGHT+0)) * 0.01f - 1.0f; f[j++] = d / (1.0f + fabsf(d)); f[j++] = (float)demo_i32(ex + 4*(NETHACK_EXTRA_WEIGHT+1)) * 0.001f; } - for (int k = 0; k < 20; k++) // role/race/gender one-hots; zeroed in the ide era - f[j++] = net->ide ? 0.f : (float)demo_i32(ex + 4*(NETHACK_EXTRA_ROLEOH + k)); + for (int k = 0; k < 20; k++) // role/race/gender one-hots; dead under the embed channel + f[j++] = 0.f; + for (int k = 0; k < 8; k++) f[j++] = (float)((demo_i32(ex + 4*NETHACK_EXTRA_INTRINS) >> k) & 1); for (int k = 0; k < DEMO_BL_FEAT; k++) f[k] = fminf(fmaxf(f[k], -1.f), 1.f); - float* blout = net->concat + DEMO_LOC_HID + DEMO_GLB_HID + DEMO_INV_POOL; - _linear(f, net->bl_w, net->bl_b, blout, 1, DEMO_BL_FEAT, 64); - _relu(blout, blout, 64); + float* blout = net->concat + DEMO_LOC_HID + DEMO_GLB_HID; + _linear(f, net->bl_w, net->bl_b, blout, 1, DEMO_BL_FEAT, DEMO_BLH); + _relu(blout, blout, DEMO_BLH); demo_msg_pool(net, obs, net->concat + DEMO_MSG_CONCAT_OFF); @@ -464,19 +756,43 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f } } float* sp = net->concat + DEMO_SPELL_CONCAT_OFF; - for (int d = 0; d < DEMO_SPKEY; d++) { - float best = -1e30f; - for (int s = 0; s < NETHACK_SPELL_SLOTS; s++) { - float acc = 0.f; - for (int k = 0; k < DEMO_SPKEY; k++) - acc += net->spk2_w[d * DEMO_SPKEY + k] * net->spkeys[s * DEMO_SPKEY + k]; - if (acc > best) best = acc; + for (int d = 0; d < 2*DEMO_SPM; d++) sp[d] = 0.0f; + int mf = 999, ml = 0, nn = 0; long mr = 99999; + for (int s = 0; s < NETHACK_SPELL_SLOTS; s++) { + const unsigned char* q = obs + NETHACK_OFF_EXTRA + + 4*(NETHACK_EXTRA_SPELL + 1 + 4*s); + int sid = demo_i32(q); + if (sid <= 0) continue; + nn++; + int lv = demo_i32(q + 4), fl = demo_i32(q + 8), kn = demo_i32(q + 12); + if (fl < mf) mf = fl; + if (lv > ml) ml = lv; + if (kn < mr) mr = kn; + // entity value: per-slot MLP over the (already relu'd) key + float m1[DEMO_SPM], v; + const float* k = net->spkeys + s * DEMO_SPKEY; + for (int d = 0; d < DEMO_SPM; d++) { + float acc = net->spm1_b[d]; + for (int c = 0; c < DEMO_SPKEY; c++) acc += net->spm1_w[d * DEMO_SPKEY + c] * k[c]; + m1[d] = acc > 0.f ? acc : 0.f; + } + for (int d = 0; d < DEMO_SPM; d++) { + float acc = net->spm2_b[d]; + for (int c = 0; c < DEMO_SPM; c++) acc += net->spm2_w[d * DEMO_SPM + c] * m1[c]; + v = acc > 0.f ? acc : 0.f; + sp[d] += 0.25f * v; + if (v > sp[DEMO_SPM + d]) sp[DEMO_SPM + d] = v; } - float v = best + net->spk2_b[d]; - sp[d] = v > 0.f ? v : 0.f; } + if (nn == 0) { mf = 100; ml = 0; mr = 20000; } // empty book -> [1,0,0,1] + float e0 = (float)mf * 0.01f, e1 = (float)ml * (1.0f/7.0f); + float e2 = (float)(nn < 8 ? nn : 8) * 0.125f, e3 = (float)mr * 0.00005f; + sp[2*DEMO_SPM + 0] = e0 < 0.f ? 0.f : (e0 > 1.f ? 1.f : e0); + sp[2*DEMO_SPM + 1] = e1 < 0.f ? 0.f : (e1 > 1.f ? 1.f : e1); + sp[2*DEMO_SPM + 2] = e2 < 0.f ? 0.f : (e2 > 1.f ? 1.f : e2); + sp[2*DEMO_SPM + 3] = e3 < 0.f ? 0.f : (e3 > 1.f ? 1.f : e3); } - if (net->ide) { // identity table tail; mirrors nh_idemb_kernel (last set bit wins) + { // identity table tail; mirrors nh_idemb_kernel (last set bit wins) float* ide = net->concat + DEMO_IDE_CONCAT_OFF; int r = 0, rc = 0, g = 0; for (int k = 0; k < 13; k++) if (demo_i32(ex + 4*(NETHACK_EXTRA_ROLEOH + k))) r = k; @@ -489,7 +805,7 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f for (int d = 0; d < DEMO_IDE_GEND; d++) *ide++ = net->ide_gend_w[g * DEMO_IDE_GEND + d]; for (int d = 0; d < DEMO_IDE_ALGN; d++) *ide++ = net->ide_algn_w[al * DEMO_IDE_ALGN + d]; } - _linear(net->concat, net->proj_w, net->proj_b, net->hidden, 1, net->concat_dim, net->hidden_size); + _linear(net->concat, net->proj_w, net->proj_b, net->hidden, 1, DEMO_CONCAT, net->hidden_size); _relu(net->hidden, net->hidden, net->hidden_size); mingru(net->mingru, net->hidden); @@ -511,6 +827,9 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f for (int k = 0; k < H; k++) acc += net->dec_q[r * H + k] * hs[k]; q[r] = acc; } + // NH_DEC_COS=1: legacy full-cosine slot logits (pre-2026-08-25 checkpoints) + static int dec_cos = -1; + if (dec_cos < 0) { const char* e = getenv("NH_DEC_COS"); dec_cos = e && e[0] && e[0] != '0'; } float kmat[DEMO_INV_FLAT], kn[NETHACK_INV_SLOTS]; for (int i = 0; i < NETHACK_INV_SLOTS; i++) { float nk = 0.0f; @@ -521,7 +840,7 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // f kmat[i * DEMO_INV_HID + r] = acc; nk += acc * acc; } - kn[i] = sqrtf(nk) + 1e-6f; + kn[i] = dec_cos ? sqrtf(nk) + 1e-6f : 1.0f; } for (int a = 0; a < NETHACK_NUM_ACTIONS; a++) net->logits[a] = tmp[a]; for (int h = 0; h < DEMO_PTR_HEADS; h++) { @@ -571,6 +890,7 @@ static void demo_restore_term(void) { } static void demo_on_signal(int sig) { + (void)sig; demo_restore_term(); _exit(128 + sig); } @@ -685,16 +1005,16 @@ static void demo_step_once(NethackNet* net, Nethack* env, float* acts_f, nethack_net_forward(net, env->agents[0].observations); for (int i = 0; i < DEMO_OD; i++) if (!env->action_mask[i]) net->logits[i] = -1e9f; - multidiscrete(net->md, net->logits, acts_f, 0, NULL); + multidiscrete(net->md, net->logits, acts_f, 0, NULL); // illegal logits already set to -1e9 above for (int h = 0; h < DEMO_NUM_HEADS; h++) env->agents[0].actions[h] = acts_f[h]; puf_step(env); if (env->agents[0].terminals[0] > 0.5f) { float d = env->log.max_depth - *ep_depth; float x = env->log.max_xp_level - *ep_xp; float g = env->log.game_time - *ep_gt; - fprintf(stderr, "episode end: score=%.0f len=%.0f max_depth=%.0f xp=%.0f game_t=%.0f " + fprintf(stderr, "episode end: role=%d score=%.0f len=%.0f max_depth=%.0f xp=%.0f game_t=%.0f " "eats=%.1f floor_eats=%.1f wears=%.1f throws=%.1f\n", - env->log.score - *ep_score, env->log.episode_length - *ep_len, d, x, g, + env->role_idx, env->log.score - *ep_score, env->log.episode_length - *ep_len, d, x, g, env->log.verb_uses[NETHACK_ACT_EAT], env->log.floor_eats, env->log.verb_uses[NETHACK_ACT_WEAR], env->log.verb_uses[NETHACK_ACT_THROW]); @@ -1118,7 +1438,11 @@ static void run_demo_auto(long max_steps, int frame_ms) { } if (frame_ms > 0) { demo_render(&env, 1000 / frame_ms, t); +#ifdef __EMSCRIPTEN__ + emscripten_sleep(frame_ms); // usleep busy-waits in wasm: yield or the page never paints +#else usleep(frame_ms * 1000); +#endif } } if (env.log.n > 0) diff --git a/ocean/nethack/nethack.cu b/ocean/nethack/nethack.cu index f9604fe25c..330f79f6a0 100644 --- a/ocean/nethack/nethack.cu +++ b/ocean/nethack/nethack.cu @@ -12,6 +12,16 @@ __global__ void nh_bias_relu_kernel( data[idx] = from_float(fmaxf(0.0f, to_float(data[idx]) + to_float(bias[idx % dim]))); } +__global__ void nh_bias_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(to_float(data[idx]) + to_float(bias[idx % dim])); +} +// GEMM + bias + relu (16 launches per forward) +#define NH_MM_BR(A, W, O, BIAS) do { puf_mm((A), (W), (O), stream); \ + int64_t nn_ = numel((O)->shape); int dd_ = (int)(O)->shape[ndim((O)->shape) - 1]; \ + nh_bias_relu_kernel<<>>((O)->data, (BIAS).data, (int)nn_, dd_); } while (0) // constants // Obs layout (must match ocean/nethack/nethack.h): // [0, 2*NH_MGRID) full 79x21 glyph grid, int16 LE (map memory included) @@ -24,14 +34,32 @@ static constexpr int NH_MGRID = NH_MAPW * NH_MAPH; static constexpr int NH_GLYPH_VOCAB = 5977; // MAX_GLYPH + 1 (NetHack 3.6.6) static constexpr int NH_PAD_GLYPH = NH_GLYPH_VOCAB - 1; // NO_GLYPH: off-map crop cells static constexpr int NH_EMBED_DIM = 32; +static constexpr int NH_TERRC_DIM = 8; // v3 crop terrain code width (own tiny table) +static constexpr int NH_RING = 9; // v3 3x3 adjacency ring, full-width monster rows static constexpr int NH_CROP = 9, NH_CHALF = 4; // NETHACK_CROP, egocentric static constexpr int NH_CGRID = NH_CROP * NH_CROP; static constexpr int NH_PW = 5, NH_PH = 5; // patch size (cells) static constexpr int NH_PX = 16, NH_PY = 5; // patch grid (ceil 79/5, 21/5) static constexpr int NH_TOK = NH_PX * NH_PY; // 80 global tokens static constexpr int NH_PCELLS = NH_PW * NH_PH; // cells per patch (off-map -> pad glyph) -static constexpr int NH_LOC_IN = NH_CGRID * NH_EMBED_DIM; +// v5 champion-candidate (ENCODER_V5.md): NH_LAB base + five gated changes, +// each independently settable for the per-change screen. +// NH_LOC3: 9-class semantic crop (croplab: danger .998, esc-dir .83; +// classes ~= depth, both help; width flat 128-512 -> 256) +static constexpr int NH_LOCC_DIM = 8; +static constexpr int NH_LOCC_CLASSES = 9; // wall floor door hazard mon item stone other offmap +// v3 local: terrain field (81x8, own table) + adjacency ring (9x32, T_mon) +// + underfoot item (32, T_item). Positions share "what", loc_w keeps "where". +static constexpr int NH_LOC_IN = NH_CGRID * NH_LOCC_DIM; static constexpr int NH_LOC_HID = 256; +// two-layer local branch (crop -> NH_LOC_H1 -> NH_LOC_HID) so bearing +// at 2-4 tiles has a nonlinearity to be represented in (audit: 0.63 decodable +// from the single linear+relu, 0.39 in the direction logits). NH_LOC3 forces +// the two-layer path at hidden 256. +static constexpr int NH_LOC_H1 = 256; // hidden width; loc_w is (NH_LOC_H1, NH_LOC_IN) +// NH_TERR: landmark table (12x4) + sector radar (8x4x17) -> MLP (terrain lab) +static constexpr int NH_TERRF = 12 * 4 + 8 * 4 * 17; // 592 features +static constexpr int NH_TERR_H1 = 256; // Global branch: per patch, embed+flatten (25 cells x 32 dims) + normalized // (dx,dy) patch-center offset from the hero -> 16 -> 128, then elementwise // MAX over the 80 tokens. The 16-dim bottleneck keeps the fused per-glyph @@ -44,6 +72,14 @@ static constexpr int NH_GLB_HID = 128; static constexpr int NH_TROW = NH_PCELLS * NH_P1; // fused-table row: per-pos 16-dim static constexpr int NH_PAD_PER_SAMPLE = NH_TOK * NH_PCELLS - NH_MGRID; static constexpr int NH_HOT_G = 10; // hot-glyph dT smem slots (10x400 int64 = 32KB) +// typed token streams (supervised-lab locked design, 2026-08). Typed K-nearest token +// streams (monsters, items; lean NH_TOK_OBS lists env-side) with deep values +// (2-layer MLP), 8-head attention pools scoring from raw token features, and +// a rank feature; inventory gets the 8-query attention pool (NH_INV_ATTN); +// aux mega-spec heads regress ~26 engine-truth targets off the concat +// (training-time only shaping; targets ride the NH_AUX_OBS extras block and +// are never featurized). All additive beside the champion branches. +// component toggles for leave-one-out ablation static constexpr int NH_BL_RAW = 27; // NLE_BLSTATS_SIZE static constexpr int NH_BL_HUNGER = 21, NH_BL_CONDITION = 25; static constexpr int NH_BL_HP = 10, NH_BL_ENE = 14; // hp/hpmax at 10/11, ene/enemax at 14/15 @@ -52,7 +88,9 @@ static constexpr int NH_OCLASSES = 18; // MAXOCLASSES static constexpr int NH_EXTRA_SHOP = 2 + NH_OCLASSES; // extra[] index of the shop pair static constexpr int NH_SPELL_SLOTS = 8; // NETHACK_SPELL_SLOTS static constexpr int NH_EX_ROLEOH = 2 + NH_OCLASSES + 2 + 1 + 4 * NH_SPELL_SLOTS + 2; -static constexpr int NH_EX_RAW = NH_EX_ROLEOH + 13 + 5 + 2; // NETHACK_EXTRA_INTS (+role/race/gender one-hots) +static constexpr int NH_EX_RAW = NH_EX_ROLEOH + 13 + 5 + 2 + + 1 +; // NETHACK_EXTRA_INTS (+role/race/gender one-hots) // blstats feature map (cumulative offsets; each block documented at its // kernel branch). hp/ene fracs are the danger ratios the linear bl_w can't // synthesize from separate cur/max scalars; dnum is one-hot because dungeon @@ -71,7 +109,13 @@ static constexpr int NH_F_SPELL = NH_F_SHOP + 2; // encumbrance ratio (softsign around the wall) + carry capacity /1000 static constexpr int NH_F_WEIGHT = NH_F_SPELL + 1; static constexpr int NH_F_ROLE = NH_F_WEIGHT + 2; // 13 role + 5 race + 2 gender -static constexpr int NH_BL_FEAT = NH_F_ROLE + 20; +static constexpr int NH_F_INTRINS = NH_F_ROLE + 20; // 8 intrinsic bits +// threat block: visible bit, proximity, bearing 8-onehot, adj/near3/vis +// counts, nearest difficulty, difficulty-vs-xplvl, nearest speed +static constexpr int NH_F_THREAT = NH_F_INTRINS + 8; +static constexpr int NH_BL_FEAT = NH_F_THREAT; +static constexpr int NH_EX_INTRINS = NH_EX_ROLEOH + 20; +static constexpr int NH_EX_THREAT = NH_EX_INTRINS + 1; static constexpr int NH_BL_DNUM = 23; static constexpr int NH_BL_HID = 64; // Inventory entity branch: 55 slot glyphs, each embed -> shared 32->32 @@ -95,26 +139,83 @@ static constexpr int NH_INV_POOL = 128; // scaled by 1/sqrt(count+1) (normalized bag / EmbeddingBag sum). The summary // is concatenated raw (signed, no relu) like the blstats raw features. static constexpr int NH_MSG_LEN = 128; // raw topline chars in obs tail -static constexpr int NH_MSG_VOCAB = 4096; // trigram hash buckets -static constexpr int NH_MSG_LOG2V = 12; // log2(NH_MSG_VOCAB) -static constexpr int NH_MSG_HID = 32; // trigram embed = message summary dim -static constexpr int NH_MSG_CONCAT_OFF = NH_LOC_HID + NH_GLB_HID + NH_INV_POOL + NH_BL_HID + NH_BL_FEAT; +#ifdef NH_MSG_V4096 +static constexpr int NH_MSG_VOCAB = 4096; // pretrained-code arms (4096x128 tables) +static constexpr int NH_MSG_LOG2V = 12; +#else +static constexpr int NH_MSG_VOCAB = 1024; // champion: ~4.6K distinct trigrams, 1024 buckets tie 4096 (probe + RL n=8) +static constexpr int NH_MSG_LOG2V = 10; // log2(NH_MSG_VOCAB) +#endif +#if defined(NH_MSG128) +static constexpr int NH_MSG_HID = 128; // pretrained-code arms +#else +static constexpr int NH_MSG_HID = 256; // champion: width ladder 32/128/256 pooled n=13, 256 ahead on mean/median/tail +#endif +// NH_INV2 drops the max-pool trunk summary (half-dead in production); +// its slice leaves the concat entirely. +static constexpr int NH_INVP_DIM = 0; +static constexpr int NH_MSG_CONCAT_OFF = NH_LOC_HID + NH_GLB_HID + NH_INVP_DIM + NH_BL_HID + NH_BL_FEAT; // spell-key path (v3 pointer): per slot, key = spk_w . [e_eff(book glyph) | // known, lev/7, fail/100, know/20000]; keys feed the CAST pointer head and a // sum-pooled 16-dim trunk summary. Empty slots are exact zeros end to end. static constexpr int NH_SPKEY = NH_INV_HID; // 16, shared key width static constexpr int NH_SPIN = NH_EMBED_DIM + 4; // 36 key inputs/slot +// spell entity block (v6): keys r = relu(spk_w x) -> per-slot MLP 16->64->64 +// -> masked sum x0.25 | masked max + 4 exact doorstep scalars [min_fail, +// max_lev, n/8, min_retention]. Same block as inventory/streams; trunk channel +// backprops through the keys (dense gradient; the old spk2 raw-key max-pool +// destroyed identity, AUC .18 vs .97 — this max is over post-depth values). +static constexpr int NH_SPM = 64; +static constexpr int NH_SPELL_SLICE = 2 * NH_SPM + 4; static constexpr int NH_SPELL_CONCAT_OFF = NH_MSG_CONCAT_OFF + NH_MSG_HID; -// identity embeddings A/B arm: 1 = explicit tables in a direct concat channel -// (the one-hot bl features go dead); 0 = the committed one-hot representation -#ifndef NH_ID_EMBED -#define NH_ID_EMBED 1 -#endif +// identity embeddings: explicit role/race/gend/align tables in a direct concat channel static constexpr int NH_IDE_ROLE = 16, NH_IDE_RACE = 8; static constexpr int NH_IDE_GEND = 8, NH_IDE_ALGN = 8; static constexpr int NH_IDE_DIM = NH_IDE_ROLE + NH_IDE_RACE + NH_IDE_GEND + NH_IDE_ALGN; -static constexpr int NH_IDE_CONCAT_OFF = NH_SPELL_CONCAT_OFF + NH_SPKEY; -static constexpr int NH_CONCAT = NH_IDE_CONCAT_OFF + (NH_ID_EMBED ? NH_IDE_DIM : 0); +static constexpr int NH_IDE_CONCAT_OFF = NH_SPELL_CONCAT_OFF + NH_SPELL_SLICE; +static constexpr int NH_GMEAN_CONCAT_OFF = NH_IDE_CONCAT_OFF + NH_IDE_DIM; +static constexpr int NH_TERR_VOCAB = 128; // cmap_index+1; 0 unseen; 127 off-map pad +static constexpr int NH_ENT_K = 32; +static constexpr int NH_ENT_F = 6; +static constexpr int NH_ENT_IN = 40; +static constexpr int NH_ENT_HID = 32; +static constexpr int NH_ENT_TAIL = 36; +// v3 typed-level: glyph ranges, tables, two K-nearest token lists +static constexpr int NH_NUMMONS = 381, NH_PET_OFF = 381, NH_DET_OFF = 762; +static constexpr int NH_BODY_OFF = 1144, NH_OBJ_LO = 1906; +static constexpr int NH_MON_ROWS = 384; // species+1 (0 = pad/empty) +static constexpr int NH_ITEM_ROWS = 840; // objects 1..453, bodies 454..834, 0 = pad +#ifndef NETHACK_V3_K +#define NETHACK_V3_K 16 +#endif +static constexpr int NH_V3K = NETHACK_V3_K; +static constexpr int NH_V3_MONF = 8, NH_V3_ITEMF = 8; +static constexpr int NH_V3_IN = 40, NH_V3_HID = 32, NH_V3_TAIL = 36; +static constexpr int NH_ITBL = NH_GLYPH_VOCAB; // inv/discovery table rows +__global__ void nh_scale_kernel(precision_t* p, float a, int64_t n) { + int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) p[i] = from_float(to_float(p[i]) * a); +} +static constexpr int NH_AP_DIM = NH_P1 + 2; +static constexpr int NH_TK_K = 8; +static constexpr int NH_TK_DIM = NH_P1 + 2; // gated rep + (dx,dy) +static constexpr int NH_ENT_CONCAT_OFF = NH_GMEAN_CONCAT_OFF; +static constexpr int NH_IVA_M = 8; // invattn query heads (invattn8) +static constexpr int NH_IVA_CONCAT_OFF = NH_ENT_CONCAT_OFF; +// lab stream dims: 16 tokens x 48 inputs (emb 32 + dx dy cheb rank + flags + +// diff speed), deep values 48 -> 64 -> 64, 8 heads x 8 dims +static constexpr int NH_LABK = NETHACK_V3_K; +static constexpr int NH_LAB_IN = 48; +static constexpr int NH_LAB_HID = 64; +static constexpr int NH_LAB_HEADS = 8; +static constexpr int NH_AUXH = 32; // aux head rows (26 used, %8 pad) +static constexpr float NH_AUX_COEF = 0.05f; +static constexpr int NH_LABM_CONCAT_OFF = NH_IVA_CONCAT_OFF + NH_IVA_M * NH_INV_HID; +static constexpr int NH_LABI_CONCAT_OFF = NH_LABM_CONCAT_OFF + NH_LAB_HID; +// NH_INV2 tail: hard wield readout (16, parameterless) + sum channel (64) +static constexpr int NH_WLD_CONCAT_OFF = NH_LABI_CONCAT_OFF + NH_LAB_HID; +static constexpr int NH_ISUM_DIM = 64; +static constexpr int NH_ISUM_CONCAT_OFF = NH_WLD_CONCAT_OFF + NH_INV_HID; static constexpr int NH_BL_OFF = 2 * NH_MGRID; // blstats offset, obs elements static constexpr int NH_INV_OFF = NH_BL_OFF + (NH_BL_RAW + NH_EX_RAW) * 4; // obs v4: per-slot identification-gated state, 8 int8 fields per slot @@ -125,8 +226,44 @@ static constexpr int NH_ST_RAW = 8; // NLE_INV_STATE_FIELDS static constexpr int NH_SFEAT = 24; // buc4 + known+spe + quan + ero2 + flags7 + tk + armcat7 // discovered-type glyphs: true otyp glyph once dknown && oc_name_known, else pad static constexpr int NH_INVTRUE_OFF = NH_INVST_OFF + NH_INV * NH_ST_RAW; +static constexpr int NH_PWORN_DIM = NH_INV_HID; // 16-d worn content mean +#ifndef NH_NO_APANEL +// accessory panel: 4 single-occupant slots [amulet | ring A | ring B | eyewear], +// each an exact [r16|sfeat24] copy of the worn slot (zeros when bare); rings in +// inventory order. otyps (onames.h, NetHack 3.6.6): rings 150-177, amulets +// 178-188, lenses/blindfold/towel 207-209. Worn bit = owornmask & W_ACCESSORY. +static constexpr int NH_PACC_DIM = 4 * (NH_INV_HID + NH_SFEAT); // 160 +#else +static constexpr int NH_PACC_DIM = 0; +#endif +static constexpr int NH_PASS_DIM = 40 + 40 + NH_PWORN_DIM + NH_PACC_DIM; +// V6-min: statistics + lookups everywhere. Inventory: r16 -> MLP 64->64 -> +// sum|max + pass-throughs (wielded/quivered [r|sfeat], worn-mean). Streams: +// tok48 -> rep16 -> MLP 64->64 -> sum|max + token-0 gate (+underfoot for items). +// No encoder attention anywhere; decoder pointers unchanged. +static constexpr int NH_MV = 64; // value width (all arms) +static constexpr int NH_MINV_DIM = 2 * NH_MV + 40 + 40 + NH_PWORN_DIM + NH_PACC_DIM; +static constexpr int NH_MINV_OFF = NH_IVA_CONCAT_OFF; // 690 +static constexpr int NH_MLM_DIM = 2 * NH_MV + NH_INV_HID; // 144 +static constexpr int NH_MLM_OFF = NH_MINV_OFF + NH_MINV_DIM; +static constexpr int NH_MLI_DIM = 2 * NH_MV + 2 * NH_INV_HID; // 160 +static constexpr int NH_MLI_OFF = NH_MLM_OFF + NH_MLM_DIM; +static constexpr int NH_CONCAT = NH_MLI_OFF + NH_MLI_DIM; // 1410 + static constexpr int NH_MSG_OFF = NH_INVTRUE_OFF + NH_INV * 2; // message block start -static constexpr int NH_OBS_SIZE = NH_MSG_OFF + NH_MSG_LEN; +static constexpr int NH_OBS_SIZE = NH_MSG_OFF + NH_MSG_LEN + + + NH_V3K * (NH_V3_MONF + NH_V3_ITEMF); +static constexpr int NH_TERR_OFF = NH_MSG_OFF + NH_MSG_LEN; +static constexpr int NH_ENTL_OFF = NH_TERR_OFF + NH_MGRID; +static constexpr int NH_OBJM_OFF = NH_TERR_OFF + NH_MGRID; // v3 (aliases ENTL; exclusive) +static constexpr int NH_VMON_OFF = NH_OBJM_OFF + NH_MGRID * 2; +static constexpr int NH_VITEM_OFF = NH_VMON_OFF + NH_V3K * NH_V3_MONF; +// lean token lists (NH_LAB / env NH_TOK_OBS): appended after every other +// optional block, no dense planes (netlib.h NETHACK_OFF_TOKM mirror) +static constexpr int NH_TOKM_OFF = NH_MSG_OFF + NH_MSG_LEN +; +static constexpr int NH_TOKI_OFF = NH_TOKM_OFF + NH_V3K * NH_V3_MONF; static constexpr int NH_SORT_BLOCKS = 256; // hist grid (smem histograms) static constexpr int NH_HOT_T = 16; // hot-glyph smem rows (16x32 int64 = 4KB) @@ -181,6 +318,11 @@ __device__ __forceinline__ void nh_fxp_atomic_add(long long* addr, float v) { __device__ __forceinline__ float nh_fxp_to_float(long long v) { return (float)((double)v * (1.0 / 16777216.0)); } +static precision_t* nh_msg_w0_dev = nullptr; // NH_MSG_FROZEN: pristine trigram table +__global__ void nh_f32_to_precision_kernel(precision_t* __restrict__ dst, const float* __restrict__ src, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) dst[i] = from_float(src[i]); +} __global__ void nh_fxp_to_precision_kernel( precision_t* __restrict__ dst, const long long* __restrict__ src, int n) { int idx = blockIdx.x * blockDim.x + threadIdx.x; @@ -236,6 +378,223 @@ __device__ __forceinline__ int nh_bl_read_i32(const precision_t* p) { } // Decode int16 LE glyph ids into an fp32 index buffer (full grid). +// v5 semantic-class LUTs (host-built once at create, before graph capture) +// (dx,dy) -> sector*4+band, dx in [-78,78], dy in [-20,20]; built on device +// with the exact atan2f expression of the original featurizer (ULP-identical). +static unsigned char* nh_secband_lut_dev = NULL; +__global__ void nh_secband_init_kernel(unsigned char* lut) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= 157 * 41) return; + int dy = idx / 157 - 20, dx = idx % 157 - 78; + int ady = dy < 0 ? -dy : dy, adx = dx < 0 ? -dx : dx; + int cheb = adx > ady ? adx : ady; + float a = atan2f((float)dy, (float)dx) + 3.14159265358979f; + int sct = ((int)(a / 0.78539816339745f)) & 7; + int band = cheb < 3 ? 0 : cheb < 7 ? 1 : cheb < 15 ? 2 : 3; + lut[idx] = (unsigned char)(sct * 4 + band); +} +static unsigned char* nh_locc_lut_dev = NULL; // glyph -> 9-class local id +static unsigned char* nh_terrc_lut_dev = NULL; // glyph -> 17-class terrain id +static void nh_v5_luts_init(void) { + if (nh_locc_lut_dev) return; + unsigned char* loc = (unsigned char*)malloc(NH_GLYPH_VOCAB); + unsigned char* ter = (unsigned char*)malloc(NH_GLYPH_VOCAB); + for (int g = 0; g < NH_GLYPH_VOCAB; g++) { + int ci = (g >= 2359 && g < 2359 + 96) ? g - 2359 : -1; + unsigned char lc = 7; // other + if ((ci >= 1 && ci <= 11) || ci == 17 || ci == 18) lc = 0; // wall/bars/tree + else if (ci >= 19 && ci <= 22) lc = 1; // floor/corridor + else if (ci >= 12 && ci <= 16) lc = 2; // door + else if (ci == 32 || ci == 34 || ci == 41 + || (ci >= 42 && ci <= 64)) lc = 3; // hazard terrain + else if (g >= 0 && g < 1144) lc = 4; // monster + else if ((g >= 1906 && g < 2359) || (g >= 1144 && g < 1525)) lc = 5; // item + else if (ci == 0) lc = 6; // unexplored + if (g == NH_PAD_GLYPH) lc = 8; // off-map + loc[g] = lc; + unsigned char tc = 255; // not counted + if (ci == 23 || ci == 25) tc = 0; // upstairs + else if (ci == 24 || ci == 26) tc = 1; // downstairs + else if (ci == 27) tc = 2; // altar + else if (ci == 31) tc = 3; // fountain + else if (ci == 30) tc = 4; // sink + else if (ci == 29) tc = 5; // throne + else if (ci >= 12 && ci <= 14) tc = 6; // open door + else if (ci == 15 || ci == 16) tc = 7; // closed door + else if (ci >= 42 && ci <= 64) tc = 8; // trap + else if (ci == 28) tc = 9; // grave + else if (ci >= 1 && ci <= 11) tc = 10; // wall + else if (ci == 17 || ci == 18) tc = 11; // bars/tree + else if (ci == 0) tc = 12; // stone + else if (ci >= 19 && ci <= 22) tc = 13; // floor + else if ((g >= 1906 && g < 2359) || (g >= 1144 && g < 1525)) tc = 14; // item + else if (g >= 0 && g < 1144) tc = (g % 381 == 267) ? 16 : 15; // shk / monster + ter[g] = tc; + } + cudaMalloc(&nh_locc_lut_dev, NH_GLYPH_VOCAB); + cudaMalloc(&nh_terrc_lut_dev, NH_GLYPH_VOCAB); + cudaMalloc(&nh_secband_lut_dev, 157 * 41); + nh_secband_init_kernel<<>>(nh_secband_lut_dev); + cudaMemcpy(nh_locc_lut_dev, loc, NH_GLYPH_VOCAB, cudaMemcpyHostToDevice); + cudaMemcpy(nh_terrc_lut_dev, ter, NH_GLYPH_VOCAB, cudaMemcpyHostToDevice); + free(loc); free(ter); +} +static unsigned char* nh_haz_lut_dev = NULL; // species -> hazard bits +static void nh_haz_lut_init(void) { + if (nh_haz_lut_dev) return; + cudaMalloc(&nh_haz_lut_dev, NH_MONS_STATIC_N); + cudaMemcpy(nh_haz_lut_dev, NH_MON_HAZ, NH_MONS_STATIC_N, cudaMemcpyHostToDevice); +} +// semantic crop gather: crop glyph ids -> class -> tiny class embedding. +// Hero cell (crop center) is class 7 (mirrors croplab). +__global__ void nh_loc3_gather_kernel(precision_t* __restrict__ x, + const precision_t* __restrict__ locc_w, const float* __restrict__ crop_glyph, + const unsigned char* __restrict__ lut, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_LOC_IN) return; + int b = t / NH_LOC_IN, j = t % NH_LOC_IN; + int p = j / NH_LOCC_DIM, d = j % NH_LOCC_DIM; + int g = (int)crop_glyph[(int64_t)b * NH_CGRID + p]; + int cls = p == (NH_CGRID / 2) ? 7 : (int)lut[g]; + x[t] = locc_w[cls * NH_LOCC_DIM + d]; +} +__global__ void nh_loc3_scatter_kernel(long long* __restrict__ acc, + const precision_t* __restrict__ dx, const float* __restrict__ crop_glyph, + const unsigned char* __restrict__ lut, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_LOC_IN) return; + int b = t / NH_LOC_IN, j = t % NH_LOC_IN; + int p = j / NH_LOCC_DIM, d = j % NH_LOCC_DIM; + int g = (int)crop_glyph[(int64_t)b * NH_CGRID + p]; + int cls = p == (NH_CGRID / 2) ? 7 : (int)lut[g]; + float v = to_float(dx[t]); + if (v != 0.0f) nh_fxp_atomic_add(&acc[cls * NH_LOCC_DIM + d], v); +} +// terrain featurize (forward-only; features carry no gradients): landmark +// table (12x[seen,dx/78,dy/20,min(d,30)/30]) + sector radar (8x4x17, +// log1p(c)/log1p(1660)). Mirrors terrlab2 exactly: bands digitize([3,7,15]), +// hero cell counts as floor, all 17 classes counted incl. landmarks. +__global__ void nh_terr_feat_kernel(precision_t* __restrict__ tf, + const float* __restrict__ gidx, const precision_t* __restrict__ obs, + const unsigned char* __restrict__ lut, const unsigned char* __restrict__ sblut, + int B) { + int b = blockIdx.x; + if (b >= B) return; + __shared__ int sec[8 * 4 * 17]; + __shared__ unsigned int lmkey[12]; // (cheb << 11) | cell: min == serial first-hit + for (int i = threadIdx.x; i < 8 * 4 * 17; i += blockDim.x) sec[i] = 0; + if (threadIdx.x < 12) lmkey[threadIdx.x] = 0xFFFFFFFFu; + __syncthreads(); + const precision_t* bl = obs + (int64_t)b * NH_OBS_SIZE + NH_BL_OFF; + int hx = nh_bl_read_i32(bl), hy = nh_bl_read_i32(bl + 4); + int hcell = hy * NH_MAPW + hx; + for (int cell = threadIdx.x; cell < NH_MGRID; cell += blockDim.x) { + int g = (int)gidx[(int64_t)b * NH_MGRID + cell]; + int tc = cell == hcell ? 13 : (int)lut[g]; + if (tc == 255) continue; + int dy = cell / NH_MAPW - hy, dx = cell % NH_MAPW - hx; + if (tc < 12) { + int ady = dy < 0 ? -dy : dy, adx = dx < 0 ? -dx : dx; + int cheb = adx > ady ? adx : ady; + atomicMin(&lmkey[tc], ((unsigned int)cheb << 11) | (unsigned int)cell); + } + atomicAdd(&sec[(int)sblut[(dy + 20) * 157 + (dx + 78)] * 17 + tc], 1); + } + __syncthreads(); + precision_t* o = tf + (int64_t)b * NH_TERRF; + if (threadIdx.x < 12) { + int t = threadIdx.x; + unsigned int key = lmkey[t]; + if (key == 0xFFFFFFFFu) { + o[t * 4 + 0] = from_float(0.0f); o[t * 4 + 1] = from_float(0.0f); + o[t * 4 + 2] = from_float(0.0f); o[t * 4 + 3] = from_float(0.0f); + } else { + int cell = (int)(key & 2047u); + int cheb = (int)(key >> 11); + int dy = cell / NH_MAPW - hy, dx = cell % NH_MAPW - hx; + o[t * 4 + 0] = from_float(1.0f); + o[t * 4 + 1] = from_float((float)dx * (1.0f / 78.0f)); + o[t * 4 + 2] = from_float((float)dy * (1.0f / 20.0f)); + o[t * 4 + 3] = from_float((float)(cheb < 30 ? cheb : 30) * (1.0f / 30.0f)); + } + } + float inv_log = 1.0f / logf(1660.0f); + for (int i = threadIdx.x; i < 8 * 4 * 17; i += blockDim.x) + o[48 + i] = from_float(log1pf((float)sec[i]) * inv_log); +} +__global__ void nh_isum_pool_kernel(precision_t* __restrict__ concat, + const precision_t* __restrict__ ih, const float* __restrict__ inv_idx, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_ISUM_DIM) return; + int b = t / NH_ISUM_DIM, d = t % NH_ISUM_DIM; + float acc = 0.0f; + for (int s = 0; s < NH_INV; s++) + if ((int)inv_idx[(int64_t)b * NH_INV + s] != NH_PAD_GLYPH) + acc += to_float(ih[((int64_t)b * NH_INV + s) * NH_ISUM_DIM + d]); + concat[(int64_t)b * NH_CONCAT + NH_ISUM_CONCAT_OFF + d] = from_float(acc * 0.2f); +} +// ---- lab arm kernels: typed streams (deep values + 8-head pools) ---- +// token builder: 48-dim = [e_eff(glyph) 32 | dx dy cheb rank | type flags | +// diff speed | pad]. gid stores the mapped GLYPH id (shared embed table). +__global__ void nh_lab_tok_kernel(precision_t* __restrict__ tok, float* __restrict__ gid, + const precision_t* __restrict__ obs, const precision_t* __restrict__ e_eff, + const unsigned char* __restrict__ haz, int list_off, int is_mon, int B) { + // thread per (b, token, out dim): pure independent writes -> bit-identical + // to the serial per-token version by construction. + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_LABK * NH_LAB_IN) return; + int d = t % NH_LAB_IN; + int bk = t / NH_LAB_IN; + int k = bk % NH_LABK; + const precision_t* e = obs + (int64_t)(bk / NH_LABK) * NH_OBS_SIZE + list_off + k * NH_V3_MONF; + int row = (int)to_float(e[0]) | ((int)to_float(e[1]) << 8); + int g = row <= 0 ? -1 + : is_mon ? row - 1 + : (row < 454 ? NH_OBJ_LO + row - 1 : NH_BODY_OFF + row - 454); + if (d == 0) gid[bk] = (float)g; + precision_t* o = tok + (int64_t)bk * NH_LAB_IN; + if (d < NH_EMBED_DIM) { + o[d] = g >= 0 ? e_eff[(int64_t)g * NH_EMBED_DIM + d] : from_float(0.0f); + return; + } + if (d >= 46) { o[d] = from_float(0.0f); return; } + int dx = (int)to_float(e[2]); if (dx >= 128) dx -= 256; + int dy = (int)to_float(e[3]); if (dy >= 128) dy -= 256; + int f4 = (int)to_float(e[4]), f5 = (int)to_float(e[5]), f6 = (int)to_float(e[6]); + int cheb = abs(dx) > abs(dy) ? abs(dx) : abs(dy); + float v = 0.0f; + switch (d) { + case 32: v = g >= 0 ? fmaxf(fminf((float)dx * (1.0f / 40.0f), 1.0f), -1.0f) : 0.0f; break; + case 33: v = g >= 0 ? fmaxf(fminf((float)dy * (1.0f / 11.0f), 1.0f), -1.0f) : 0.0f; break; + case 34: v = g >= 0 ? fminf((float)cheb, 15.0f) * (1.0f / 15.0f) : 0.0f; break; + case 35: v = g >= 0 ? (float)k * (1.0f / 15.0f) : 0.0f; break; + case 36: v = is_mon ? (g >= 0 && (f4 & 1) ? 1.0f : 0.0f) + : (g >= 0 && (f5 & 1) ? 1.0f : 0.0f); break; + case 37: v = is_mon ? (g >= 0 && (f4 & 8) ? 1.0f : 0.0f) + : (g >= 0 && (f5 & 2) ? 1.0f : 0.0f); break; + case 38: v = is_mon && g >= 0 && (f4 & 4) ? 1.0f : 0.0f; break; + case 39: v = is_mon && g >= 0 && cheb <= 1 ? 1.0f : 0.0f; break; + case 40: v = is_mon && g >= 0 ? fminf((float)f5 * 0.04f, 1.0f) : 0.0f; break; + case 41: v = is_mon && g >= 0 ? fminf((float)f6 * (1.0f / 24.0f), 1.0f) : 0.0f; break; + case 42: case 43: case 44: case 45: { + int hb = is_mon && row > 0 && haz != NULL ? (int)haz[(row - 1) % 381] : 0; + v = (hb & (1 << (d - 42))) ? 1.0f : 0.0f; break; + } + } + o[d] = from_float(v); +} +// embed-table scatter for a stream's dtok buffer (first 32 dims only) +__global__ void nh_lab_dE_scatter_kernel(long long* __restrict__ dE_i, + const precision_t* __restrict__ dtok, const float* __restrict__ gid, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_LABK * NH_EMBED_DIM) return; + int bk = t / NH_EMBED_DIM, d = t % NH_EMBED_DIM; + int g = (int)gid[bk]; + if (g < 0) return; + float v = to_float(dtok[(int64_t)bk * NH_LAB_IN + d]); + if (v != 0.0f) nh_fxp_atomic_add(&dE_i[(int64_t)g * NH_EMBED_DIM + d], v); +} + __global__ void nh_decode_kernel( float* __restrict__ idx, const precision_t* __restrict__ obs, int B) { int t = blockIdx.x * blockDim.x + threadIdx.x; @@ -344,7 +703,7 @@ __global__ void nh_patch_max_kernel( const precision_t* __restrict__ T, const precision_t* __restrict__ b1, const precision_t* __restrict__ w_xy, const precision_t* __restrict__ w2, const precision_t* __restrict__ b2, const float* __restrict__ idx, - const precision_t* __restrict__ obs, int B) { + const precision_t* __restrict__ obs, int pad, int B) { __shared__ float w2s[NH_GLB_HID * NH_P1]; __shared__ float t16s[NH_TOK * NH_P1]; __shared__ float hero[2]; @@ -369,7 +728,7 @@ __global__ void nh_patch_max_kernel( #pragma unroll for (int pos = 0; pos < NH_PCELLS; pos++) { int r = r0 + pos / NH_PW, c = c0 + pos % NH_PW; - int g = (r < NH_MAPH && c < NH_MAPW) ? (int)gi[r * NH_MAPW + c] : NH_PAD_GLYPH; + int g = (r < NH_MAPH && c < NH_MAPW) ? (int)gi[r * NH_MAPW + c] : pad; acc += to_float(T[(int64_t)g * NH_TROW + pos * NH_P1 + k]); } acc = fmaxf(acc, 0.0f); @@ -400,10 +759,217 @@ __global__ void nh_patch_max_kernel( // relu-masked (and b2's grad accumulated) by nh_relu_bias_bwd. dW2 and dt16 // accumulate in fixed-point smem (deterministic), dt16 is relu-masked against // the saved t16 and written back over it. + +__global__ void nh_fill_kernel(precision_t* p, float v, int n); // defined with the decoder below + + + + +// ---- V6-min kernels: statistics (sum|max) + lookups, no attention ---- +// sum+max over n entities, thread per (b, value dim). vid: pad value semantics +// padv >= 0 -> invalid when id == padv (inventory); padv < 0 -> invalid when id < 0 (gid). + + + +__global__ void nh_min_summax_kernel(precision_t* __restrict__ concat, + float* __restrict__ amax, const precision_t* __restrict__ v, + const float* __restrict__ vid, float eps, int n, int padv, int off, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_MV) return; + int b = t / NH_MV, d = t % NH_MV; + float sm = 0.0f, mx = -1e30f; + int bm = -1; + for (int k = 0; k < n; k++) { + float id = vid[(int64_t)b * n + k]; + int bad = padv >= 0 ? ((int)id == padv) : (id < 0.0f); + if (bad) continue; + float vv = to_float(v[((int64_t)b * n + k) * NH_MV + d]); + sm += vv; + if (vv > mx) { mx = vv; bm = k; } + } + amax[t] = (float)bm; + precision_t* dst = concat + (int64_t)b * NH_CONCAT + off; + dst[d] = from_float(sm * eps); + dst[NH_MV + d] = from_float(bm >= 0 ? mx : 0.0f); +} +// dv = broadcast(sum grad) + scatter(max grad); thread per (b, entity, dim). +__global__ void nh_min_dv_kernel(precision_t* __restrict__ dv, + const precision_t* __restrict__ grad_concat, const float* __restrict__ amax, + const float* __restrict__ vid, float eps, int n, int padv, int off, int B) { + int64_t t = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (t >= (int64_t)B * n * NH_MV) return; + int d = t % NH_MV; + int k = (t / NH_MV) % n; + int b = t / ((int64_t)n * NH_MV); + float id = vid[(int64_t)b * n + k]; + int bad = padv >= 0 ? ((int)id == padv) : (id < 0.0f); + const precision_t* gc = grad_concat + (int64_t)b * NH_CONCAT + off; + float g = 0.0f; + if (!bad) { + g = eps * to_float(gc[d]); + if ((int)amax[(int64_t)b * NH_MV + d] == k) g += to_float(gc[NH_MV + d]); + } + dv[t] = from_float(g); +} +// stream gates fwd: token-0 rep (both streams) + underfoot sum (items). +__global__ void nh_min_sgate_kernel(precision_t* __restrict__ concat, + const precision_t* __restrict__ rp, const precision_t* __restrict__ tok, + const float* __restrict__ gid, int underfoot, int off, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_INV_HID) return; + int b = t / NH_INV_HID, d = t % NH_INV_HID; + precision_t* dst = concat + (int64_t)b * NH_CONCAT + off; + float g0 = gid[(int64_t)b * NH_LABK] >= 0.0f + ? to_float(rp[((int64_t)b * NH_LABK) * NH_INV_HID + d]) : 0.0f; + dst[d] = from_float(g0); + if (underfoot) { + float uf = 0.0f; + for (int k = 0; k < NH_LABK; k++) { + if (gid[(int64_t)b * NH_LABK + k] < 0.0f) continue; + if (to_float(tok[((int64_t)b * NH_LABK + k) * NH_LAB_IN + 36]) > 0.5f) + uf += to_float(rp[((int64_t)b * NH_LABK + k) * NH_INV_HID + d]); + } + dst[NH_INV_HID + d] = from_float(uf); + } +} +// stream gates bwd: adds into drp (after its first write). +__global__ void nh_min_sgate_bwd_kernel(precision_t* __restrict__ drp, + const precision_t* __restrict__ grad_concat, const precision_t* __restrict__ tok, + const float* __restrict__ gid, int underfoot, int off, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_INV_HID) return; + int b = t / NH_INV_HID, d = t % NH_INV_HID; + const precision_t* gc = grad_concat + (int64_t)b * NH_CONCAT + off; + if (gid[(int64_t)b * NH_LABK] >= 0.0f) { + precision_t* dr = drp + ((int64_t)b * NH_LABK) * NH_INV_HID + d; + *dr = from_float(to_float(*dr) + to_float(gc[d])); + } + if (underfoot) { + for (int k = 0; k < NH_LABK; k++) { + if (gid[(int64_t)b * NH_LABK + k] < 0.0f) continue; + if (to_float(tok[((int64_t)b * NH_LABK + k) * NH_LAB_IN + 36]) > 0.5f) { + precision_t* dr = drp + ((int64_t)b * NH_LABK + k) * NH_INV_HID + d; + *dr = from_float(to_float(*dr) + to_float(gc[NH_INV_HID + d])); + } + } + } +} +// ---- inventory pass-through kernels (shared by GEN and MIN arms) ---- +// pass-throughs: wielded [r|sfeat], quivered [r|sfeat], worn-profile mean r. +// thread per (b, dim): each output dim replays the serial version's exact +// per-dim op sequence (bf16 round-trip accumulation in slot order), so the +// rewrite is bit-identical to the old one-thread-per-sample kernel. +#ifndef NH_NO_APANEL +// accessory slot of inventory slot k: 0 amulet, 1/2 worn rings by slot order, +// 3 eyewear; -1 otherwise. Shared by forward and backward (same decision). +__device__ __forceinline__ int nh_acc_slot(const float* __restrict__ vid, + const precision_t* __restrict__ sfeat, int b, int k) { + int g = (int)vid[(int64_t)b * NH_INV + k]; + if (g == NH_PAD_GLYPH) return -1; + if (to_float(sfeat[((int64_t)b * NH_INV + k) * NH_SFEAT + 9]) <= 0.5f) return -1; + int ot = g - NH_GLYPH_OBJ_OFF, lo, hi, base, cap; + if (ot >= 178 && ot <= 188) { lo = 178; hi = 188; base = 0; cap = 1; } + else if (ot >= 207 && ot <= 209) { lo = 207; hi = 209; base = 3; cap = 1; } + else if (ot >= 150 && ot <= 177) { lo = 150; hi = 177; base = 1; cap = 2; } + else return -1; + int n = 0; // ordinal among earlier worn slots of the same kind (first owner wins) + for (int j = 0; j < k; j++) { + int gj = (int)vid[(int64_t)b * NH_INV + j]; + if (gj == NH_PAD_GLYPH) continue; + int oj = gj - NH_GLYPH_OBJ_OFF; + if (oj >= lo && oj <= hi + && to_float(sfeat[((int64_t)b * NH_INV + j) * NH_SFEAT + 9]) > 0.5f) n++; + } + return n < cap ? base + n : -1; +} +#endif +__global__ void nh_pass_kernel(precision_t* __restrict__ concat, + const precision_t* __restrict__ inv_out, const precision_t* __restrict__ sfeat, + const float* __restrict__ vid, int nopass, int passoff, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_PASS_DIM) return; + int b = t / NH_PASS_DIM, d = t % NH_PASS_DIM; + precision_t* dst = concat + (int64_t)b * NH_CONCAT + passoff + d; + *dst = from_float(0.0f); + if (nopass == 7) return; +#ifndef NH_NO_APANEL + if (d >= 80 + NH_PWORN_DIM) { // accessory panel: first slot owning index a + if (nopass & 4) return; + int j3 = d - 80 - NH_PWORN_DIM; + int a = j3 / (NH_INV_HID + NH_SFEAT), e = j3 % (NH_INV_HID + NH_SFEAT); + for (int k = 0; k < NH_INV; k++) { + if (nh_acc_slot(vid, sfeat, b, k) != a) continue; + const precision_t* f = sfeat + ((int64_t)b * NH_INV + k) * NH_SFEAT; + const precision_t* r = inv_out + ((int64_t)b * NH_INV + k) * NH_INV_HID; + *dst = e < NH_INV_HID ? r[e] : f[e - NH_INV_HID]; + return; + } + return; + } +#endif + int region = d < 40 ? 0 : d < 80 ? 1 : 2; // wield | quiver | worn + int j = region == 0 ? d : region == 1 ? d - 40 : d - 80; + if (region == 0 && (nopass & 1)) return; + if (region == 1 && (nopass & 2)) return; + if (region == 2 && (nopass & 4)) return; + int fbit = region == 0 ? 10 : region == 1 ? 12 : 9; + float wsum = 0.0f; int nworn = 0; + for (int k = 0; k < NH_INV; k++) { + if ((int)vid[(int64_t)b * NH_INV + k] == NH_PAD_GLYPH) continue; + const precision_t* f = sfeat + ((int64_t)b * NH_INV + k) * NH_SFEAT; + if (to_float(f[fbit]) <= 0.5f) continue; + const precision_t* r = inv_out + ((int64_t)b * NH_INV + k) * NH_INV_HID; + if (region == 2) { nworn++; wsum += to_float(r[j]); } + else { + float v = j < NH_INV_HID ? to_float(r[j]) : to_float(f[j - NH_INV_HID]); + *dst = from_float(to_float(*dst) + v); + } + } + if (region == 2 && nworn > 0) *dst = from_float(wsum / (float)nworn); +} +// pass-through backward (runs AFTER the v-path GEMM's first write of inv_grad): +// adds r-grads at flagged slots. One thread per sample; races impossible. +// thread per (b, slot, dim): per-slot add order (wield, quiver, worn) matches +// the serial version exactly per dim -> bit-identical. +__global__ void nh_passbwd_kernel(precision_t* __restrict__ inv_grad, + const precision_t* __restrict__ grad_concat, const precision_t* __restrict__ sfeat, + const float* __restrict__ vid, int nopass, int passoff, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_INV * NH_INV_HID) return; + int d = t % NH_INV_HID; + int k = (t / NH_INV_HID) % NH_INV; + int b = t / (NH_INV * NH_INV_HID); + if ((int)vid[(int64_t)b * NH_INV + k] == NH_PAD_GLYPH) return; + const precision_t* gp = grad_concat + (int64_t)b * NH_CONCAT + passoff; + const precision_t* f = sfeat + ((int64_t)b * NH_INV + k) * NH_SFEAT; + precision_t* dr = inv_grad + ((int64_t)b * NH_INV + k) * NH_INV_HID + d; + if (!(nopass & 1) && to_float(f[10]) > 0.5f) + *dr = from_float(to_float(*dr) + to_float(gp[d])); + if (!(nopass & 2) && to_float(f[12]) > 0.5f) + *dr = from_float(to_float(*dr) + to_float(gp[40 + d])); + if (!(nopass & 4) && to_float(f[9]) > 0.5f) { + int nworn = 0; + for (int j = 0; j < NH_INV; j++) { + if ((int)vid[(int64_t)b * NH_INV + j] == NH_PAD_GLYPH) continue; + if (to_float(sfeat[((int64_t)b * NH_INV + j) * NH_SFEAT + 9]) > 0.5f) nworn++; + } + if (nworn > 0) + *dr = from_float(to_float(*dr) + to_float(gp[80 + d]) / (float)nworn); + } +#ifndef NH_NO_APANEL + if (!(nopass & 4)) { + int a = nh_acc_slot(vid, sfeat, b, k); + if (a >= 0) + *dr = from_float(to_float(*dr) + + to_float(gp[80 + NH_PWORN_DIM + a * (NH_INV_HID + NH_SFEAT) + d])); + } +#endif +} __global__ void nh_patch_max_bwd_kernel( precision_t* __restrict__ t16_io, long long* __restrict__ dw2_acc, const precision_t* __restrict__ dglb, const precision_t* __restrict__ w2, - const int* __restrict__ argmax, int B) { + const int* __restrict__ argmax, const precision_t* __restrict__ dmean, + const precision_t* __restrict__ dt_extra, int B) { __shared__ float w2s[NH_GLB_HID * NH_P1]; __shared__ float t16s[NH_TOK * NH_P1]; __shared__ long long dt16s[NH_TOK * NH_P1]; @@ -439,8 +1005,10 @@ __global__ void nh_patch_max_bwd_kernel( if (dw2s[i] != 0) atomicAdd((unsigned long long*)&dw2_acc[i], (unsigned long long)dw2s[i]); for (int i = threadIdx.x; i < NH_TOK * NH_P1; i += blockDim.x) { - float v = t16s[i] > 0.0f ? nh_fxp_to_float(dt16s[i]) : 0.0f; - t16_io[(int64_t)b * (NH_TOK * NH_P1) + i] = from_float(v); + float v = nh_fxp_to_float(dt16s[i]); + if (dmean) v += to_float(dmean[(int64_t)b * NH_P1 + (i % NH_P1)]) * (1.0f / NH_TOK); // mean path: uniform over tokens + if (dt_extra) v += to_float(dt_extra[(int64_t)b * (NH_TOK * NH_P1) + i]); // topk path + t16_io[(int64_t)b * (NH_TOK * NH_P1) + i] = from_float(t16s[i] > 0.0f ? v : 0.0f); } } @@ -504,11 +1072,12 @@ __global__ void nh_blstats_kernel( f = d / (1.0f + fabsf(d)); } else f = (float)v * 0.001f; - } else { + } else if (j < NH_F_INTRINS) { // role/race/gender one-hots, already 0/1 // dead under NH_ID_EMBED: identity flows via the embed channel - f = NH_ID_EMBED ? 0.0f - : (float)nh_bl_read_i32(ex + 4*(NH_EX_ROLEOH + (j - NH_F_ROLE))); + f = 0.0f; // identity flows via the embed channel + } else { + f = (float)((nh_bl_read_i32(ex + 4*NH_EX_INTRINS) >> (j - NH_F_INTRINS)) & 1); } // strict [-1,1]: bounds deep-play excursions (AC -15 -> -1.5, hp 300 -> // 1.5, stacked inv counts) — validated neutral-now, deep-safe (n=4) @@ -594,14 +1163,17 @@ __global__ void nh_spkey_kernel(precision_t* __restrict__ keys, precision_t* __restrict__ sp_in, float* __restrict__ sp_idx, const precision_t* __restrict__ spk_w, const precision_t* __restrict__ e_eff, const precision_t* __restrict__ obs, int B) { + // thread per (b, slot, key row): rebuilds the float input (cached loads), + // computes one matvec row in the original MAC order -> bit-identical. int t = blockIdx.x * blockDim.x + threadIdx.x; - if (t >= B * NH_SPELL_SLOTS) return; - int b = t / NH_SPELL_SLOTS, s = t % NH_SPELL_SLOTS; + if (t >= B * NH_SPELL_SLOTS * NH_SPKEY) return; + int r = t % NH_SPKEY; + int bs = t / NH_SPKEY; + int b = bs / NH_SPELL_SLOTS, s = bs % NH_SPELL_SLOTS; const precision_t* src = obs + (int64_t)b * NH_OBS_SIZE + NH_BL_OFF + 4 * (NH_BL_RAW + NH_EXTRA_SHOP + 2 + 1 + 4 * s); int id = (int)to_float(src[0]) | ((int)to_float(src[1]) << 8); int g = id > 0 ? min(id + 1906, NH_GLYPH_VOCAB - 1) : -1; - sp_idx[(int64_t)t] = (float)g; float in[NH_SPIN]; for (int d = 0; d < NH_EMBED_DIM; d++) in[d] = g >= 0 ? to_float(e_eff[(int64_t)g * NH_EMBED_DIM + d]) : 0.0f; @@ -613,15 +1185,15 @@ __global__ void nh_spkey_kernel(precision_t* __restrict__ keys, in[NH_EMBED_DIM + 1] = fminf((float)lev * 0.142857f, 1.0f); in[NH_EMBED_DIM + 2] = fminf((float)fail * 0.01f, 1.0f); in[NH_EMBED_DIM + 3] = fminf((float)know * 0.00005f, 1.0f); - precision_t* inb = sp_in + (int64_t)t * NH_SPIN; - for (int c = 0; c < NH_SPIN; c++) inb[c] = from_float(in[c]); - precision_t* kb = keys + (int64_t)t * NH_SPKEY; - for (int r = 0; r < NH_SPKEY; r++) { - float acc = 0.0f; - for (int c = 0; c < NH_SPIN; c++) - acc += to_float(spk_w[r * NH_SPIN + c]) * in[c]; - kb[r] = from_float(fmaxf(acc, 0.0f)); // relu'd slot rep (inv1 idiom) + if (r == 0) { + sp_idx[(int64_t)bs] = (float)g; + precision_t* inb = sp_in + (int64_t)bs * NH_SPIN; + for (int c = 0; c < NH_SPIN; c++) inb[c] = from_float(in[c]); } + float acc = 0.0f; + for (int c = 0; c < NH_SPIN; c++) + acc += to_float(spk_w[r * NH_SPIN + c]) * in[c]; + keys[(int64_t)bs * NH_SPKEY + r] = from_float(fmaxf(acc, 0.0f)); } // inventory-style pool: project each relu'd slot rep, max over slots, bias, @@ -651,6 +1223,47 @@ __global__ void nh_sppool_kernel(precision_t* __restrict__ concat, concat[(int64_t)b * NH_CONCAT + NH_SPELL_CONCAT_OFF + d] = from_float(v); } +// spell sum-channel forward: masked sum x0.2 of relu'd key projections + 4 +// exact doorstep scalars (all in [0,1]; no spells -> minfail 1, maxlev 0, +// n 0, minret 1). +__global__ void nh_sp_doorstep_kernel(precision_t* __restrict__ concat, + const precision_t* __restrict__ obs, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * 4) return; + int b = t / 4, d = 2 * NH_SPM + t % 4; + float v; + { + const precision_t* ex = obs + (int64_t)b * NH_OBS_SIZE + NH_BL_OFF + NH_BL_RAW * 4; + int mf = 100, ml = 0, n = 0, mr = 20000; + for (int s = 0; s < NH_SPELL_SLOTS; s++) { + const precision_t* q = ex + 4 * (NH_EXTRA_SHOP + 2 + 1 + 4 * s); + int sid = nh_bl_read_i32(q); + if (sid <= 0) continue; + n++; + int lv = nh_bl_read_i32(q + 4), fl = nh_bl_read_i32(q + 8), kn = nh_bl_read_i32(q + 12); + if (fl < mf) mf = fl; + if (lv > ml) ml = lv; + if (kn < mr) mr = kn; + } + int j = d - 2 * NH_SPM; + v = j == 0 ? (float)mf * 0.01f + : j == 1 ? (float)ml * (1.0f / 7.0f) + : j == 2 ? (float)(n > 8 ? 8 : n) * 0.125f + : (float)mr * 0.00005f; + v = fminf(fmaxf(v, 0.0f), 1.0f); + } + concat[(int64_t)b * NH_CONCAT + NH_SPELL_CONCAT_OFF + d] = from_float(v); +} +// dkeys = trunk-MLP grads (pre-staged in dkeys by the spm1 backward GEMM) +// + pointer grads, gated by the key relu +__global__ void nh_sp2_dk_kernel(precision_t* __restrict__ dkeys, + const precision_t* __restrict__ ptr_dkeys, const precision_t* __restrict__ keys, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_SPELL_SLOTS * NH_SPKEY) return; + float v = to_float(dkeys[t]) + (ptr_dkeys ? to_float(ptr_dkeys[t]) : 0.0f); + if (to_float(keys[t]) <= 0.0f) v = 0.0f; + dkeys[t] = from_float(v); +} // backward: scatter the concat-grad spell slice into the fxp dE staging // spell-key backward, stage 1: total per-slot rep grad = pool grad routed // through relu + argmax + spk2 projection, plus the pointer's key grads; @@ -717,6 +1330,7 @@ __global__ void nh_spkey_dE_kernel(long long* __restrict__ dE_i, if (acc != 0.0f) nh_fxp_atomic_add(&dE_i[(int64_t)g * NH_EMBED_DIM + d], acc); } + __global__ void nh_concat_kernel( precision_t* __restrict__ out, const precision_t* __restrict__ loc, const precision_t* __restrict__ glb, const precision_t* __restrict__ inv, @@ -730,12 +1344,12 @@ __global__ void nh_concat_kernel( val = loc[(int64_t)b * NH_LOC_HID + c]; else if (c < NH_LOC_HID + NH_GLB_HID) val = glb[(int64_t)b * NH_GLB_HID + (c - NH_LOC_HID)]; - else if (c < NH_LOC_HID + NH_GLB_HID + NH_INV_POOL) + else if (c < NH_LOC_HID + NH_GLB_HID + NH_INVP_DIM) val = inv[(int64_t)b * NH_INV_POOL + (c - NH_LOC_HID - NH_GLB_HID)]; - else if (c < NH_LOC_HID + NH_GLB_HID + NH_INV_POOL + NH_BL_HID) - val = bl_out[(int64_t)b * NH_BL_HID + (c - NH_LOC_HID - NH_GLB_HID - NH_INV_POOL)]; + else if (c < NH_LOC_HID + NH_GLB_HID + NH_INVP_DIM + NH_BL_HID) + val = bl_out[(int64_t)b * NH_BL_HID + (c - NH_LOC_HID - NH_GLB_HID - NH_INVP_DIM)]; else if (c < NH_MSG_CONCAT_OFF) - val = bl_feats[(int64_t)b * NH_BL_FEAT + (c - NH_LOC_HID - NH_GLB_HID - NH_INV_POOL - NH_BL_HID)]; + val = bl_feats[(int64_t)b * NH_BL_FEAT + (c - NH_LOC_HID - NH_GLB_HID - NH_INVP_DIM - NH_BL_HID)]; else if (c < NH_SPELL_CONCAT_OFF) val = msg[(int64_t)b * NH_MSG_HID + (c - NH_MSG_CONCAT_OFF)]; else @@ -793,8 +1407,10 @@ __global__ void nh_inv_sfeat_kernel(precision_t* __restrict__ out, for (int c = 0; c < 4; c++) f[c] = from_float(st[0] == c ? 1.0f : 0.0f); int spe_known = st[1] != -128; f[4] = from_float((float)spe_known); - f[5] = from_float(spe_known ? (float)st[1] * 0.1f : 0.0f); - f[6] = from_float(log1pf(fmaxf((float)st[2], 0.0f)) * 0.5f); // guard vs unclamped export drift + // strict [-1,1]: spe clamp/7; quantity LINEAR capped (log-squash + // regressed ammo-count .88->.76 across two seeds — counts want scale) + f[5] = from_float(spe_known ? fmaxf(fminf((float)st[1], 7.0f), -7.0f) * (1.0f / 7.0f) : 0.0f); + f[6] = from_float(fminf(fmaxf((float)st[2], 0.0f), 30.0f) * (1.0f / 30.0f)); f[7] = from_float((float)st[3] * (1.0f / 3.0f)); f[8] = from_float((float)st[4] * (1.0f / 3.0f)); for (int c = 0; c < 7; c++) @@ -949,6 +1565,33 @@ __global__ void nh_relu_bias_bwd_kernel( if (sdata[j] != 0) atomicAdd((unsigned long long*)&bias_acc[j], (unsigned long long)sdata[j]); } +// bias-grad column sums without a relu mask (NH_ENT_LINPOOL: linear last entity layer) +__global__ void nh_bias_bwd_kernel( + precision_t* __restrict__ grad, const precision_t* __restrict__ out, + long long* __restrict__ bias_acc, int64_t total, int dim) { + extern __shared__ long long sdata[]; + for (int j = threadIdx.x; j < dim; j += blockDim.x) sdata[j] = 0; + __syncthreads(); + int64_t i0 = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + int64_t stride = (int64_t)gridDim.x * blockDim.x; + float acc = 0.0f; + for (int64_t i = i0; i < total; i += stride) acc += to_float(grad[i]); + if (acc != 0.0f) nh_fxp_atomic_add(&sdata[(int)(i0 % dim)], acc); + __syncthreads(); + for (int j = threadIdx.x; j < dim; j += blockDim.x) + if (sdata[j] != 0) atomicAdd((unsigned long long*)&bias_acc[j], (unsigned long long)sdata[j]); +} +// NH_ENT_LINPOOL: the last per-entity layer (inventory / monster / item / spell MLPs) +// is linear before sum|max pooling; default keeps the relu. +#ifdef NH_ENT_LINPOOL +#define NH_ENT_LAST(A, W, O, BIAS) do { puf_mm((A), (W), (O), stream); \ + int64_t nn_ = numel((O)->shape); int dd_ = (int)(O)->shape[ndim((O)->shape) - 1]; \ + nh_bias_kernel<<>>((O)->data, (BIAS).data, (int)nn_, dd_); } while (0) +#define NH_ENT_LAST_BWD nh_bias_bwd_kernel +#else +#define NH_ENT_LAST(A, W, O, BIAS) NH_MM_BR(A, W, O, BIAS) +#define NH_ENT_LAST_BWD nh_relu_bias_bwd_kernel +#endif static inline int nh_colsum_grid(int64_t total, int dim) { int64_t g = (total + BLOCK_SIZE - 1) / BLOCK_SIZE; if (g > 1024) g = 1024; @@ -1069,7 +1712,7 @@ __global__ void nh_dE_scatter_kernel( __global__ void nh_dT_patch_scatter_kernel( long long* __restrict__ dT_i, const precision_t* __restrict__ dt16, const float* __restrict__ idx, const int* __restrict__ hot_map, - const int* __restrict__ hot_list, const int* __restrict__ hot_n, int B) { + const int* __restrict__ hot_list, const int* __restrict__ hot_n, int pad, int B) { extern __shared__ long long acc_s[]; // NH_HOT_G x NH_TROW for (int i = threadIdx.x; i < NH_HOT_G * NH_TROW; i += blockDim.x) acc_s[i] = 0; @@ -1089,7 +1732,7 @@ __global__ void nh_dT_patch_scatter_kernel( #pragma unroll for (int pos = 0; pos < NH_PCELLS; pos++) { int r = r0 + pos / NH_PW, c = c0 + pos % NH_PW; - int gl = (r < NH_MAPH && c < NH_MAPW) ? (int)gi[r * NH_MAPW + c] : NH_PAD_GLYPH; + int gl = (r < NH_MAPH && c < NH_MAPW) ? (int)gi[r * NH_MAPW + c] : pad; int slot = hot_map[gl]; if (slot >= 0) atomicAdd((unsigned long long*)&acc_s[slot * NH_TROW + pos * NH_P1 + k], q); else atomicAdd((unsigned long long*)&dT_i[(int64_t)gl * NH_TROW + pos * NH_P1 + k], q); @@ -1140,12 +1783,11 @@ __global__ void nh_fxp_add_rows_kernel( // Seed the grid-glyph histogram with the static pad-glyph count (edge-patch // positions past the map read the pad row; the hist over glyph_idx can't see // them, but the rows-cast guard and hot selection must). -__global__ void nh_count_pad_kernel(int* __restrict__ counts, int B) { +__global__ void nh_count_pad_kernel(int* __restrict__ counts, int pad, int B) { if (threadIdx.x == 0 && blockIdx.x == 0) - counts[NH_PAD_GLYPH] += NH_PAD_PER_SAMPLE * B; + counts[pad] += NH_PAD_PER_SAMPLE * B; } -#if NH_ID_EMBED // identity embeddings: indices recovered from the one-hot obs block (align // from blstats), table rows copied raw into the concat tail (bl-feats idiom) __global__ void nh_idemb_kernel(precision_t* __restrict__ concat, @@ -1200,21 +1842,25 @@ __global__ void nh_idemb_grad_kernel(precision_t* __restrict__ role_g, } out[row * wdt + d] = from_float(acc); } -#endif // encoder structs struct NethackEncoderWeights { Prec embed_w, ekind_w, esub_w, loc_w, loc_b; - Prec glb1_w, glb1_xy, glb1_b, glb2_w, glb2_b; - Prec inv1_w, inv1_b, inv1s_w, invt_w, inv2_w, inv2_b; + Prec loc2_w, loc2_b; // (NH_LOC_HID, NH_LOC_H1), (NH_LOC_HID) + Prec terr1_w, terr1_b, terr2_w, terr2_b; // terrain MLP 592->256->128 + Prec locc_w; // (NH_LOCC_CLASSES, NH_LOCC_DIM) local class table + Prec inv1_w, inv1_b, inv1s_w, invt_w; Prec bl_w, bl_b, proj_w, proj_b; Prec msg_w; // trigram embedding table (NH_MSG_VOCAB, NH_MSG_HID) Prec spk_w; // spell slot-rep projection (NH_SPKEY, NH_SPIN) - Prec spk2_w, spk2_b; // spell pool projection + bias (inv2 idiom) -#if NH_ID_EMBED + Prec spm1_w, spm1_b, spm2_w, spm2_b; // spell per-slot MLP 16->64->64 Prec ide_role_w, ide_race_w, ide_gend_w, ide_algn_w; // identity tables -#endif + Prec mv1_w, mv1_b, mv2_w, mv2_b; // inventory MLP 16->64->64 + Prec mr_w, mr_b; // monster rep 48->16 + Prec mm1_w, mm1_b, mm2_w, mm2_b; // monster MLP 16->64->64 + Prec ir_w, ir_b; // item rep 48->16 + Prec im1_w, im1_b, im2_w, im2_b; // item MLP 16->64->64 int obs_size, hidden; }; @@ -1222,52 +1868,60 @@ struct NethackEncoderActivations { Float glyph_idx, crop_glyph; // decoded grid + crop glyph ids Prec e_eff; // materialized E_res + E_kind + E_sub Prec x_local; // crop embeds (grad aliases it) - Prec w_perm, glyph_T; // fused embed+flatten+glb1 table - Prec t16; // relu'd patch tokens (dt16 overwrites) - Prec dxy; // per-token hero offsets (w_xy wgrad) - Int tok_argmax; // winning token per (sample, out dim) + Prec terr_tf; // (B, NH_TERRF) featurized terrain (fwd-only input) + Prec terr_h, terr_dh; // relu'd hidden (B, NH_TERR_H1) + its grad + Long terr1b_acc; // fixed-point terr1_b accumulator + Prec terr1_wgrad, terr1_bgrad, terr2_wgrad, terr2_bgrad; + Long locc_acc; // fixed-point class-table accumulator + Prec locc_wgrad; Float inv_idx; // inventory slot glyph ids Float spell_idx; // per-slot book glyphs (-1 = empty slot) Prec spk_in, spk_keys; // spell-key inputs (B, 8*36) + relu'd reps (B, 8*16) - Prec spk_dkeys; // total per-slot rep grads (pool + pointer) - Int spk_amax; // pool argmax slot per (sample, pool dim) - Prec spk_pool; // pooled forward values (relu gate for backward; - // concat is grad-aliased and unreadable then) + Prec spk_dkeys; // per-slot key grads (pointer; +pool under !SPELL2) + Prec sph1, spdh1, spv, spdv; // (B, 8*NH_SPM) spell MLP hidden/values + grads + Float spvmax; // (B, NH_SPM) argmax winners + Long spm1b_acc, spm2b_acc; + Prec spm1_wgrad, spm1_bgrad, spm2_wgrad, spm2_bgrad; Float invt_idx; // discovered-type glyph ids (pad = unknown) Prec inv_sfeat; // per-slot state features (B, 55*NH_SFEAT) Prec inv_T, inv_out; // fused inv table + relu'd flat slots Prec invt_T; // fused discovered-type table - Prec inv_pool; // pooled 128-dim trunk summary - Int inv_amax; // winning slot per (sample, pool dim) Prec loc_out, glb_out; Prec bl_feats, bl_out; Float msg_ids; // per-position trigram bucket ids (-1 pad) Prec msg_out; // normalized trigram-bag summary (B, NH_MSG_HID) Prec concat, out; Prec loc_grad, glb_grad, inv_grad, bl_grad; // contiguous concat slices - Prec inv_pool_grad; // pooled-summary slice of concat grad - Prec dT, dw_perm; // dT table + permuted glb1 wgrad + Prec mh1, mdh1, mvv, mdv; // inv MLP hidden + values (+grads) + Float mvmax; // (B,64) inv argmax + Prec mrp, mdrp, mh1m, mdh1m, mvm, mdvm; Float mvmaxm; // monster stream + Prec irp, idrp, mh1i, mdh1i, mvi, mdvi; Float mvmaxi; // item stream + Long mv1b_acc, mv2b_acc, mrb_acc, mm1b_acc, mm2b_acc, irb_acc, im1b_acc, im2b_acc; + Prec mv1_wgrad, mv1_bgrad, mv2_wgrad, mv2_bgrad; + Prec mr_wgrad, mr_bgrad, mm1_wgrad, mm1_bgrad, mm2_wgrad, mm2_bgrad; + Prec ir_wgrad, ir_bgrad, im1_wgrad, im1_bgrad, im2_wgrad, im2_bgrad; Prec dTinv, dE_tmp; // inv-table grad + its dE staging Prec dTtrue; // discovered-type table grad - Long dT_i, dTinv_i, dTtrue_i; // fixed-point dT scatter staging + Long dTinv_i, dTtrue_i; // fixed-point dT scatter staging Long dE_i; // fixed-point local embed-grad staging - Long dw2_acc; // fixed-point glb2 wgrad staging - Long dw2i_acc; // fixed-point inv2 wgrad staging Long dmsg_acc; // fixed-point trigram-table wgrad staging Long bias_acc; // fixed-point bias grads: proj | loc | glb2 | bl | glb1 | inv1 | inv2 - Int sort_local, sort_grid; // counts | hot_map | hot_list | hot_n Prec embed_wgrad, ekind_wgrad, esub_wgrad, loc_wgrad, loc_bgrad; - Prec glb1_wgrad, glb1_xygrad, glb1_bgrad, glb2_wgrad, glb2_bgrad; - Prec inv1_wgrad, inv1_bgrad, inv1s_wgrad, invt_wgrad, inv2_wgrad, inv2_bgrad; + Prec loc2_wgrad, loc2_bgrad; + Prec loc_h1, loc_h1_grad; // relu'd hidden (B, NH_LOC_H1) and its grad + Long loc1b_acc; // fixed-point bias accumulator for loc_b + Prec inv1_wgrad, inv1_bgrad, inv1s_wgrad, invt_wgrad; Prec bl_wgrad, bl_bgrad, proj_wgrad, proj_bgrad; - Prec msg_wgrad, spk_wgrad, spk2_wgrad, spk2_bgrad; -#if NH_ID_EMBED + Prec msg_wgrad, spk_wgrad; Float ide_idx; // per-sample [role, race, gend, align] saved for backward Prec ide_role_wgrad, ide_race_wgrad, ide_gend_wgrad, ide_algn_wgrad; -#endif + Prec lm_tok; Float lm_gid; // monster stream fwd + Prec li_tok; Float li_gid; // item stream fwd }; static NethackEncoderWeights* nethack_encoder_create(int obs_size, int hidden) { + nh_v5_luts_init(); + nh_haz_lut_init(); if (obs_size != NH_OBS_SIZE) { fprintf(stderr, "nethack encoder: obs size %d != expected %d " "(env obs layout out of sync with ocean/nethack/nethack.cu?)\n", @@ -1306,20 +1960,18 @@ static Prec nethack_encoder_forward(void* w, void* activations, Prec input, cuda a->glyph_idx.data, input.data, B); nh_crop_kernel<<>>( a->crop_glyph.data, a->glyph_idx.data, input.data, B); - - nh_local_gather_kernel<<>>( - a->x_local.data, a->e_eff.data, a->crop_glyph.data, B); - puf_mm(&a->x_local, &ew->loc_w, &a->loc_out, stream); - nh_bias_relu_kernel<<>>( - a->loc_out.data, ew->loc_b.data, B * NH_LOC_HID, NH_LOC_HID); - - nh_permute_g1_kernel<<>>( - a->w_perm.data, ew->glb1_w.data); - puf_mm(&a->e_eff, &a->w_perm, &a->glyph_T, stream); - nh_patch_max_kernel<<>>( - a->glb_out.data, a->t16.data, a->dxy.data, a->tok_argmax.data, - a->glyph_T.data, ew->glb1_b.data, ew->glb1_xy.data, ew->glb2_w.data, - ew->glb2_b.data, a->glyph_idx.data, input.data, B); + nh_loc3_gather_kernel<<>>( + a->x_local.data, ew->locc_w.data, a->crop_glyph.data, nh_locc_lut_dev, B); + NH_MM_BR(&a->x_local, &ew->loc_w, &a->loc_h1, ew->loc_b); + NH_MM_BR(&a->loc_h1, &ew->loc2_w, &a->loc_out, ew->loc2_b); + + // terrain branch replaces the patch encoder: featurize (fwd-only, no + // input grads) -> 592 -> 256 -> 128 into the glb slot + nh_terr_feat_kernel<<>>( + a->terr_tf.data, a->glyph_idx.data, input.data, nh_terrc_lut_dev, + nh_secband_lut_dev, B); + NH_MM_BR(&a->terr_tf, &ew->terr1_w, &a->terr_h, ew->terr1_b); + NH_MM_BR(&a->terr_h, &ew->terr2_w, &a->glb_out, ew->terr2_b); nh_inv_decode_kernel<<>>( a->inv_idx.data, input.data, B, NH_INV_OFF); @@ -1334,63 +1986,125 @@ static Prec nethack_encoder_forward(void* w, void* activations, Prec input, cuda nh_inv_gather_kernel<<>>( a->inv_out.data, a->inv_T.data, ew->inv1_b.data, ew->inv1s_w.data, a->inv_sfeat.data, a->inv_idx.data, a->invt_T.data, a->invt_idx.data, B); - nh_inv_max_kernel<<>>( - a->inv_pool.data, a->inv_amax.data, a->inv_out.data, - ew->inv2_w.data, ew->inv2_b.data, B); + { Prec invf = {.data = a->inv_out.data, .shape = {B * NH_INV, NH_INV_HID}}; + Prec h1f = {.data = a->mh1.data, .shape = {B * NH_INV, NH_MV}}; + NH_MM_BR(&invf, &ew->mv1_w, &h1f, ew->mv1_b); } + { Prec h1f = {.data = a->mh1.data, .shape = {B * NH_INV, NH_MV}}; + Prec vf = {.data = a->mvv.data, .shape = {B * NH_INV, NH_MV}}; + NH_ENT_LAST(&h1f, &ew->mv2_w, &vf, ew->mv2_b); } + nh_min_summax_kernel<<>>( + a->concat.data, a->mvmax.data, a->mvv.data, a->inv_idx.data, + 0.2f, NH_INV, NH_PAD_GLYPH, NH_MINV_OFF, B); +#ifdef NH_NOPASS + nh_pass_kernel<<>>( + a->concat.data, a->inv_out.data, a->inv_sfeat.data, a->inv_idx.data, + 7, NH_MINV_OFF + 2 * NH_MV, B); +#else + nh_pass_kernel<<>>( + a->concat.data, a->inv_out.data, a->inv_sfeat.data, a->inv_idx.data, + 0, NH_MINV_OFF + 2 * NH_MV, B); +#endif nh_blstats_kernel<<>>( a->bl_feats.data, input.data, B); - puf_mm(&a->bl_feats, &ew->bl_w, &a->bl_out, stream); - nh_bias_relu_kernel<<>>( - a->bl_out.data, ew->bl_b.data, B * NH_BL_HID, NH_BL_HID); + NH_MM_BR(&a->bl_feats, &ew->bl_w, &a->bl_out, ew->bl_b); +#ifdef NH_MSG_FROZEN + cudaMemcpyAsync(ew->msg_w.data, nh_msg_w0_dev, (size_t)NH_MSG_VOCAB * NH_MSG_HID * sizeof(precision_t), cudaMemcpyDeviceToDevice, stream); +#endif nh_msg_ids_kernel<<>>( a->msg_ids.data, input.data, B); nh_msg_pool_kernel<<>>( a->msg_out.data, ew->msg_w.data, a->msg_ids.data, B); nh_concat_kernel<<>>( - a->concat.data, a->loc_out.data, a->glb_out.data, a->inv_pool.data, + a->concat.data, a->loc_out.data, a->glb_out.data, + (const precision_t*)NULL, // invpool slice is width-0 a->bl_out.data, a->bl_feats.data, a->msg_out.data, B); - nh_spkey_kernel<<>>( + { // typed streams: tok -> deep values (48->64->64) -> 8-head pool + nh_lab_tok_kernel<<>>( + a->lm_tok.data, a->lm_gid.data, input.data, a->e_eff.data, + nh_haz_lut_dev, + NH_TOKM_OFF, 1, B); + { Prec tokf = {.data = a->lm_tok.data, .shape = {B * NH_LABK, NH_LAB_IN}}; + Prec rpf = {.data = a->mrp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + NH_MM_BR(&tokf, &ew->mr_w, &rpf, ew->mr_b); } + { Prec rpf = {.data = a->mrp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + Prec h1f = {.data = a->mh1m.data, .shape = {B * NH_LABK, NH_MV}}; + NH_MM_BR(&rpf, &ew->mm1_w, &h1f, ew->mm1_b); } + { Prec h1f = {.data = a->mh1m.data, .shape = {B * NH_LABK, NH_MV}}; + Prec vf = {.data = a->mvm.data, .shape = {B * NH_LABK, NH_MV}}; + NH_ENT_LAST(&h1f, &ew->mm2_w, &vf, ew->mm2_b); } + nh_min_summax_kernel<<>>( + a->concat.data, a->mvmaxm.data, a->mvm.data, a->lm_gid.data, + 0.25f, NH_LABK, -1, NH_MLM_OFF, B); + nh_min_sgate_kernel<<>>( + a->concat.data, a->mrp.data, a->lm_tok.data, a->lm_gid.data, + 0, NH_MLM_OFF + 2 * NH_MV, B); + nh_lab_tok_kernel<<>>( + a->li_tok.data, a->li_gid.data, input.data, a->e_eff.data, NULL, NH_TOKI_OFF, 0, B); + { Prec tokf = {.data = a->li_tok.data, .shape = {B * NH_LABK, NH_LAB_IN}}; + Prec rpf = {.data = a->irp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + NH_MM_BR(&tokf, &ew->ir_w, &rpf, ew->ir_b); } + { Prec rpf = {.data = a->irp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + Prec h1f = {.data = a->mh1i.data, .shape = {B * NH_LABK, NH_MV}}; + NH_MM_BR(&rpf, &ew->im1_w, &h1f, ew->im1_b); } + { Prec h1f = {.data = a->mh1i.data, .shape = {B * NH_LABK, NH_MV}}; + Prec vf = {.data = a->mvi.data, .shape = {B * NH_LABK, NH_MV}}; + NH_ENT_LAST(&h1f, &ew->im2_w, &vf, ew->im2_b); } + nh_min_summax_kernel<<>>( + a->concat.data, a->mvmaxi.data, a->mvi.data, a->li_gid.data, + 0.25f, NH_LABK, -1, NH_MLI_OFF, B); + nh_min_sgate_kernel<<>>( + a->concat.data, a->irp.data, a->li_tok.data, a->li_gid.data, + 1, NH_MLI_OFF + 2 * NH_MV, B); + } + nh_spkey_kernel<<>>( a->spk_keys.data, a->spk_in.data, a->spell_idx.data, ew->spk_w.data, a->e_eff.data, input.data, B); - nh_sppool_kernel<<>>( - a->concat.data, a->spk_pool.data, a->spk_amax.data, a->spk_keys.data, - ew->spk2_w.data, ew->spk2_b.data, B); -#if NH_ID_EMBED + { Prec kf = {.data = a->spk_keys.data, .shape = {B * NH_SPELL_SLOTS, NH_SPKEY}}; + Prec hf = {.data = a->sph1.data, .shape = {B * NH_SPELL_SLOTS, NH_SPM}}; + NH_MM_BR(&kf, &ew->spm1_w, &hf, ew->spm1_b); } + { Prec hf = {.data = a->sph1.data, .shape = {B * NH_SPELL_SLOTS, NH_SPM}}; + Prec vf = {.data = a->spv.data, .shape = {B * NH_SPELL_SLOTS, NH_SPM}}; + NH_ENT_LAST(&hf, &ew->spm2_w, &vf, ew->spm2_b); } + nh_min_summax_kernel<<>>( + a->concat.data, a->spvmax.data, a->spv.data, a->spell_idx.data, + 0.25f, NH_SPELL_SLOTS, -1, NH_SPELL_CONCAT_OFF, B); + nh_sp_doorstep_kernel<<>>( + a->concat.data, input.data, B); nh_idemb_kernel<<>>( a->concat.data, a->ide_idx.data, input.data, ew->ide_role_w.data, ew->ide_race_w.data, ew->ide_gend_w.data, ew->ide_algn_w.data, B); -#endif - puf_mm(&a->concat, &ew->proj_w, &a->out, stream); - nh_bias_relu_kernel<<hidden), BLOCK_SIZE, 0, stream>>>( - a->out.data, ew->proj_b.data, B * ew->hidden, ew->hidden); + NH_MM_BR(&a->concat, &ew->proj_w, &a->out, ew->proj_b); return a->out; } +// packed bias-acc slot widths: glb1 slot vanishes under TERR (terr1_b has its +// own acc), inv2 slot vanishes under INV2 (max-pool path deleted) +#define NH_BACC_GLB1 0 +#define NH_BACC_INVP 0 static void nethack_encoder_backward(void* w, void* activations, Prec grad, cudaStream_t stream) { NethackEncoderWeights* ew = (NethackEncoderWeights*)w; NethackEncoderActivations* a = (NethackEncoderActivations*)activations; int B = grad.shape[0], H = ew->hidden; - // fixed-point bias-grad accumulators: [proj H | loc 256 | glb2 128 | bl 64 | glb1 16 | inv1 32 | inv2 128] + // fixed-point bias-grad accumulators: [proj H | loc 256 | glb2/terr2 128 | bl 64 | glb1 16? | inv1 | inv2 128?] long long* bacc = (long long*)a->bias_acc.data; - cudaMemsetAsync(bacc, 0, (H + NH_LOC_HID + NH_GLB_HID + NH_BL_HID + NH_P1 + NH_INV_HID + NH_INV_POOL) * sizeof(long long), stream); + cudaMemsetAsync(bacc, 0, (H + NH_LOC_HID + NH_GLB_HID + NH_BL_HID + NH_BACC_GLB1 + NH_INV_HID + NH_BACC_INVP) * sizeof(long long), stream); nh_relu_bias_bwd_kernel<<>>( grad.data, a->out.data, bacc, (int64_t)B * H, H); puf_mm_tn(&grad, &a->concat, &a->proj_wgrad, stream); + Prec grad_concat = {.data = a->concat.data, .shape = {B, NH_CONCAT}}; puf_mm_nn(&grad, &ew->proj_w, &grad_concat, stream); -#if NH_ID_EMBED nh_idemb_grad_kernel<<>>( a->ide_role_wgrad.data, a->ide_race_wgrad.data, a->ide_gend_wgrad.data, a->ide_algn_wgrad.data, grad_concat.data, a->ide_idx.data, B); -#endif // Local view: wgrad against saved x_local, then the input grad overwrites // x_local in place before scattering into the embed table. @@ -1399,9 +2113,18 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda nh_relu_bias_bwd_kernel<<>>( a->loc_grad.data, a->loc_out.data, bacc + H, (int64_t)B * NH_LOC_HID, NH_LOC_HID); Prec locg = {.data = a->loc_grad.data, .shape = {B, NH_LOC_HID}}; - puf_mm_tn(&locg, &a->x_local, &a->loc_wgrad, stream); + // second layer: dW2 = dout^T h1, dh1 = dout W2 (in place over h1 after wgrad), relu mask + own bias acc + puf_mm_tn(&locg, &a->loc_h1, &a->loc2_wgrad, stream); + Prec dh1 = {.data = a->loc_h1_grad.data, .shape = {B, NH_LOC_H1}}; + puf_mm_nn(&locg, &ew->loc2_w, &dh1, stream); + cudaMemsetAsync(a->loc1b_acc.data, 0, NH_LOC_H1 * sizeof(long long), stream); + nh_relu_bias_bwd_kernel<<>>( + dh1.data, a->loc_h1.data, (long long*)a->loc1b_acc.data, (int64_t)B * NH_LOC_H1, NH_LOC_H1); + nh_fxp_to_precision_kernel<<>>( + a->loc_bgrad.data, (long long*)a->loc1b_acc.data, NH_LOC_H1); + puf_mm_tn(&dh1, &a->x_local, &a->loc_wgrad, stream); Prec dx_local = {.data = a->x_local.data, .shape = {B, NH_LOC_IN}}; - puf_mm_nn(&locg, &ew->loc_w, &dx_local, stream); + puf_mm_nn(&dh1, &ew->loc_w, &dx_local, stream); // Global view: relu mask + b2 grad, then the fused max backward (dW2 via // fixed-point staging, dt16 overwrites t16), then b1's column sum. @@ -1409,64 +2132,91 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda a->glb_grad.data, grad_concat.data, B, NH_CONCAT, NH_LOC_HID, NH_GLB_HID); nh_relu_bias_bwd_kernel<<>>( a->glb_grad.data, a->glb_out.data, bacc + H + NH_LOC_HID, (int64_t)B * NH_GLB_HID, NH_GLB_HID); - cudaMemsetAsync(a->dw2_acc.data, 0, NH_GLB_HID * NH_P1 * sizeof(long long), stream); - nh_patch_max_bwd_kernel<<>>( - a->t16.data, (long long*)a->dw2_acc.data, a->glb_grad.data, - ew->glb2_w.data, a->tok_argmax.data, B); - nh_fxp_to_precision_kernel<<>>( - a->glb2_wgrad.data, (long long*)a->dw2_acc.data, NH_GLB_HID * NH_P1); - nh_col_sum_kernel<<>>( - bacc + H + NH_LOC_HID + NH_GLB_HID + NH_BL_HID, a->t16.data, (int64_t)B * NH_TOK * NH_P1, NH_P1); - // (dx,dy) weight slice: dW_xy = dt16^T @ dxy (tall-K, 16x2 output) - Prec dt16v = {.data = a->t16.data, .shape = {B * NH_TOK, NH_P1}}; - Prec dxyv = {.data = a->dxy.data, .shape = {B * NH_TOK, 2}}; - puf_mm_tn(&dt16v, &dxyv, &a->glb1_xygrad, stream); + // terrain MLP backward: features are inputs (no grad past terr_tf); the + // bacc glb2 slot above carries terr2_b's grad. + { Prec glbg = {.data = a->glb_grad.data, .shape = {B, NH_GLB_HID}}; + puf_mm_tn(&glbg, &a->terr_h, &a->terr2_wgrad, stream); + Prec tdh = {.data = a->terr_dh.data, .shape = {B, NH_TERR_H1}}; + puf_mm_nn(&glbg, &ew->terr2_w, &tdh, stream); + cudaMemsetAsync(a->terr1b_acc.data, 0, NH_TERR_H1 * sizeof(long long), stream); + nh_relu_bias_bwd_kernel<<>>( + tdh.data, a->terr_h.data, (long long*)a->terr1b_acc.data, (int64_t)B * NH_TERR_H1, NH_TERR_H1); + nh_fxp_to_precision_kernel<<>>( + a->terr1_bgrad.data, (long long*)a->terr1b_acc.data, NH_TERR_H1); + puf_mm_tn(&tdh, &a->terr_tf, &a->terr1_wgrad, stream); } // Inventory branch: slice the pooled-summary grad, relu-mask it (inv2 // bias grad rides along), backprop the fused max into inv_grad (dW2 via // fixed-point staging), then the per-slot relu mask + inv1 bias, dT_inv // scatter by slot glyph, dinv1_w = dT_inv^T @ E; dE added at the end. - nh_slice_kernel<<>>( - a->inv_pool_grad.data, grad_concat.data, B, NH_CONCAT, NH_LOC_HID + NH_GLB_HID, NH_INV_POOL); - nh_relu_bias_bwd_kernel<<>>( - a->inv_pool_grad.data, a->inv_pool.data, - bacc + H + NH_LOC_HID + NH_GLB_HID + NH_BL_HID + NH_P1 + NH_INV_HID, - (int64_t)B * NH_INV_POOL, NH_INV_POOL); - cudaMemsetAsync(a->dw2i_acc.data, 0, NH_INV_POOL * NH_INV_HID * sizeof(long long), stream); - nh_inv_max_bwd_kernel<<>>( - a->inv_grad.data, (long long*)a->dw2i_acc.data, a->inv_pool_grad.data, - ew->inv2_w.data, a->inv_out.data, a->inv_amax.data, B); - nh_fxp_to_precision_kernel<<>>( - a->inv2_wgrad.data, (long long*)a->dw2i_acc.data, NH_INV_POOL * NH_INV_HID); + // sum channel backward: dih = 0.2*occ*g_isum -> relu gate (+isum bias acc) + // -> isum wgrad; its dinv via mm_nn is the FIRST writer of inv_grad + // (mirrors the max-bwd's write-all contract); wield readout adds next. + cudaMemsetAsync(a->mv1b_acc.data, 0, NH_MV * sizeof(long long), stream); + cudaMemsetAsync(a->mv2b_acc.data, 0, NH_MV * sizeof(long long), stream); + nh_min_dv_kernel<<>>( + a->mdv.data, grad_concat.data, a->mvmax.data, a->inv_idx.data, + 0.2f, NH_INV, NH_PAD_GLYPH, NH_MINV_OFF, B); + NH_ENT_LAST_BWD<<>>( + a->mdv.data, a->mvv.data, (long long*)a->mv2b_acc.data, + (int64_t)B * NH_INV * NH_MV, NH_MV); + nh_fxp_to_precision_kernel<<>>( + a->mv2_bgrad.data, (long long*)a->mv2b_acc.data, NH_MV); + { Prec dvf = {.data = a->mdv.data, .shape = {B * NH_INV, NH_MV}}; + Prec h1f = {.data = a->mh1.data, .shape = {B * NH_INV, NH_MV}}; + puf_mm_tn(&dvf, &h1f, &a->mv2_wgrad, stream); + Prec dh1f = {.data = a->mdh1.data, .shape = {B * NH_INV, NH_MV}}; + puf_mm_nn(&dvf, &ew->mv2_w, &dh1f, stream); } + nh_relu_bias_bwd_kernel<<>>( + a->mdh1.data, a->mh1.data, (long long*)a->mv1b_acc.data, + (int64_t)B * NH_INV * NH_MV, NH_MV); + nh_fxp_to_precision_kernel<<>>( + a->mv1_bgrad.data, (long long*)a->mv1b_acc.data, NH_MV); + { Prec dh1f = {.data = a->mdh1.data, .shape = {B * NH_INV, NH_MV}}; + Prec invf = {.data = a->inv_out.data, .shape = {B * NH_INV, NH_INV_HID}}; + puf_mm_tn(&dh1f, &invf, &a->mv1_wgrad, stream); + Prec dinvf = {.data = a->inv_grad.data, .shape = {B * NH_INV, NH_INV_HID}}; + puf_mm_nn(&dh1f, &ew->mv1_w, &dinvf, stream); } // FIRST writer of inv_grad +#ifdef NH_NOPASS + nh_passbwd_kernel<<>>( + a->inv_grad.data, grad_concat.data, a->inv_sfeat.data, a->inv_idx.data, + 7, NH_MINV_OFF + 2 * NH_MV, B); +#else + nh_passbwd_kernel<<>>( + a->inv_grad.data, grad_concat.data, a->inv_sfeat.data, a->inv_idx.data, + 0, NH_MINV_OFF + 2 * NH_MV, B); +#endif // pointer-decoder key grads: second consumer of inv_out, summed before // the relu mask (both paths read the post-relu slot vectors) if (nh_ptr_keygrad != NULL) nh_add_inplace_kernel<<>>( a->inv_grad.data, nh_ptr_keygrad->data, B * NH_INV_FLAT); + // attention tail: third consumer of inv_out, also pre-relu-mask + // no encoder attention under MIN nh_relu_bias_bwd_kernel<<>>( - a->inv_grad.data, a->inv_out.data, bacc + H + NH_LOC_HID + NH_GLB_HID + NH_BL_HID + NH_P1, + a->inv_grad.data, a->inv_out.data, bacc + H + NH_LOC_HID + NH_GLB_HID + NH_BL_HID + NH_BACC_GLB1, (int64_t)B * NH_INV_FLAT, NH_INV_HID); // state-path weight grad: dW_s = dslot^T @ sfeat over the B*55 slot rows Prec dsflat = {.data = a->inv_grad.data, .shape = {B * NH_INV, NH_INV_HID}}; Prec sfflat = {.data = a->inv_sfeat.data, .shape = {B * NH_INV, NH_SFEAT}}; puf_mm_tn(&dsflat, &sfflat, &a->inv1s_wgrad, stream); - cudaMemsetAsync(a->dTinv_i.data, 0, (size_t)NH_GLYPH_VOCAB * NH_INV_HID * sizeof(long long), stream); + cudaMemsetAsync(a->dTinv_i.data, 0, (size_t)NH_ITBL * NH_INV_HID * sizeof(long long), stream); nh_dTinv_scatter_kernel<<>>( (long long*)a->dTinv_i.data, a->inv_grad.data, a->inv_idx.data, (int64_t)B * NH_INV_FLAT, 0); nh_fxp_to_precision_kernel<<>>( - a->dTinv.data, (long long*)a->dTinv_i.data, NH_GLYPH_VOCAB * NH_INV_HID); + a->dTinv.data, (long long*)a->dTinv_i.data, NH_ITBL * NH_INV_HID); puf_mm_tn(&a->dTinv, &a->e_eff, &a->inv1_wgrad, stream); // discovered-type channel: same fused-table backward keyed by true glyph - cudaMemsetAsync(a->dTtrue_i.data, 0, (size_t)NH_GLYPH_VOCAB * NH_INV_HID * sizeof(long long), stream); + cudaMemsetAsync(a->dTtrue_i.data, 0, (size_t)NH_ITBL * NH_INV_HID * sizeof(long long), stream); nh_dTinv_scatter_kernel<<>>( (long long*)a->dTtrue_i.data, a->inv_grad.data, a->invt_idx.data, (int64_t)B * NH_INV_FLAT, 1); nh_fxp_to_precision_kernel<<>>( - a->dTtrue.data, (long long*)a->dTtrue_i.data, NH_GLYPH_VOCAB * NH_INV_HID); + a->dTtrue.data, (long long*)a->dTtrue_i.data, NH_ITBL * NH_INV_HID); puf_mm_tn(&a->dTtrue, &a->e_eff, &a->invt_wgrad, stream); // Blstats branch (raw-feature slice of concat has no upstream params) nh_slice_kernel<<>>( - a->bl_grad.data, grad_concat.data, B, NH_CONCAT, NH_LOC_HID + NH_GLB_HID + NH_INV_POOL, NH_BL_HID); + a->bl_grad.data, grad_concat.data, B, NH_CONCAT, NH_LOC_HID + NH_GLB_HID + NH_INVP_DIM, NH_BL_HID); nh_relu_bias_bwd_kernel<<>>( a->bl_grad.data, a->bl_out.data, bacc + H + NH_LOC_HID + NH_GLB_HID, (int64_t)B * NH_BL_HID, NH_BL_HID); Prec blg = {.data = a->bl_grad.data, .shape = {B, NH_BL_HID}}; @@ -1483,43 +2233,16 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda // Global branch to the embed table + glb1: scatter dt16 occurrences into // dT, then dE = dT @ W' and dW' = dT^T @ E (the fused-table backward). - int* counts_g = a->sort_grid.data; - int* hot_map_g = counts_g + NH_GLYPH_VOCAB; - int* hot_list_g = hot_map_g + NH_GLYPH_VOCAB; - int* hot_n_g = hot_list_g + NH_HOT_G; - cudaMemsetAsync(counts_g, 0, NH_GLYPH_VOCAB * sizeof(int), stream); - cudaMemsetAsync(hot_map_g, 0xFF, NH_GLYPH_VOCAB * sizeof(int), stream); - cudaMemsetAsync(hot_n_g, 0, sizeof(int), stream); - nh_hist_kernel<<>>(counts_g, a->glyph_idx.data, B * NH_MGRID); - nh_count_pad_kernel<<<1, 1, 0, stream>>>(counts_g, B); - nh_hot_select_kernel<<<1, 1024, 0, stream>>>(hot_map_g, hot_list_g, hot_n_g, counts_g, NH_HOT_G); - nh_dT_patch_scatter_kernel<<<128, 1024, NH_HOT_G * NH_TROW * sizeof(long long), stream>>>( - (long long*)a->dT_i.data, a->t16.data, a->glyph_idx.data, - hot_map_g, hot_list_g, hot_n_g, B); - int dT_n = NH_GLYPH_VOCAB * NH_TROW; - nh_fxp_to_precision_rows_kernel<<>>( - a->dT.data, (long long*)a->dT_i.data, counts_g, hot_map_g, NH_TROW, dT_n); - puf_mm_nn(&a->dT, &a->w_perm, &a->embed_wgrad, stream); // dE = dT @ W' - puf_mm_tn(&a->dT, &a->e_eff, &a->dw_perm, stream); // dW' = dT^T @ E_eff - nh_unpermute_g1_kernel<<>>( - a->glb1_wgrad.data, a->dw_perm.data); - - // Local branch adds its embed grads on top of the GEMM-produced dE. - int* counts_l = a->sort_local.data; - int* hot_map_l = counts_l + NH_GLYPH_VOCAB; - int* hot_list_l = hot_map_l + NH_GLYPH_VOCAB; - int* hot_n_l = hot_list_l + NH_HOT_T; - cudaMemsetAsync(counts_l, 0, NH_GLYPH_VOCAB * sizeof(int), stream); - cudaMemsetAsync(hot_map_l, 0xFF, NH_GLYPH_VOCAB * sizeof(int), stream); - cudaMemsetAsync(hot_n_l, 0, sizeof(int), stream); - nh_hist_kernel<<>>(counts_l, a->crop_glyph.data, B * NH_CGRID); - nh_hot_select_kernel<<<1, 1024, 0, stream>>>(hot_map_l, hot_list_l, hot_n_l, counts_l, NH_HOT_T); - nh_dE_scatter_kernel<<<1024, BLOCK_SIZE, NH_HOT_T * NH_EMBED_DIM * sizeof(long long), stream>>>( - (long long*)a->dE_i.data, dx_local.data, a->crop_glyph.data, - hot_map_l, hot_list_l, hot_n_l, (int64_t)B * NH_CGRID); + // patch machinery gone: seed the embed grad at zero, downstream adds only + cudaMemsetAsync(a->embed_wgrad.data, 0, (size_t)NH_GLYPH_VOCAB * NH_EMBED_DIM * sizeof(precision_t), stream); + int dE_n = NH_GLYPH_VOCAB * NH_EMBED_DIM; - nh_fxp_add_rows_kernel<<>>( - a->embed_wgrad.data, (long long*)a->dE_i.data, counts_l, hot_map_l, NH_EMBED_DIM, dE_n); + // class-crop backward: dx_local scatters into the 9x8 class table only + cudaMemsetAsync(a->locc_acc.data, 0, NH_LOCC_CLASSES * NH_LOCC_DIM * sizeof(long long), stream); + nh_loc3_scatter_kernel<<>>( + (long long*)a->locc_acc.data, dx_local.data, a->crop_glyph.data, nh_locc_lut_dev, B); + nh_fxp_to_precision_kernel<<>>( + a->locc_wgrad.data, (long long*)a->locc_acc.data, NH_LOCC_CLASSES * NH_LOCC_DIM); // Inventory branch adds its embed grads last: dE += dT_inv @ inv1_w. puf_mm_nn(&a->dTinv, &ew->inv1_w, &a->dE_tmp, stream); @@ -1532,13 +2255,36 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda // spell-embed channel: scatter its concat-grad slice into dE (reuse dE_i) cudaMemsetAsync(a->dE_i.data, 0, (size_t)NH_GLYPH_VOCAB * NH_EMBED_DIM * sizeof(long long), stream); - nh_spkey_dk_kernel<<>>( - a->spk_dkeys.data, grad_concat.data, a->spk_pool.data, a->spk_amax.data, - ew->spk2_w.data, a->spk_keys.data, - nh_ptr_spkeygrad != NULL ? nh_ptr_spkeygrad->data : NULL, B); - nh_spk2_grad_kernel<<>>( - a->spk2_wgrad.data, a->spk2_bgrad.data, grad_concat.data, a->spk_pool.data, - a->spk_amax.data, a->spk_keys.data, B); + // spell block backward: dv (sum+max routing) -> MLP chain -> trunk dkeys + // staged into spk_dkeys, then pointer grads added + key relu gate. + nh_min_dv_kernel<<>>( + a->spdv.data, grad_concat.data, a->spvmax.data, a->spell_idx.data, + 0.25f, NH_SPELL_SLOTS, -1, NH_SPELL_CONCAT_OFF, B); + cudaMemsetAsync(a->spm1b_acc.data, 0, NH_SPM * sizeof(long long), stream); + cudaMemsetAsync(a->spm2b_acc.data, 0, NH_SPM * sizeof(long long), stream); + NH_ENT_LAST_BWD<<>>( + a->spdv.data, a->spv.data, (long long*)a->spm2b_acc.data, + (int64_t)B * NH_SPELL_SLOTS * NH_SPM, NH_SPM); + nh_fxp_to_precision_kernel<<>>( + a->spm2_bgrad.data, (long long*)a->spm2b_acc.data, NH_SPM); + { Prec dvf = {.data = a->spdv.data, .shape = {B * NH_SPELL_SLOTS, NH_SPM}}; + Prec hf = {.data = a->sph1.data, .shape = {B * NH_SPELL_SLOTS, NH_SPM}}; + puf_mm_tn(&dvf, &hf, &a->spm2_wgrad, stream); + Prec dhf = {.data = a->spdh1.data, .shape = {B * NH_SPELL_SLOTS, NH_SPM}}; + puf_mm_nn(&dvf, &ew->spm2_w, &dhf, stream); } + nh_relu_bias_bwd_kernel<<>>( + a->spdh1.data, a->sph1.data, (long long*)a->spm1b_acc.data, + (int64_t)B * NH_SPELL_SLOTS * NH_SPM, NH_SPM); + nh_fxp_to_precision_kernel<<>>( + a->spm1_bgrad.data, (long long*)a->spm1b_acc.data, NH_SPM); + { Prec dhf = {.data = a->spdh1.data, .shape = {B * NH_SPELL_SLOTS, NH_SPM}}; + Prec kf = {.data = a->spk_keys.data, .shape = {B * NH_SPELL_SLOTS, NH_SPKEY}}; + puf_mm_tn(&dhf, &kf, &a->spm1_wgrad, stream); + Prec dkf = {.data = a->spk_dkeys.data, .shape = {B * NH_SPELL_SLOTS, NH_SPKEY}}; + puf_mm_nn(&dhf, &ew->spm1_w, &dkf, stream); } // trunk dkeys, FIRST writer + nh_sp2_dk_kernel<<>>( + a->spk_dkeys.data, nh_ptr_spkeygrad != NULL ? nh_ptr_spkeygrad->data : NULL, + a->spk_keys.data, B); nh_spkey_dE_kernel<<>>( (long long*)a->dE_i.data, a->spk_dkeys.data, ew->spk_w.data, a->spell_idx.data, B); { Prec dkf = {.data = a->spk_dkeys.data, .shape = {B * NH_SPELL_SLOTS, NH_SPKEY}}; @@ -1548,6 +2294,100 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda a->dE_tmp.data, (long long*)a->dE_i.data, NH_GLYPH_VOCAB * NH_EMBED_DIM); nh_add_inplace_kernel<<>>( a->embed_wgrad.data, a->dE_tmp.data, dE_n); + { // typed stream backward: pool -> score/value grads -> shared-embed scatter + // monster stream + cudaMemsetAsync(a->mrb_acc.data, 0, NH_INV_HID * sizeof(long long), stream); + cudaMemsetAsync(a->mm1b_acc.data, 0, NH_MV * sizeof(long long), stream); + cudaMemsetAsync(a->mm2b_acc.data, 0, NH_MV * sizeof(long long), stream); + nh_min_dv_kernel<<>>( + a->mdvm.data, grad_concat.data, a->mvmaxm.data, a->lm_gid.data, + 0.25f, NH_LABK, -1, NH_MLM_OFF, B); + NH_ENT_LAST_BWD<<>>( + a->mdvm.data, a->mvm.data, (long long*)a->mm2b_acc.data, + (int64_t)B * NH_LABK * NH_MV, NH_MV); + nh_fxp_to_precision_kernel<<>>( + a->mm2_bgrad.data, (long long*)a->mm2b_acc.data, NH_MV); + { Prec dvf = {.data = a->mdvm.data, .shape = {B * NH_LABK, NH_MV}}; + Prec h1f = {.data = a->mh1m.data, .shape = {B * NH_LABK, NH_MV}}; + puf_mm_tn(&dvf, &h1f, &a->mm2_wgrad, stream); + Prec dh1f = {.data = a->mdh1m.data, .shape = {B * NH_LABK, NH_MV}}; + puf_mm_nn(&dvf, &ew->mm2_w, &dh1f, stream); } + nh_relu_bias_bwd_kernel<<>>( + a->mdh1m.data, a->mh1m.data, (long long*)a->mm1b_acc.data, + (int64_t)B * NH_LABK * NH_MV, NH_MV); + nh_fxp_to_precision_kernel<<>>( + a->mm1_bgrad.data, (long long*)a->mm1b_acc.data, NH_MV); + { Prec dh1f = {.data = a->mdh1m.data, .shape = {B * NH_LABK, NH_MV}}; + Prec rpf = {.data = a->mrp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + puf_mm_tn(&dh1f, &rpf, &a->mm1_wgrad, stream); + Prec drpf = {.data = a->mdrp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + puf_mm_nn(&dh1f, &ew->mm1_w, &drpf, stream); } // FIRST writer of mdrp + nh_min_sgate_bwd_kernel<<>>( + a->mdrp.data, grad_concat.data, a->lm_tok.data, a->lm_gid.data, + 0, NH_MLM_OFF + 2 * NH_MV, B); + nh_relu_bias_bwd_kernel<<>>( + a->mdrp.data, a->mrp.data, (long long*)a->mrb_acc.data, + (int64_t)B * NH_LABK * NH_INV_HID, NH_INV_HID); + nh_fxp_to_precision_kernel<<>>( + a->mr_bgrad.data, (long long*)a->mrb_acc.data, NH_INV_HID); + Prec mtokf = {.data = a->lm_tok.data, .shape = {B * NH_LABK, NH_LAB_IN}}; + { Prec drpf = {.data = a->mdrp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + puf_mm_tn(&drpf, &mtokf, &a->mr_wgrad, stream); + Prec dtokf = {.data = a->lm_tok.data, .shape = {B * NH_LABK, NH_LAB_IN}}; // in place; tok done + puf_mm_nn(&drpf, &ew->mr_w, &dtokf, stream); } + // item stream + cudaMemsetAsync(a->irb_acc.data, 0, NH_INV_HID * sizeof(long long), stream); + cudaMemsetAsync(a->im1b_acc.data, 0, NH_MV * sizeof(long long), stream); + cudaMemsetAsync(a->im2b_acc.data, 0, NH_MV * sizeof(long long), stream); + nh_min_dv_kernel<<>>( + a->mdvi.data, grad_concat.data, a->mvmaxi.data, a->li_gid.data, + 0.25f, NH_LABK, -1, NH_MLI_OFF, B); + NH_ENT_LAST_BWD<<>>( + a->mdvi.data, a->mvi.data, (long long*)a->im2b_acc.data, + (int64_t)B * NH_LABK * NH_MV, NH_MV); + nh_fxp_to_precision_kernel<<>>( + a->im2_bgrad.data, (long long*)a->im2b_acc.data, NH_MV); + { Prec dvf = {.data = a->mdvi.data, .shape = {B * NH_LABK, NH_MV}}; + Prec h1f = {.data = a->mh1i.data, .shape = {B * NH_LABK, NH_MV}}; + puf_mm_tn(&dvf, &h1f, &a->im2_wgrad, stream); + Prec dh1f = {.data = a->mdh1i.data, .shape = {B * NH_LABK, NH_MV}}; + puf_mm_nn(&dvf, &ew->im2_w, &dh1f, stream); } + nh_relu_bias_bwd_kernel<<>>( + a->mdh1i.data, a->mh1i.data, (long long*)a->im1b_acc.data, + (int64_t)B * NH_LABK * NH_MV, NH_MV); + nh_fxp_to_precision_kernel<<>>( + a->im1_bgrad.data, (long long*)a->im1b_acc.data, NH_MV); + { Prec dh1f = {.data = a->mdh1i.data, .shape = {B * NH_LABK, NH_MV}}; + Prec rpf = {.data = a->irp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + puf_mm_tn(&dh1f, &rpf, &a->im1_wgrad, stream); + Prec drpf = {.data = a->idrp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + puf_mm_nn(&dh1f, &ew->im1_w, &drpf, stream); } // FIRST writer of idrp + nh_min_sgate_bwd_kernel<<>>( + a->idrp.data, grad_concat.data, a->li_tok.data, a->li_gid.data, + 1, NH_MLI_OFF + 2 * NH_MV, B); + nh_relu_bias_bwd_kernel<<>>( + a->idrp.data, a->irp.data, (long long*)a->irb_acc.data, + (int64_t)B * NH_LABK * NH_INV_HID, NH_INV_HID); + nh_fxp_to_precision_kernel<<>>( + a->ir_bgrad.data, (long long*)a->irb_acc.data, NH_INV_HID); + Prec itokf = {.data = a->li_tok.data, .shape = {B * NH_LABK, NH_LAB_IN}}; + { Prec drpf = {.data = a->idrp.data, .shape = {B * NH_LABK, NH_INV_HID}}; + puf_mm_tn(&drpf, &itokf, &a->ir_wgrad, stream); + Prec dtokf = {.data = a->li_tok.data, .shape = {B * NH_LABK, NH_LAB_IN}}; // in place; tok done + puf_mm_nn(&drpf, &ew->ir_w, &dtokf, stream); } + // shared-embed scatter: value-path dtok (in tok buffers) + score-path + // dtok (dts buffers) for both streams, then one flush + add + cudaMemsetAsync(a->dE_i.data, 0, (size_t)NH_GLYPH_VOCAB * NH_EMBED_DIM * sizeof(long long), stream); + nh_lab_dE_scatter_kernel<<>>( + (long long*)a->dE_i.data, a->lm_tok.data, a->lm_gid.data, B); + nh_lab_dE_scatter_kernel<<>>( + (long long*)a->dE_i.data, a->li_tok.data, a->li_gid.data, B); + nh_fxp_to_precision_kernel<<>>( + a->dE_tmp.data, (long long*)a->dE_i.data, NH_GLYPH_VOCAB * NH_EMBED_DIM); + nh_add_inplace_kernel<<>>( + a->embed_wgrad.data, a->dE_tmp.data, dE_n); + } + // dE_eff complete: embed_wgrad doubles as dE_res (identity path); factor // tables get deterministic CSR row sums of the same buffer @@ -1556,11 +2396,15 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda nh_esub_grad_kernel<<>>( a->esub_wgrad.data, a->embed_wgrad.data); - nh_bias_flush_kernel<<>>( - bacc, a->proj_bgrad.data, H, a->loc_bgrad.data, NH_LOC_HID, - a->glb2_bgrad.data, NH_GLB_HID, a->bl_bgrad.data, NH_BL_HID, - a->glb1_bgrad.data, NH_P1, a->inv1_bgrad.data, NH_INV_HID, - a->inv2_bgrad.data, NH_INV_POOL); + nh_bias_flush_kernel<<>>( + bacc, a->proj_bgrad.data, H, + a->loc2_bgrad.data, + NH_LOC_HID, + a->terr2_bgrad.data, + NH_GLB_HID, a->bl_bgrad.data, NH_BL_HID, + (precision_t*)NULL, 0, + a->inv1_bgrad.data, NH_INV_HID, + (precision_t*)NULL, 0); } static void nethack_encoder_init_weights(void* w, uint64_t* seed, cudaStream_t stream) { @@ -1572,34 +2416,68 @@ static void nethack_encoder_init_weights(void* w, uint64_t* seed, cudaStream_t s cudaMemsetAsync(ew->esub_w.data, 0, numel(ew->esub_w.shape) * sizeof(precision_t), stream); puf_kaiming_init(&ew->loc_w, 1.0f, (*seed)++, stream); cudaMemsetAsync(ew->loc_b.data, 0, numel(ew->loc_b.shape) * sizeof(precision_t), stream); - puf_kaiming_init(&ew->glb1_w, 1.0f, (*seed)++, stream); - // zero: kaiming's fan_in=2 would run 20x hotter than the glyph slice - cudaMemsetAsync(ew->glb1_xy.data, 0, numel(ew->glb1_xy.shape) * sizeof(precision_t), stream); - cudaMemsetAsync(ew->glb1_b.data, 0, numel(ew->glb1_b.shape) * sizeof(precision_t), stream); - puf_kaiming_init(&ew->glb2_w, 1.0f, (*seed)++, stream); - cudaMemsetAsync(ew->glb2_b.data, 0, numel(ew->glb2_b.shape) * sizeof(precision_t), stream); + puf_kaiming_init(&ew->loc2_w, 1.0f, (*seed)++, stream); + cudaMemsetAsync(ew->loc2_b.data, 0, numel(ew->loc2_b.shape) * sizeof(precision_t), stream); + puf_kaiming_init(&ew->terr1_w, 1.0f, (*seed)++, stream); + cudaMemsetAsync(ew->terr1_b.data, 0, numel(ew->terr1_b.shape) * sizeof(precision_t), stream); + puf_kaiming_init(&ew->terr2_w, 1.0f, (*seed)++, stream); + cudaMemsetAsync(ew->terr2_b.data, 0, numel(ew->terr2_b.shape) * sizeof(precision_t), stream); + puf_normal_init(&ew->locc_w, 1.0f, (*seed)++, stream); // class table = embedding idiom puf_kaiming_init(&ew->inv1_w, 1.0f, (*seed)++, stream); cudaMemsetAsync(ew->inv1_b.data, 0, numel(ew->inv1_b.shape) * sizeof(precision_t), stream); puf_kaiming_init(&ew->inv1s_w, 1.0f, (*seed)++, stream); // zero: discovered-type channel starts as an exact no-op (ekind_w idiom) cudaMemsetAsync(ew->invt_w.data, 0, numel(ew->invt_w.shape) * sizeof(precision_t), stream); - puf_kaiming_init(&ew->inv2_w, 1.0f, (*seed)++, stream); - cudaMemsetAsync(ew->inv2_b.data, 0, numel(ew->inv2_b.shape) * sizeof(precision_t), stream); puf_kaiming_init(&ew->bl_w, 1.0f, (*seed)++, stream); cudaMemsetAsync(ew->bl_b.data, 0, numel(ew->bl_b.shape) * sizeof(precision_t), stream); puf_kaiming_init(&ew->proj_w, 1.0f, (*seed)++, stream); cudaMemsetAsync(ew->proj_b.data, 0, numel(ew->proj_b.shape) * sizeof(precision_t), stream); puf_normal_init(&ew->msg_w, 1.0f, (*seed)++, stream); // trigram embedding + { // NH_MSG_INIT=: pretrained trigram table (float32, NH_MSG_VOCAB x NH_MSG_HID row-major) + const char* mp = getenv("NH_MSG_INIT"); + size_t mn = (size_t)NH_MSG_VOCAB * NH_MSG_HID; + if (mp && *mp) { + FILE* mf = fopen(mp, "rb"); + float* mh = (float*)malloc(mn * sizeof(float)); + if (!mf || fread(mh, sizeof(float), mn, mf) != mn) { fprintf(stderr, "NH_MSG_INIT: cannot read %s\n", mp); exit(1); } + fclose(mf); + float* md; cudaMalloc(&md, mn * sizeof(float)); + cudaMemcpyAsync(md, mh, mn * sizeof(float), cudaMemcpyHostToDevice, stream); + nh_f32_to_precision_kernel<<>>(ew->msg_w.data, md, (int)mn); + cudaStreamSynchronize(stream); cudaFree(md); free(mh); + fprintf(stderr, "NH_MSG_INIT: loaded %s (%zu values)\n", mp, mn); + } +#ifdef NH_MSG_FROZEN + if (!nh_msg_w0_dev) cudaMalloc(&nh_msg_w0_dev, mn * sizeof(precision_t)); + cudaMemcpyAsync(nh_msg_w0_dev, ew->msg_w.data, mn * sizeof(precision_t), cudaMemcpyDeviceToDevice, stream); +#endif + } puf_kaiming_init(&ew->spk_w, 1.0f, (*seed)++, stream); // spell slot-rep projection - puf_kaiming_init(&ew->spk2_w, 1.0f, (*seed)++, stream); - cudaMemsetAsync(ew->spk2_b.data, 0, numel(ew->spk2_b.shape) * sizeof(precision_t), stream); -#if NH_ID_EMBED + puf_kaiming_init(&ew->spm1_w, 1.0f, (*seed)++, stream); + cudaMemsetAsync(ew->spm1_b.data, 0, numel(ew->spm1_b.shape) * sizeof(precision_t), stream); + puf_kaiming_init(&ew->spm2_w, 1.0f, (*seed)++, stream); + cudaMemsetAsync(ew->spm2_b.data, 0, numel(ew->spm2_b.shape) * sizeof(precision_t), stream); // zero: the identity channel starts as an exact no-op (ekind_w idiom) cudaMemsetAsync(ew->ide_role_w.data, 0, numel(ew->ide_role_w.shape) * sizeof(precision_t), stream); cudaMemsetAsync(ew->ide_race_w.data, 0, numel(ew->ide_race_w.shape) * sizeof(precision_t), stream); cudaMemsetAsync(ew->ide_gend_w.data, 0, numel(ew->ide_gend_w.shape) * sizeof(precision_t), stream); cudaMemsetAsync(ew->ide_algn_w.data, 0, numel(ew->ide_algn_w.shape) * sizeof(precision_t), stream); -#endif + puf_kaiming_init(&ew->mv1_w, 1.0f, (*seed)++, stream); + puf_kaiming_init(&ew->mv2_w, 1.0f, (*seed)++, stream); + puf_kaiming_init(&ew->mr_w, 1.0f, (*seed)++, stream); + puf_kaiming_init(&ew->mm1_w, 1.0f, (*seed)++, stream); + puf_kaiming_init(&ew->mm2_w, 1.0f, (*seed)++, stream); + puf_kaiming_init(&ew->ir_w, 1.0f, (*seed)++, stream); + puf_kaiming_init(&ew->im1_w, 1.0f, (*seed)++, stream); + puf_kaiming_init(&ew->im2_w, 1.0f, (*seed)++, stream); + cudaMemsetAsync(ew->mv1_b.data, 0, numel(ew->mv1_b.shape) * sizeof(precision_t), stream); + cudaMemsetAsync(ew->mv2_b.data, 0, numel(ew->mv2_b.shape) * sizeof(precision_t), stream); + cudaMemsetAsync(ew->mr_b.data, 0, numel(ew->mr_b.shape) * sizeof(precision_t), stream); + cudaMemsetAsync(ew->mm1_b.data, 0, numel(ew->mm1_b.shape) * sizeof(precision_t), stream); + cudaMemsetAsync(ew->mm2_b.data, 0, numel(ew->mm2_b.shape) * sizeof(precision_t), stream); + cudaMemsetAsync(ew->ir_b.data, 0, numel(ew->ir_b.shape) * sizeof(precision_t), stream); + cudaMemsetAsync(ew->im1_b.data, 0, numel(ew->im1_b.shape) * sizeof(precision_t), stream); + cudaMemsetAsync(ew->im2_b.data, 0, numel(ew->im2_b.shape) * sizeof(precision_t), stream); } // Param and grad registration orders must match pairwise (muon walks both flat). @@ -1608,51 +2486,69 @@ static void nethack_encoder_reg_params(void* w, Allocator* alloc) { ew->embed_w = {.shape = {NH_GLYPH_VOCAB, NH_EMBED_DIM}}; ew->ekind_w = {.shape = {NH_NKIND, NH_EMBED_DIM}}; // 14x32=448, mult of 8 ew->esub_w = {.shape = {NH_NSUB, NH_EMBED_DIM}}; // 944x32=30208, mult of 8 - ew->loc_w = {.shape = {NH_LOC_HID, NH_LOC_IN}}; - ew->loc_b = {.shape = {NH_LOC_HID}}; - ew->glb1_w = {.shape = {NH_P1, NH_PCELLS * NH_EMBED_DIM}}; - ew->glb1_xy = {.shape = {NH_P1, 2}}; - ew->glb1_b = {.shape = {NH_P1}}; - ew->glb2_w = {.shape = {NH_GLB_HID, NH_P1}}; - ew->glb2_b = {.shape = {NH_GLB_HID}}; + ew->loc_w = {.shape = {NH_LOC_H1, NH_LOC_IN}}; + ew->loc_b = {.shape = {NH_LOC_H1}}; + ew->loc2_w = {.shape = {NH_LOC_HID, NH_LOC_H1}}; + ew->loc2_b = {.shape = {NH_LOC_HID}}; + ew->terr1_w = {.shape = {NH_TERR_H1, NH_TERRF}}; // 256x592=151552, mult of 8 + ew->terr1_b = {.shape = {NH_TERR_H1}}; + ew->terr2_w = {.shape = {NH_GLB_HID, NH_TERR_H1}}; // 128x256, mult of 8 + ew->terr2_b = {.shape = {NH_GLB_HID}}; + ew->locc_w = {.shape = {NH_LOCC_CLASSES, NH_LOCC_DIM}}; // 9x8=72, mult of 8 ew->inv1_w = {.shape = {NH_INV_HID, NH_EMBED_DIM}}; ew->inv1_b = {.shape = {NH_INV_HID}}; ew->inv1s_w = {.shape = {NH_INV_HID, NH_SFEAT}}; ew->invt_w = {.shape = {NH_INV_HID, NH_EMBED_DIM}}; // 16x32=512, mult of 8 - ew->inv2_w = {.shape = {NH_INV_POOL, NH_INV_HID}}; - ew->inv2_b = {.shape = {NH_INV_POOL}}; + // 64x16=1024, mult of 8; ew->bl_w = {.shape = {NH_BL_HID, NH_BL_FEAT}}; ew->bl_b = {.shape = {NH_BL_HID}}; ew->proj_w = {.shape = {ew->hidden, NH_CONCAT}}; ew->proj_b = {.shape = {ew->hidden}}; ew->msg_w = {.shape = {NH_MSG_VOCAB, NH_MSG_HID}}; // 4096x32=131072, mult of 8 ew->spk_w = {.shape = {NH_SPKEY, NH_SPIN}}; // 16x36=576, mult of 8 - ew->spk2_w = {.shape = {NH_SPKEY, NH_SPKEY}}; // 16x16=256, mult of 8 - ew->spk2_b = {.shape = {NH_SPKEY}}; // 16, mult of 8 -#if NH_ID_EMBED + ew->spm1_w = {.shape = {NH_SPM, NH_SPKEY}}; // 64x16, mult of 8 + ew->spm1_b = {.shape = {NH_SPM}}; + ew->spm2_w = {.shape = {NH_SPM, NH_SPM}}; + ew->spm2_b = {.shape = {NH_SPM}}; ew->ide_role_w = {.shape = {13, NH_IDE_ROLE}}; // 208, mult of 8 ew->ide_race_w = {.shape = {5, NH_IDE_RACE}}; // 40 ew->ide_gend_w = {.shape = {2, NH_IDE_GEND}}; // 16 ew->ide_algn_w = {.shape = {3, NH_IDE_ALGN}}; // 24 -#endif + // 3072, mult of 8; + // 384; alloc_register(alloc,&ew->embed_w); alloc_register(alloc,&ew->ekind_w); alloc_register(alloc,&ew->esub_w); alloc_register(alloc,&ew->loc_w); alloc_register(alloc,&ew->loc_b); - alloc_register(alloc,&ew->glb1_w); alloc_register(alloc,&ew->glb1_xy); - alloc_register(alloc,&ew->glb1_b); - alloc_register(alloc,&ew->glb2_w); alloc_register(alloc,&ew->glb2_b); + alloc_register(alloc,&ew->loc2_w); alloc_register(alloc,&ew->loc2_b); + alloc_register(alloc,&ew->terr1_w); alloc_register(alloc,&ew->terr1_b); + alloc_register(alloc,&ew->terr2_w); alloc_register(alloc,&ew->terr2_b); + alloc_register(alloc,&ew->locc_w); alloc_register(alloc,&ew->inv1_w); alloc_register(alloc,&ew->inv1_b); alloc_register(alloc,&ew->inv1s_w); alloc_register(alloc,&ew->invt_w); - alloc_register(alloc,&ew->inv2_w); alloc_register(alloc,&ew->inv2_b); alloc_register(alloc,&ew->bl_w); alloc_register(alloc,&ew->bl_b); alloc_register(alloc,&ew->proj_w); alloc_register(alloc,&ew->proj_b); alloc_register(alloc,&ew->msg_w); alloc_register(alloc,&ew->spk_w); - alloc_register(alloc,&ew->spk2_w); alloc_register(alloc,&ew->spk2_b); -#if NH_ID_EMBED + alloc_register(alloc,&ew->spm1_w); alloc_register(alloc,&ew->spm1_b); + alloc_register(alloc,&ew->spm2_w); alloc_register(alloc,&ew->spm2_b); alloc_register(alloc,&ew->ide_role_w); alloc_register(alloc,&ew->ide_race_w); alloc_register(alloc,&ew->ide_gend_w); alloc_register(alloc,&ew->ide_algn_w); -#endif + ew->mv1_w = {.shape = {NH_MV, NH_INV_HID}}; ew->mv1_b = {.shape = {NH_MV}}; + ew->mv2_w = {.shape = {NH_MV, NH_MV}}; ew->mv2_b = {.shape = {NH_MV}}; + ew->mr_w = {.shape = {NH_INV_HID, NH_LAB_IN}}; ew->mr_b = {.shape = {NH_INV_HID}}; + ew->mm1_w = {.shape = {NH_MV, NH_INV_HID}}; ew->mm1_b = {.shape = {NH_MV}}; + ew->mm2_w = {.shape = {NH_MV, NH_MV}}; ew->mm2_b = {.shape = {NH_MV}}; + ew->ir_w = {.shape = {NH_INV_HID, NH_LAB_IN}}; ew->ir_b = {.shape = {NH_INV_HID}}; + ew->im1_w = {.shape = {NH_MV, NH_INV_HID}}; ew->im1_b = {.shape = {NH_MV}}; + ew->im2_w = {.shape = {NH_MV, NH_MV}}; ew->im2_b = {.shape = {NH_MV}}; + alloc_register(alloc,&ew->mv1_w); alloc_register(alloc,&ew->mv1_b); + alloc_register(alloc,&ew->mv2_w); alloc_register(alloc,&ew->mv2_b); + alloc_register(alloc,&ew->mr_w); alloc_register(alloc,&ew->mr_b); + alloc_register(alloc,&ew->mm1_w); alloc_register(alloc,&ew->mm1_b); + alloc_register(alloc,&ew->mm2_w); alloc_register(alloc,&ew->mm2_b); + alloc_register(alloc,&ew->ir_w); alloc_register(alloc,&ew->ir_b); + alloc_register(alloc,&ew->im1_w); alloc_register(alloc,&ew->im1_b); + alloc_register(alloc,&ew->im2_w); alloc_register(alloc,&ew->im2_b); } static void nethack_encoder_reg_train(void* w, void* activations, Allocator* acts, Allocator* grads, int B_TT) { @@ -1663,21 +2559,21 @@ static void nethack_encoder_reg_train(void* w, void* activations, Allocator* act a->crop_glyph = {.shape = {B_TT, NH_CGRID}}; a->e_eff = {.shape = {NH_GLYPH_VOCAB, NH_EMBED_DIM}}; a->x_local = {.shape = {B_TT, NH_LOC_IN}}; - a->w_perm = {.shape = {NH_TROW, NH_EMBED_DIM}}; - a->glyph_T = {.shape = {NH_GLYPH_VOCAB, NH_TROW}}; - a->t16 = {.shape = {B_TT, NH_TOK * NH_P1}}; - a->dxy = {.shape = {B_TT, NH_TOK * 2}}; - a->tok_argmax = {.shape = {B_TT, NH_GLB_HID}}; + a->terr_tf = {.shape = {B_TT, NH_TERRF}}; + a->terr_h = {.shape = {B_TT, NH_TERR_H1}}; + a->terr_dh = {.shape = {B_TT, NH_TERR_H1}}; + a->terr1b_acc = {.shape = {NH_TERR_H1}}; + a->locc_acc = {.shape = {NH_LOCC_CLASSES * NH_LOCC_DIM}}; a->inv_idx = {.shape = {B_TT, NH_INV}}; a->spell_idx = {.shape = {B_TT, 8}}; a->invt_idx = {.shape = {B_TT, NH_INV}}; a->inv_sfeat = {.shape = {B_TT, NH_INV * NH_SFEAT}}; - a->inv_T = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; - a->invt_T = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; + a->inv_T = {.shape = {NH_ITBL, NH_INV_HID}}; + a->invt_T = {.shape = {NH_ITBL, NH_INV_HID}}; a->inv_out = {.shape = {B_TT, NH_INV_FLAT}}; - a->inv_pool = {.shape = {B_TT, NH_INV_POOL}}; - a->inv_amax = {.shape = {B_TT, NH_INV_POOL}}; a->loc_out = {.shape = {B_TT, NH_LOC_HID}}; + a->loc_h1 = {.shape = {B_TT, NH_LOC_H1}}; a->loc_h1_grad = {.shape = {B_TT, NH_LOC_H1}}; + a->loc1b_acc = {.shape = {NH_LOC_H1}}; a->glb_out = {.shape = {B_TT, NH_GLB_HID}}; a->bl_feats = {.shape = {B_TT, NH_BL_FEAT}}; a->bl_out = {.shape = {B_TT, NH_BL_HID}}; @@ -1686,115 +2582,155 @@ static void nethack_encoder_reg_train(void* w, void* activations, Allocator* act a->spk_in = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPIN}}; a->spk_keys = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPKEY}}; a->spk_dkeys = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPKEY}}; - a->spk_amax = {.shape = {B_TT, NH_SPKEY}}; - a->spk_pool = {.shape = {B_TT, NH_SPKEY}}; -#if NH_ID_EMBED + a->sph1 = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPM}}; + a->spdh1 = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPM}}; + a->spv = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPM}}; + a->spdv = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPM}}; + a->spvmax = {.shape = {B_TT, NH_SPM}}; + a->spm1b_acc = {.shape = {NH_SPM}}; a->spm2b_acc = {.shape = {NH_SPM}}; a->ide_idx = {.shape = {B_TT, 4}}; -#endif + a->lm_tok = {.shape = {B_TT, NH_LABK * NH_LAB_IN}}; + a->lm_gid = {.shape = {B_TT, NH_LABK}}; + a->li_tok = {.shape = {B_TT, NH_LABK * NH_LAB_IN}}; + a->li_gid = {.shape = {B_TT, NH_LABK}}; a->concat = {.shape = {B_TT, NH_CONCAT}}; a->out = {.shape = {B_TT, ew->hidden}}; alloc_register(acts,&a->glyph_idx); alloc_register(acts,&a->crop_glyph); alloc_register(acts,&a->e_eff); alloc_register(acts,&a->x_local); - alloc_register(acts,&a->w_perm); alloc_register(acts,&a->glyph_T); - alloc_register(acts,&a->t16); alloc_register(acts,&a->dxy); - alloc_register(acts,&a->tok_argmax); + alloc_register(acts,&a->terr_tf); alloc_register(acts,&a->terr_h); + alloc_register(acts,&a->terr_dh); alloc_register(acts,&a->terr1b_acc); + alloc_register(acts,&a->locc_acc); alloc_register(acts,&a->inv_idx); alloc_register(acts,&a->invt_idx); alloc_register(acts,&a->spell_idx); alloc_register(acts,&a->spk_in); alloc_register(acts,&a->spk_keys); - alloc_register(acts,&a->spk_dkeys); alloc_register(acts,&a->spk_amax); - alloc_register(acts,&a->spk_pool); -#if NH_ID_EMBED + alloc_register(acts,&a->spk_dkeys); + alloc_register(acts,&a->sph1); alloc_register(acts,&a->spdh1); + alloc_register(acts,&a->spv); alloc_register(acts,&a->spdv); + alloc_register(acts,&a->spvmax); + alloc_register(acts,&a->spm1b_acc); alloc_register(acts,&a->spm2b_acc); alloc_register(acts,&a->ide_idx); -#endif alloc_register(acts,&a->inv_sfeat); alloc_register(acts,&a->inv_T); alloc_register(acts,&a->invt_T); alloc_register(acts,&a->inv_out); - alloc_register(acts,&a->inv_pool); alloc_register(acts,&a->inv_amax); alloc_register(acts,&a->loc_out); alloc_register(acts,&a->glb_out); alloc_register(acts,&a->bl_feats); alloc_register(acts,&a->bl_out); alloc_register(acts,&a->msg_ids); alloc_register(acts,&a->msg_out); + alloc_register(acts,&a->lm_tok); + alloc_register(acts,&a->lm_gid); + alloc_register(acts,&a->li_tok); + alloc_register(acts,&a->li_gid); alloc_register(acts,&a->concat); alloc_register(acts,&a->out); a->loc_grad = {.shape = {B_TT, NH_LOC_HID}}; a->glb_grad = {.shape = {B_TT, NH_GLB_HID}}; a->inv_grad = {.shape = {B_TT, NH_INV_FLAT}}; - a->inv_pool_grad = {.shape = {B_TT, NH_INV_POOL}}; a->bl_grad = {.shape = {B_TT, NH_BL_HID}}; - a->dT = {.shape = {NH_GLYPH_VOCAB, NH_TROW}}; - a->dT_i = {.shape = {NH_GLYPH_VOCAB, NH_TROW}}; - a->dTinv = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; - a->dTinv_i = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; - a->dTtrue = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; - a->dTtrue_i = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; + a->dTinv = {.shape = {NH_ITBL, NH_INV_HID}}; + a->dTinv_i = {.shape = {NH_ITBL, NH_INV_HID}}; + a->dTtrue = {.shape = {NH_ITBL, NH_INV_HID}}; + a->dTtrue_i = {.shape = {NH_ITBL, NH_INV_HID}}; a->dE_tmp = {.shape = {NH_GLYPH_VOCAB, NH_EMBED_DIM}}; - a->dw_perm = {.shape = {NH_TROW, NH_EMBED_DIM}}; a->dE_i = {.shape = {NH_GLYPH_VOCAB, NH_EMBED_DIM}}; - a->dw2_acc = {.shape = {NH_GLB_HID * NH_P1}}; - a->dw2i_acc = {.shape = {NH_INV_POOL * NH_INV_HID}}; + a->loc2_wgrad = {.shape = {NH_LOC_HID, NH_LOC_H1}}; + a->loc2_bgrad = {.shape = {NH_LOC_HID}}; a->dmsg_acc = {.shape = {NH_MSG_VOCAB * NH_MSG_HID}}; - a->sort_local = {.shape = {2 * NH_GLYPH_VOCAB + NH_HOT_T + 1}}; - a->sort_grid = {.shape = {2 * NH_GLYPH_VOCAB + NH_HOT_G + 1}}; - a->bias_acc = {.shape = {ew->hidden + NH_LOC_HID + NH_GLB_HID + NH_BL_HID + NH_P1 + NH_INV_HID + NH_INV_POOL}}; + a->bias_acc = {.shape = {ew->hidden + NH_LOC_HID + NH_GLB_HID + NH_BL_HID + NH_P1 + NH_INV_HID + NH_INV_POOL}}; // superset of packed slots alloc_register(acts,&a->loc_grad); alloc_register(acts,&a->glb_grad); - alloc_register(acts,&a->inv_grad); alloc_register(acts,&a->inv_pool_grad); + a->mh1 = {.shape = {B_TT, NH_INV * NH_MV}}; a->mdh1 = {.shape = {B_TT, NH_INV * NH_MV}}; + a->mvv = {.shape = {B_TT, NH_INV * NH_MV}}; a->mdv = {.shape = {B_TT, NH_INV * NH_MV}}; + a->mvmax = {.shape = {B_TT, NH_MV}}; + a->mrp = {.shape = {B_TT, NH_LABK * NH_INV_HID}}; a->mdrp = {.shape = {B_TT, NH_LABK * NH_INV_HID}}; + a->mh1m = {.shape = {B_TT, NH_LABK * NH_MV}}; a->mdh1m = {.shape = {B_TT, NH_LABK * NH_MV}}; + a->mvm = {.shape = {B_TT, NH_LABK * NH_MV}}; a->mdvm = {.shape = {B_TT, NH_LABK * NH_MV}}; + a->mvmaxm = {.shape = {B_TT, NH_MV}}; + a->irp = {.shape = {B_TT, NH_LABK * NH_INV_HID}}; a->idrp = {.shape = {B_TT, NH_LABK * NH_INV_HID}}; + a->mh1i = {.shape = {B_TT, NH_LABK * NH_MV}}; a->mdh1i = {.shape = {B_TT, NH_LABK * NH_MV}}; + a->mvi = {.shape = {B_TT, NH_LABK * NH_MV}}; a->mdvi = {.shape = {B_TT, NH_LABK * NH_MV}}; + a->mvmaxi = {.shape = {B_TT, NH_MV}}; + a->mv1b_acc = {.shape = {NH_MV}}; a->mv2b_acc = {.shape = {NH_MV}}; + a->mrb_acc = {.shape = {NH_INV_HID}}; a->mm1b_acc = {.shape = {NH_MV}}; a->mm2b_acc = {.shape = {NH_MV}}; + a->irb_acc = {.shape = {NH_INV_HID}}; a->im1b_acc = {.shape = {NH_MV}}; a->im2b_acc = {.shape = {NH_MV}}; + a->mv1_wgrad = {.shape = {NH_MV, NH_INV_HID}}; a->mv1_bgrad = {.shape = {NH_MV}}; + a->mv2_wgrad = {.shape = {NH_MV, NH_MV}}; a->mv2_bgrad = {.shape = {NH_MV}}; + a->mr_wgrad = {.shape = {NH_INV_HID, NH_LAB_IN}}; a->mr_bgrad = {.shape = {NH_INV_HID}}; + a->mm1_wgrad = {.shape = {NH_MV, NH_INV_HID}}; a->mm1_bgrad = {.shape = {NH_MV}}; + a->mm2_wgrad = {.shape = {NH_MV, NH_MV}}; a->mm2_bgrad = {.shape = {NH_MV}}; + a->ir_wgrad = {.shape = {NH_INV_HID, NH_LAB_IN}}; a->ir_bgrad = {.shape = {NH_INV_HID}}; + a->im1_wgrad = {.shape = {NH_MV, NH_INV_HID}}; a->im1_bgrad = {.shape = {NH_MV}}; + a->im2_wgrad = {.shape = {NH_MV, NH_MV}}; a->im2_bgrad = {.shape = {NH_MV}}; + alloc_register(acts,&a->mh1); alloc_register(acts,&a->mdh1); + alloc_register(acts,&a->mvv); alloc_register(acts,&a->mdv); alloc_register(acts,&a->mvmax); + alloc_register(acts,&a->mrp); alloc_register(acts,&a->mdrp); + alloc_register(acts,&a->mh1m); alloc_register(acts,&a->mdh1m); + alloc_register(acts,&a->mvm); alloc_register(acts,&a->mdvm); alloc_register(acts,&a->mvmaxm); + alloc_register(acts,&a->irp); alloc_register(acts,&a->idrp); + alloc_register(acts,&a->mh1i); alloc_register(acts,&a->mdh1i); + alloc_register(acts,&a->mvi); alloc_register(acts,&a->mdvi); alloc_register(acts,&a->mvmaxi); + alloc_register(acts,&a->mv1b_acc); alloc_register(acts,&a->mv2b_acc); + alloc_register(acts,&a->mrb_acc); alloc_register(acts,&a->mm1b_acc); alloc_register(acts,&a->mm2b_acc); + alloc_register(acts,&a->irb_acc); alloc_register(acts,&a->im1b_acc); alloc_register(acts,&a->im2b_acc); + alloc_register(acts,&a->inv_grad); alloc_register(acts,&a->bl_grad); - alloc_register(acts,&a->dT); alloc_register(acts,&a->dT_i); alloc_register(acts,&a->dTinv); alloc_register(acts,&a->dTinv_i); alloc_register(acts,&a->dTtrue); alloc_register(acts,&a->dTtrue_i); - alloc_register(acts,&a->dE_tmp); - alloc_register(acts,&a->dw_perm); alloc_register(acts,&a->dE_i); - alloc_register(acts,&a->dw2_acc); alloc_register(acts,&a->dw2i_acc); + alloc_register(acts,&a->dE_tmp); alloc_register(acts,&a->dE_i); + alloc_register(acts,&a->loc_h1); alloc_register(acts,&a->loc_h1_grad); + alloc_register(acts,&a->loc1b_acc); alloc_register(acts,&a->dmsg_acc); - alloc_register(acts,&a->sort_local); alloc_register(acts,&a->sort_grid); alloc_register(acts,&a->bias_acc); a->embed_wgrad = {.shape = {NH_GLYPH_VOCAB, NH_EMBED_DIM}}; a->ekind_wgrad = {.shape = {NH_NKIND, NH_EMBED_DIM}}; a->esub_wgrad = {.shape = {NH_NSUB, NH_EMBED_DIM}}; - a->loc_wgrad = {.shape = {NH_LOC_HID, NH_LOC_IN}}; - a->loc_bgrad = {.shape = {NH_LOC_HID}}; - a->glb1_wgrad = {.shape = {NH_P1, NH_PCELLS * NH_EMBED_DIM}}; - a->glb1_xygrad = {.shape = {NH_P1, 2}}; - a->glb1_bgrad = {.shape = {NH_P1}}; - a->glb2_wgrad = {.shape = {NH_GLB_HID, NH_P1}}; - a->glb2_bgrad = {.shape = {NH_GLB_HID}}; + a->loc_wgrad = {.shape = {NH_LOC_H1, NH_LOC_IN}}; + a->loc_bgrad = {.shape = {NH_LOC_H1}}; + a->terr1_wgrad = {.shape = {NH_TERR_H1, NH_TERRF}}; + a->terr1_bgrad = {.shape = {NH_TERR_H1}}; + a->terr2_wgrad = {.shape = {NH_GLB_HID, NH_TERR_H1}}; + a->terr2_bgrad = {.shape = {NH_GLB_HID}}; + a->locc_wgrad = {.shape = {NH_LOCC_CLASSES, NH_LOCC_DIM}}; a->inv1_wgrad = {.shape = {NH_INV_HID, NH_EMBED_DIM}}; a->inv1_bgrad = {.shape = {NH_INV_HID}}; a->inv1s_wgrad = {.shape = {NH_INV_HID, NH_SFEAT}}; a->invt_wgrad = {.shape = {NH_INV_HID, NH_EMBED_DIM}}; - a->inv2_wgrad = {.shape = {NH_INV_POOL, NH_INV_HID}}; - a->inv2_bgrad = {.shape = {NH_INV_POOL}}; a->bl_wgrad = {.shape = {NH_BL_HID, NH_BL_FEAT}}; a->bl_bgrad = {.shape = {NH_BL_HID}}; a->proj_wgrad = {.shape = {ew->hidden, NH_CONCAT}}; a->proj_bgrad = {.shape = {ew->hidden}}; a->msg_wgrad = {.shape = {NH_MSG_VOCAB, NH_MSG_HID}}; a->spk_wgrad = {.shape = {NH_SPKEY, NH_SPIN}}; - a->spk2_wgrad = {.shape = {NH_SPKEY, NH_SPKEY}}; - a->spk2_bgrad = {.shape = {NH_SPKEY}}; -#if NH_ID_EMBED + a->spm1_wgrad = {.shape = {NH_SPM, NH_SPKEY}}; + a->spm1_bgrad = {.shape = {NH_SPM}}; + a->spm2_wgrad = {.shape = {NH_SPM, NH_SPM}}; + a->spm2_bgrad = {.shape = {NH_SPM}}; a->ide_role_wgrad = {.shape = {13, NH_IDE_ROLE}}; a->ide_race_wgrad = {.shape = {5, NH_IDE_RACE}}; a->ide_gend_wgrad = {.shape = {2, NH_IDE_GEND}}; a->ide_algn_wgrad = {.shape = {3, NH_IDE_ALGN}}; -#endif alloc_register(grads,&a->embed_wgrad); alloc_register(grads,&a->ekind_wgrad); alloc_register(grads,&a->esub_wgrad); alloc_register(grads,&a->loc_wgrad); alloc_register(grads,&a->loc_bgrad); - alloc_register(grads,&a->glb1_wgrad); alloc_register(grads,&a->glb1_xygrad); - alloc_register(grads,&a->glb1_bgrad); - alloc_register(grads,&a->glb2_wgrad); alloc_register(grads,&a->glb2_bgrad); + alloc_register(grads,&a->loc2_wgrad); alloc_register(grads,&a->loc2_bgrad); + alloc_register(grads,&a->terr1_wgrad); alloc_register(grads,&a->terr1_bgrad); + alloc_register(grads,&a->terr2_wgrad); alloc_register(grads,&a->terr2_bgrad); + alloc_register(grads,&a->locc_wgrad); alloc_register(grads,&a->inv1_wgrad); alloc_register(grads,&a->inv1_bgrad); alloc_register(grads,&a->inv1s_wgrad); alloc_register(grads,&a->invt_wgrad); - alloc_register(grads,&a->inv2_wgrad); alloc_register(grads,&a->inv2_bgrad); alloc_register(grads,&a->bl_wgrad); alloc_register(grads,&a->bl_bgrad); alloc_register(grads,&a->proj_wgrad); alloc_register(grads,&a->proj_bgrad); alloc_register(grads,&a->msg_wgrad); alloc_register(grads,&a->spk_wgrad); - alloc_register(grads,&a->spk2_wgrad); alloc_register(grads,&a->spk2_bgrad); -#if NH_ID_EMBED + alloc_register(grads,&a->spm1_wgrad); alloc_register(grads,&a->spm1_bgrad); + alloc_register(grads,&a->spm2_wgrad); alloc_register(grads,&a->spm2_bgrad); alloc_register(grads,&a->ide_role_wgrad); alloc_register(grads,&a->ide_race_wgrad); alloc_register(grads,&a->ide_gend_wgrad); alloc_register(grads,&a->ide_algn_wgrad); -#endif + alloc_register(grads,&a->mv1_wgrad); alloc_register(grads,&a->mv1_bgrad); + alloc_register(grads,&a->mv2_wgrad); alloc_register(grads,&a->mv2_bgrad); + alloc_register(grads,&a->mr_wgrad); alloc_register(grads,&a->mr_bgrad); + alloc_register(grads,&a->mm1_wgrad); alloc_register(grads,&a->mm1_bgrad); + alloc_register(grads,&a->mm2_wgrad); alloc_register(grads,&a->mm2_bgrad); + alloc_register(grads,&a->ir_wgrad); alloc_register(grads,&a->ir_bgrad); + alloc_register(grads,&a->im1_wgrad); alloc_register(grads,&a->im1_bgrad); + alloc_register(grads,&a->im2_wgrad); alloc_register(grads,&a->im2_bgrad); nh_enc_last = a; } @@ -1805,21 +2741,25 @@ static void nethack_encoder_reg_rollout(void* w, void* activations, Allocator* a a->crop_glyph = {.shape = {B, NH_CGRID}}; a->e_eff = {.shape = {NH_GLYPH_VOCAB, NH_EMBED_DIM}}; a->x_local = {.shape = {B, NH_LOC_IN}}; - a->w_perm = {.shape = {NH_TROW, NH_EMBED_DIM}}; - a->glyph_T = {.shape = {NH_GLYPH_VOCAB, NH_TROW}}; - a->t16 = {.shape = {B, NH_TOK * NH_P1}}; - a->dxy = {.shape = {B, NH_TOK * 2}}; - a->tok_argmax = {.shape = {B, NH_GLB_HID}}; + a->terr_tf = {.shape = {B, NH_TERRF}}; + a->terr_h = {.shape = {B, NH_TERR_H1}}; a->inv_idx = {.shape = {B, NH_INV}}; a->invt_idx = {.shape = {B, NH_INV}}; a->spell_idx = {.shape = {B, 8}}; a->inv_sfeat = {.shape = {B, NH_INV * NH_SFEAT}}; - a->inv_T = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; - a->invt_T = {.shape = {NH_GLYPH_VOCAB, NH_INV_HID}}; + a->inv_T = {.shape = {NH_ITBL, NH_INV_HID}}; + a->invt_T = {.shape = {NH_ITBL, NH_INV_HID}}; a->inv_out = {.shape = {B, NH_INV_FLAT}}; - a->inv_pool = {.shape = {B, NH_INV_POOL}}; - a->inv_amax = {.shape = {B, NH_INV_POOL}}; a->loc_out = {.shape = {B, NH_LOC_HID}}; + a->loc_h1 = {.shape = {B, NH_LOC_H1}}; + a->mh1 = {.shape = {B, NH_INV * NH_MV}}; + a->mvv = {.shape = {B, NH_INV * NH_MV}}; a->mvmax = {.shape = {B, NH_MV}}; + a->mrp = {.shape = {B, NH_LABK * NH_INV_HID}}; + a->mh1m = {.shape = {B, NH_LABK * NH_MV}}; + a->mvm = {.shape = {B, NH_LABK * NH_MV}}; a->mvmaxm = {.shape = {B, NH_MV}}; + a->irp = {.shape = {B, NH_LABK * NH_INV_HID}}; + a->mh1i = {.shape = {B, NH_LABK * NH_MV}}; + a->mvi = {.shape = {B, NH_LABK * NH_MV}}; a->mvmaxi = {.shape = {B, NH_MV}}; a->glb_out = {.shape = {B, NH_GLB_HID}}; a->bl_feats = {.shape = {B, NH_BL_FEAT}}; a->bl_out = {.shape = {B, NH_BL_HID}}; @@ -1827,33 +2767,43 @@ static void nethack_encoder_reg_rollout(void* w, void* activations, Allocator* a a->msg_out = {.shape = {B, NH_MSG_HID}}; a->spk_in = {.shape = {B, NH_SPELL_SLOTS * NH_SPIN}}; a->spk_keys = {.shape = {B, NH_SPELL_SLOTS * NH_SPKEY}}; - a->spk_amax = {.shape = {B, NH_SPKEY}}; - a->spk_pool = {.shape = {B, NH_SPKEY}}; -#if NH_ID_EMBED + a->sph1 = {.shape = {B, NH_SPELL_SLOTS * NH_SPM}}; + a->spv = {.shape = {B, NH_SPELL_SLOTS * NH_SPM}}; + a->spvmax = {.shape = {B, NH_SPM}}; a->ide_idx = {.shape = {B, 4}}; -#endif + a->lm_tok = {.shape = {B, NH_LABK * NH_LAB_IN}}; + a->lm_gid = {.shape = {B, NH_LABK}}; + a->li_tok = {.shape = {B, NH_LABK * NH_LAB_IN}}; + a->li_gid = {.shape = {B, NH_LABK}}; a->concat = {.shape = {B, NH_CONCAT}}; a->out = {.shape = {B, ew->hidden}}; alloc_register(alloc,&a->glyph_idx); alloc_register(alloc,&a->crop_glyph); alloc_register(alloc,&a->e_eff); alloc_register(alloc,&a->x_local); - alloc_register(alloc,&a->w_perm); alloc_register(alloc,&a->glyph_T); - alloc_register(alloc,&a->t16); alloc_register(alloc,&a->dxy); - alloc_register(alloc,&a->tok_argmax); + alloc_register(alloc,&a->terr_tf); alloc_register(alloc,&a->terr_h); alloc_register(alloc,&a->inv_idx); alloc_register(alloc,&a->invt_idx); alloc_register(alloc,&a->spell_idx); -#if NH_ID_EMBED alloc_register(alloc,&a->ide_idx); -#endif alloc_register(alloc,&a->inv_sfeat); alloc_register(alloc,&a->inv_T); alloc_register(alloc,&a->invt_T); alloc_register(alloc,&a->inv_out); - alloc_register(alloc,&a->inv_pool); alloc_register(alloc,&a->inv_amax); alloc_register(alloc,&a->loc_out); alloc_register(alloc,&a->glb_out); + alloc_register(alloc,&a->loc_h1); + alloc_register(alloc,&a->mh1); alloc_register(alloc,&a->mvv); alloc_register(alloc,&a->mvmax); + alloc_register(alloc,&a->mrp); alloc_register(alloc,&a->mh1m); + alloc_register(alloc,&a->mvm); alloc_register(alloc,&a->mvmaxm); + alloc_register(alloc,&a->irp); alloc_register(alloc,&a->mh1i); + alloc_register(alloc,&a->mvi); alloc_register(alloc,&a->mvmaxi); alloc_register(alloc,&a->bl_feats); alloc_register(alloc,&a->bl_out); alloc_register(alloc,&a->msg_ids); alloc_register(alloc,&a->msg_out); alloc_register(alloc,&a->spk_in); alloc_register(alloc,&a->spk_keys); - alloc_register(alloc,&a->spk_amax); alloc_register(alloc,&a->spk_pool); + alloc_register(alloc,&a->sph1); + alloc_register(alloc,&a->spv); + alloc_register(alloc,&a->spvmax); + alloc_register(alloc,&a->lm_tok); + alloc_register(alloc,&a->lm_gid); + alloc_register(alloc,&a->li_tok); + alloc_register(alloc,&a->li_gid); alloc_register(alloc,&a->concat); alloc_register(alloc,&a->out); nh_enc_last = a; } @@ -1929,7 +2879,7 @@ struct NethackDecoderActivations { Prec dtmp, dq; Prec keygrad; // (B, NH_INV_FLAT) -> encoder inv slice Prec kmat; // (B, NH_INV_FLAT) projected keys - Prec kn, qn; // key norms (B, NH_INV), query norms (B, NH_QHEADS) + Prec qn; // query norms (B, NH_QHEADS) Prec slot_logits; // (B, NH_SLOT_OD) tau_h * cos Prec dkmat; // backward scratch Prec spdk; // spell-key grads from the pointer (B, 8*NH_SPKEY) @@ -1977,10 +2927,12 @@ __global__ void nh_ptr_rownorm_kernel(precision_t* __restrict__ n, n[r] = from_float(sqrtf(acc) + 1e-6f); } -// slot logit = exp(ltau_h) * cos(q_h, k_i), one thread per (sample, head, slot) +// slot logit = exp(ltau_h) * (qhat_h . k_i): query-only normalization — key +// magnitude reaches the logit (decoder lab: composite gear selection +3..21pp +// vs full cosine; adopted 2026-08-25, pair +386). One thread per (b, head, slot). __global__ void nh_ptr3_cos_kernel(precision_t* __restrict__ slot_logits, const precision_t* __restrict__ q, const precision_t* __restrict__ qn, - const precision_t* __restrict__ kmat, const precision_t* __restrict__ kn, + const precision_t* __restrict__ kmat, const precision_t* __restrict__ tau, int B) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= B * NH_SLOT_OD) return; @@ -1992,7 +2944,7 @@ __global__ void nh_ptr3_cos_kernel(precision_t* __restrict__ slot_logits, for (int k = 0; k < NH_INV_HID; k++) dot += to_float(qb[k]) * to_float(ki[k]); slot_logits[idx] = from_float(expf(to_float(tau[h])) * dot / - (to_float(qn[(int64_t)b * NH_QHEADS + h]) * to_float(kn[(int64_t)b * NH_INV + i]))); + to_float(qn[(int64_t)b * NH_QHEADS + h])); } __global__ void nh_dec_dtmp_kernel(precision_t* __restrict__ dtmp, @@ -2018,7 +2970,7 @@ __global__ void nh_ptr3_dq_kernel(precision_t* __restrict__ dq, long long* __restrict__ tau_acc, const precision_t* __restrict__ g, const precision_t* __restrict__ out, const precision_t* __restrict__ q, const precision_t* __restrict__ qn, const precision_t* __restrict__ kmat, - const precision_t* __restrict__ kn, const precision_t* __restrict__ tau, int B) { + const precision_t* __restrict__ tau, int B) { int bh = blockIdx.x * blockDim.x + threadIdx.x; if (bh >= B * NH_HEADS) return; int b = bh / NH_HEADS, h = bh % NH_HEADS; @@ -2037,10 +2989,9 @@ __global__ void nh_ptr3_dq_kernel(precision_t* __restrict__ dq, if (gi == 0.0f) continue; float cosv = to_float(out[gbase + i]) / tauv; dtau += gi * cosv; - float knv = to_float(kn[(int64_t)b * NH_INV + i]); const precision_t* ki = kmat + ((int64_t)b * NH_INV + i) * NH_INV_HID; for (int k = 0; k < NH_INV_HID; k++) - dv[k] += tauv * gi * to_float(ki[k]) / knv; + dv[k] += tauv * gi * to_float(ki[k]); } float vdv = 0.0f; for (int k = 0; k < NH_INV_HID; k++) vdv += vhat[k] * dv[k]; @@ -2073,7 +3024,7 @@ __global__ void nh_spq_bwd_kernel(precision_t* __restrict__ dq, __global__ void nh_ptr3_dkmat_kernel(precision_t* __restrict__ dkmat, const precision_t* __restrict__ g, const precision_t* __restrict__ out, const precision_t* __restrict__ q, const precision_t* __restrict__ qn, - const precision_t* __restrict__ kmat, const precision_t* __restrict__ kn, + const precision_t* __restrict__ kmat, const precision_t* __restrict__ tau, int B) { int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; if (idx >= (int64_t)B * NH_INV_FLAT) return; @@ -2081,22 +3032,20 @@ __global__ void nh_ptr3_dkmat_kernel(precision_t* __restrict__ dkmat, int64_t b = bi / NH_INV; int i = (int)(bi % NH_INV); int k = (int)(idx % NH_INV_HID); - float knv = to_float(kn[bi]); - float uk = to_float(kmat[idx]) / knv; float acc = 0.0f; for (int h = 0; h < NH_HEADS; h++) { int64_t gi_idx = b * (NH_DEC_OD + 1) + NH_ACTIONS + h * NH_INV + i; float gi = to_float(g[gi_idx]); if (gi == 0.0f) continue; float tauv = expf(to_float(tau[h])); - float cosv = to_float(out[gi_idx]) / tauv; float vk = to_float(q[((int64_t)b * NH_QHEADS + h) * NH_INV_HID + k]) / to_float(qn[(int64_t)b * NH_QHEADS + h]); - acc += tauv * gi * (vk - uk * cosv); + acc += tauv * gi * vk; } - dkmat[idx] = from_float(acc / knv); + dkmat[idx] = from_float(acc); } + static Prec nethack_decoder_forward(void* w, void* activations, Prec input, cudaStream_t stream) { NethackDecoderWeights* dw = (NethackDecoderWeights*)w; NethackDecoderActivations* a = (NethackDecoderActivations*)activations; @@ -2110,10 +3059,8 @@ static Prec nethack_decoder_forward(void* w, void* activations, Prec input, cuda puf_mm(&sflat, &dw->k_w, &kflat, stream); nh_ptr_rownorm_kernel<<>>( a->qn.data, a->q.data, B * NH_QHEADS); - nh_ptr_rownorm_kernel<<>>( - a->kn.data, a->kmat.data, B * NH_INV); nh_ptr3_cos_kernel<<>>( - a->slot_logits.data, a->q.data, a->qn.data, a->kmat.data, a->kn.data, dw->tau.data, B); + a->slot_logits.data, a->q.data, a->qn.data, a->kmat.data, dw->tau.data, B); nh_dec_assemble_kernel<<>>( a->out.data, a->tmp.data, a->slot_logits.data, a->q.data, ea->spk_keys.data, B); return a->out; @@ -2121,6 +3068,7 @@ static Prec nethack_decoder_forward(void* w, void* activations, Prec input, cuda static Prec nethack_decoder_backward(void* w, void* activations, Float grad_logits, Float grad_logstd, Float grad_value, cudaStream_t stream) { + (void)grad_logstd; NethackDecoderWeights* dw = (NethackDecoderWeights*)w; NethackDecoderActivations* a = (NethackDecoderActivations*)activations; int B = a->saved_input.shape[0]; @@ -2132,14 +3080,14 @@ static Prec nethack_decoder_backward(void* w, void* activations, cudaMemsetAsync(a->tau_acc.data, 0, NH_TAU_PAD * sizeof(long long), stream); nh_ptr3_dq_kernel<<>>( a->dq.data, (long long*)a->tau_acc.data, a->grad_out.data, a->out.data, - a->q.data, a->qn.data, a->kmat.data, a->kn.data, dw->tau.data, B); + a->q.data, a->qn.data, a->kmat.data, dw->tau.data, B); nh_fxp_to_precision_kernel<<<1, 32, 0, stream>>>( a->tau_grad.data, (long long*)a->tau_acc.data, NH_TAU_PAD); nh_spq_bwd_kernel<<>>( a->dq.data, a->spdk.data, a->grad_out.data, a->q.data, ea->spk_keys.data, B); nh_ptr3_dkmat_kernel<<>>( a->dkmat.data, a->grad_out.data, a->out.data, a->q.data, a->qn.data, - a->kmat.data, a->kn.data, dw->tau.data, B); + a->kmat.data, dw->tau.data, B); // dK = dkmat^T @ s ; keygrad (ds, into the encoder inv slice) = dkmat @ K Prec dkflat = {.data = a->dkmat.data, .shape = {B * NH_INV, NH_INV_HID}}; Prec sflat = {.data = ea->inv_out.data, .shape = {B * NH_INV, NH_INV_HID}}; @@ -2195,7 +3143,6 @@ static void nethack_decoder_reg_train(void* w, void* activations, Allocator* act a->dq = {.shape = {B_TT, NH_QDIM}}; a->keygrad = {.shape = {B_TT, NH_INV_FLAT}}; a->kmat = {.shape = {B_TT, NH_INV_FLAT}}; - a->kn = {.shape = {B_TT, NH_INV}}; a->qn = {.shape = {B_TT, NH_QHEADS}}; a->slot_logits = {.shape = {B_TT, NH_SLOT_OD}}; a->dkmat = {.shape = {B_TT, NH_INV_FLAT}}; @@ -2210,7 +3157,7 @@ static void nethack_decoder_reg_train(void* w, void* activations, Allocator* act alloc_register(acts,&a->grad_input); alloc_register(acts,&a->grad_input2); alloc_register(acts,&a->grad_out); alloc_register(acts,&a->dtmp); alloc_register(acts,&a->dq); alloc_register(acts,&a->keygrad); - alloc_register(acts,&a->kmat); alloc_register(acts,&a->kn); + alloc_register(acts,&a->kmat); alloc_register(acts,&a->qn); alloc_register(acts,&a->slot_logits); alloc_register(acts,&a->dkmat); alloc_register(acts,&a->tau_acc); alloc_register(acts,&a->spdk); @@ -2222,20 +3169,19 @@ static void nethack_decoder_reg_train(void* w, void* activations, Allocator* act } static void nethack_decoder_reg_rollout(void* w, void* activations, Allocator* alloc, int B) { + (void)w; // rollout shapes are all compile-time constants NethackDecoderActivations* a = (NethackDecoderActivations*)activations; a->enc = nh_enc_last; a->out = {.shape = {B, NH_DEC_OD + 1}}; a->tmp = {.shape = {B, NH_DEC_PAD}}; a->q = {.shape = {B, NH_QDIM}}; a->kmat = {.shape = {B, NH_INV_FLAT}}; - a->kn = {.shape = {B, NH_INV}}; a->qn = {.shape = {B, NH_QHEADS}}; a->slot_logits = {.shape = {B, NH_SLOT_OD}}; alloc_register(alloc,&a->out); alloc_register(alloc,&a->tmp); alloc_register(alloc,&a->q); alloc_register(alloc,&a->kmat); - alloc_register(alloc,&a->kn); alloc_register(alloc,&a->qn); alloc_register(alloc,&a->slot_logits); } diff --git a/ocean/nethack/nethack.h b/ocean/nethack/nethack.h index 09900feb86..8396323db4 100644 --- a/ocean/nethack/nethack.h +++ b/ocean/nethack/nethack.h @@ -36,10 +36,11 @@ extern int nle_container_at(nle_ctx_t*); extern int nle_food_underfoot(nle_ctx_t*); extern int nle_discoveries(nle_ctx_t*); extern int nle_peaceful_at(nle_ctx_t*, int, int); -extern int nle_spellprot(nle_ctx_t*); +extern int nle_lnc_bits(nle_ctx_t*); extern void nle_weight(nle_ctx_t*, int*, int*); extern int nle_spells(nle_ctx_t*, short*, signed char*, signed char*, int*, int); extern int nle_cast_blocked(nle_ctx_t*); +extern int nle_intrinsics(nle_ctx_t*); extern void nle_end(nle_ctx_t*); extern void nle_identity(nle_ctx_t*, int*, int*, int*, int*); #ifdef __cplusplus @@ -102,24 +103,23 @@ struct Env { long prev_exp; long prev_gold; long start_gold; - float prev_ac_led; // ac-delta ledger: last ledgered AC (durable-weighted) - float ac_account; // accrued unpaid ac-delta reward (carries) long prev_time; int prev_depth; unsigned prev_floor; // dnum << 8 | dlevel at last reward; guards path attribution int disc0; // discoveries count at reset (episode delta = types learned) unsigned long long engid_tested; // letters engrave-tested this episode + unsigned char terr_mem[NH_GRID]; // remembered terrain (cmap_index+1; 0 unseen) + unsigned short obj_mem[NH_GRID]; // v3: remembered floor object (item row+1; 0 none) + unsigned terr_floor; // dnum<<8|dlevel the memory belongs to // reward coefs float gold_coef; + float score_coef; // in-game score delta per step (0 = off) float exp_coef; float descent_coef; float floor_coef; float xp_coef; float scout_coef; - float ac_coef; - float scout_ready; // 0 = off; else a tile pays pro-rata to xp level vs depth - float ac_nospell; // unpaid fraction of protection-spell AC (1 = durable AC only) float death_penalty; float mask_search20; // 1 removes SEARCH20 from the action space float mask_run; // 1 removes RUN from the action space @@ -132,6 +132,7 @@ struct Env { #include "macros.h" + // init // demo-only obs planes; NULL in training (fills skipped) @@ -177,20 +178,9 @@ static void nethack_init_settings(Nethack* env) { } env->settings.spawn_monsters = 1; env->settings.underfoot_glyphs = 1; // underfoot shows objects - const char* role = getenv("NH_ROLE"); - char optbuf[512], rcbuf[512]; - const char* opts; - if (role && role[0]) { - snprintf(optbuf, sizeof(optbuf), - "name:Agent,role:%s,race:random,gender:random,align:random," - NETHACK_OPTIONS_TAIL "!status_updates", role); - opts = optbuf; - } else { - opts = nethack_options_override ? nethack_options_override - : NETHACK_DEFAULT_OPTIONS; - } snprintf(env->settings.options, sizeof(env->settings.options), "@%s", - nethack_rc_path_opts(rcbuf, sizeof(rcbuf), opts)); + nethack_rc_path(nethack_options_override + ? nethack_options_override : NETHACK_DEFAULT_OPTIONS)); env->settings.fix_moon_phase = true; // moon phase from seed } @@ -204,8 +194,16 @@ void init(Nethack* env) { // masking + static int nethack_slot_usable(const Nethack* env, const Verb* verb, int i) { if (!(verb->item_classes & (1u << env->inv_oclasses[i]))) return 0; + // APPLY on a container (box/bag) is a silent zero-turn no-op: the macro + // drives no put-in/take-out menus, so it is an absorbing spam loop, not a + // choice. Bag of tricks (spawns monsters, costs a turn) stays legal. + if (verb == &NETHACK_VERBS[NETHACK_ACT_APPLY]) { + int a = env->inv_glyphs[i] - NH_GLYPH_OBJ_OFF; + if (a >= 189 /* LARGE_BOX */ && a <= 194 /* BAG_OF_HOLDING */) return 0; + } // READ hygiene: blind reads refuse for free, and re-reading a still-fresh // book is a multi-turn re-study furnace; low-retention refresh stays legal if (verb->item_classes == ((1u << 9) | (1u << 10))) { @@ -444,9 +442,8 @@ static void nethack_compute_mask(Nethack* env) { // observations static void nethack_pack_obs(Nethack* env) { - obs_t* obs_buf = env->agents[0].observations; - memcpy(obs_buf + NETHACK_OFF_GLYPHS, env->glyphs, sizeof(env->glyphs)); - unsigned char* bl = obs_buf + NETHACK_OFF_BLSTATS; + memcpy(((obs_t*)env->agents[0].observations) + NETHACK_OFF_GLYPHS, env->glyphs, sizeof(env->glyphs)); + unsigned char* bl = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_BLSTATS; for (int i = 0; i < NLE_BLSTATS_SIZE; i++) { uint32_t v = (uint32_t)(int32_t)env->blstats[i]; bl[4*i + 0] = (unsigned char)(v & 0xffu); @@ -478,6 +475,7 @@ static void nethack_pack_obs(Nethack* env) { q[3] = known ? env->spell_knows[s] : 0; } + extra[NETHACK_EXTRA_INTRINS] = nle_intrinsics(env->ctx); extra[NETHACK_EXTRA_ROLEOH + env->role_idx] = 1; extra[NETHACK_EXTRA_ROLEOH + 13 + env->race_idx] = 1; extra[NETHACK_EXTRA_ROLEOH + 18 + env->gend_idx] = 1; @@ -495,7 +493,7 @@ static void nethack_pack_obs(Nethack* env) { extra[NETHACK_EXTRA_SHOP + 1] = (price > 0) ? (int32_t)(gold >= price ? 100 : (gold * 100) / price) : 0; - unsigned char* ex = obs_buf + NETHACK_OFF_EXTRA; + unsigned char* ex = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_EXTRA; for (int i = 0; i < NETHACK_EXTRA_INTS; i++) { uint32_t v = (uint32_t)extra[i]; ex[4*i + 0] = (unsigned char)(v & 0xffu); @@ -510,7 +508,7 @@ static void nethack_pack_obs(Nethack* env) { const char* e = getenv("NH_DISC_SWAP"); dsw = e && e[0] && e[0] != '0'; } - unsigned char* iv = obs_buf + NETHACK_OFF_INV; + unsigned char* iv = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_INV; for (int i = 0; i < NETHACK_INV_SLOTS; i++) { uint16_t g = env->inv_oclasses[i] < NETHACK_NUM_OCLASSES ? (dsw && env->inv_true[i] != NETHACK_PAD_GLYPH @@ -520,9 +518,9 @@ static void nethack_pack_obs(Nethack* env) { iv[2*i + 1] = (unsigned char)((g >> 8) & 0xffu); } // item state - memcpy(obs_buf + NETHACK_OFF_INVST, env->inv_state, sizeof(env->inv_state)); + memcpy(((obs_t*)env->agents[0].observations) + NETHACK_OFF_INVST, env->inv_state, sizeof(env->inv_state)); // discovered-type glyphs (engine pads with NO_GLYPH == NETHACK_PAD_GLYPH) - unsigned char* it = obs_buf + NETHACK_OFF_INVTRUE; + unsigned char* it = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_INVTRUE; for (int i = 0; i < NETHACK_INV_SLOTS; i++) { uint16_t g = !dsw && env->inv_oclasses[i] < NETHACK_NUM_OCLASSES ? (uint16_t)env->inv_true[i] : (uint16_t)NETHACK_PAD_GLYPH; @@ -545,8 +543,84 @@ static void nethack_pack_obs(Nethack* env) { env->stall_ctr = 0; env->stall_prev_turn = turn; } + { // typed level planes: terrain memory, floor-object memory, two token lists + unsigned fl = (unsigned)((env->blstats[NLE_BL_DNUM] << 8) | env->blstats[NLE_BL_DLEVEL]); + if (fl != env->terr_floor) { + memset(env->terr_mem, 0, sizeof(env->terr_mem)); + memset(env->obj_mem, 0, sizeof(env->obj_mem)); + env->terr_floor = fl; + } + long hx = env->blstats[NLE_BL_X], hy = env->blstats[NLE_BL_Y]; + int mon_cell[256], mon_g[256], nmon = 0; + for (int cell = 0; cell < NH_GRID; cell++) { + int g = env->glyphs[cell]; + int r = cell / NH_COLS, c = cell % NH_COLS; + if (g >= NETHACK_GLYPH_CMAP_OFF && g < NETHACK_GLYPH_CMAP_OFF + 96) { + env->terr_mem[cell] = (unsigned char)(g - NETHACK_GLYPH_CMAP_OFF + 1); + env->obj_mem[cell] = 0; // bare floor shown: no object here + } else if (g >= NETHACK_GLYPH_OBJ_LO && g < NETHACK_GLYPH_OBJ_HI) { + env->obj_mem[cell] = (unsigned short)(g - NETHACK_GLYPH_OBJ_LO + 1); + } else if (g >= NETHACK_GLYPH_BODY_OFF && g < NETHACK_GLYPH_BODY_HI) { + env->obj_mem[cell] = (unsigned short)(g - NETHACK_GLYPH_BODY_OFF + 454); + } else if (((g >= 0 && g < NETHACK_GLYPH_BODY_OFF && !(r == hy && c == hx))) + && nmon < 256) { // monsters, pets and detected alike + mon_cell[nmon] = cell; mon_g[nmon] = g; nmon++; + // monster on top: obj_mem/terr_mem memories stay as last seen + } + } + // lean token block: lists only, memory planes stay env-internal + unsigned char* mp = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_TOKM; + // monster tokens: K nearest by Chebyshev (selection sort) + memset(mp, 0, NETHACK_V3_K * NETHACK_V3_MONF); + for (int k = 0; k < NETHACK_V3_K && k < nmon; k++) { + int best = -1, bd = 1 << 30; + for (int i = 0; i < nmon; i++) { + if (mon_cell[i] < 0) continue; + int dx = mon_cell[i] % NH_COLS - (int)hx, dy = mon_cell[i] / NH_COLS - (int)hy; + int d = abs(dx) > abs(dy) ? abs(dx) : abs(dy); + if (d < bd) { bd = d; best = i; } + } + if (best < 0) break; + int cell = mon_cell[best], g = mon_g[best]; mon_cell[best] = -1; + int dx = cell % NH_COLS - (int)hx, dy = cell / NH_COLS - (int)hy; + int det = g >= NETHACK_GLYPH_DET_OFF; + int pet = g >= NETHACK_GLYPH_PET_OFF && g < NETHACK_GLYPH_DET_OFF; + int sp = g % NETHACK_NUMMONS; // mon/pet/detect ranges all stride NUMMONS + int peace = pet || nle_peaceful_at(env->ctx, cell % NH_COLS + 1, cell / NH_COLS); + unsigned char* e = mp + k * NETHACK_V3_MONF; + e[0] = (unsigned char)((sp + 1) & 0xff); e[1] = (unsigned char)(((sp + 1) >> 8) & 0xff); + e[2] = (unsigned char)(signed char)dx; e[3] = (unsigned char)(signed char)dy; + e[4] = (unsigned char)((peace ? 0 : 1) | (peace ? 2 : 0) | (det ? 4 : 0) | (pet ? 8 : 0)); + e[5] = NH_MON_DIFF[sp]; e[6] = NH_MON_SPEED[sp]; e[7] = 0; + } + // item tokens from the object memory (includes remembered/occluded items) + int it_cell[256], nit = 0; + for (int cell = 0; cell < NH_GRID && nit < 256; cell++) + if (env->obj_mem[cell]) it_cell[nit++] = cell; + unsigned char* ip = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_TOKI; + memset(ip, 0, NETHACK_V3_K * NETHACK_V3_ITEMF); + for (int k = 0; k < NETHACK_V3_K && k < nit; k++) { + int best = -1, bd = 1 << 30; + for (int i = 0; i < nit; i++) { + if (it_cell[i] < 0) continue; + int dx = it_cell[i] % NH_COLS - (int)hx, dy = it_cell[i] / NH_COLS - (int)hy; + int d = abs(dx) > abs(dy) ? abs(dx) : abs(dy); + if (d < bd) { bd = d; best = i; } + } + if (best < 0) break; + int cell = it_cell[best]; it_cell[best] = -1; + int row = env->obj_mem[cell]; + int dx = cell % NH_COLS - (int)hx, dy = cell / NH_COLS - (int)hy; + unsigned char* e = ip + k * NETHACK_V3_ITEMF; + e[0] = (unsigned char)(row & 0xff); e[1] = (unsigned char)((row >> 8) & 0xff); + e[2] = (unsigned char)(signed char)dx; e[3] = (unsigned char)(signed char)dy; + e[4] = 0; // oclass placeholder (item row already encodes type) + e[5] = (unsigned char)(((dx == 0 && dy == 0) ? 1 : 0) | (row >= 454 ? 2 : 0)); + e[6] = 0; e[7] = 0; + } + } // topline - unsigned char* mv = obs_buf + NETHACK_OFF_MSG; + unsigned char* mv = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_MSG; size_t mlen = strnlen((const char*)env->message, NETHACK_MSG_LEN); memcpy(mv, env->message, mlen); if (mlen < (size_t)NETHACK_MSG_LEN) memset(mv + mlen, 0, NETHACK_MSG_LEN - mlen); @@ -571,7 +645,6 @@ static void nethack_add_log(Nethack* env, int how) { // how: nle how_done, -1 = env->log.depth_15 += env->stats.max_depth >= 15 ? 1.0f : 0.0f; env->log.mines_depth += (float)__builtin_popcountll(env->stats.floors_bits[2]); env->log.sokoban_depth += (float)__builtin_popcountll(env->stats.floors_bits[4]); - env->log.scout_held += (float)env->stats.scout_held; env->log.enhances += (float)env->stats.enhances; env->log.floor_eats += (float)env->stats.floor_eats; env->log.reads_scroll += (float)env->stats.reads_scroll; @@ -585,6 +658,8 @@ static void nethack_add_log(Nethack* env, int how) { // how: nle how_done, -1 = ? (float)env->stats.burdened_steps / (float)env->stats.length : 0.0f; env->log.game_time += (float)env->prev_time; env->log.max_xp_level += (float)env->stats.max_xp; + env->log.gold += (float)(env->stats.last_gold - env->start_gold); + env->log.exp_points += (float)env->stats.last_exp; env->log.episode_return += env->stats.ret; env->log.episode_length += env->stats.length; if (how == -1) env->log.truncated += 1.0f; @@ -597,6 +672,7 @@ static void nethack_add_log(Nethack* env, int how) { // how: nle how_done, -1 = env->log.death_mon_level += (float)env->internal[NETHACK_INTERNAL_KILLER_MLEV]; if (how >= 0) env->log.death_ac += (float)env->stats.last_ac; env->log.n += 1.0f; + } // reset @@ -615,8 +691,12 @@ static void nethack_do_reset(Nethack* env) { // seed advance env->seed = env->seed * 6364136223846793005UL + 1442695040888963407UL; - // engine-random character per reset; identity read back after start - if (env->multi_role != 0.0f) { + // engine-random character per reset; identity read back after start. + // Demo NH_ROLE sets nethack_options_override; keep it across deaths. + if (nethack_options_override) { + snprintf(env->settings.options, sizeof(env->settings.options), "@%s", + nethack_rc_path(nethack_options_override)); + } else if (env->multi_role != 0.0f) { char rcp[512]; snprintf(env->settings.options, sizeof(env->settings.options), "@%s", nethack_rc_path_opts(rcp, sizeof(rcp), @@ -655,12 +735,13 @@ static void nethack_do_reset(Nethack* env) { env->engid_tested = 0; env->enh_ready = 0; memset(&env->stats, 0, sizeof(env->stats)); + memset(env->terr_mem, 0, sizeof(env->terr_mem)); + memset(env->obj_mem, 0, sizeof(env->obj_mem)); + env->terr_floor = 0xFFFFFFFFu; env->stats.max_depth = env->prev_depth; env->stats.max_xp = (int)env->blstats[NLE_BL_XP]; env->stats.min_ac = (int)env->blstats[NLE_BL_AC]; env->stats.last_ac = (int)env->blstats[NLE_BL_AC]; - env->prev_ac_led = (float)env->blstats[NLE_BL_AC]; - env->ac_account = 0.0f; nethack_pack_obs(env); } @@ -678,11 +759,8 @@ static void nethack_update_stats(Nethack* env) { env->prev_depth = (int)env->blstats[NLE_BL_DEPTH]; } -// Fractional scout claim, keyed by (dnum, dlevel). A tile pays its full -// scout_coef only once the hero's xp level covers depth * scout_ready; below -// that it pays pro-rata and the remainder stays claimable by a stronger -// visit. Total over all visits is capped at 1.0, so revisiting cannot farm -// it. scout_ready <= 0 restores plain first-visit semantics. +// First-visit scout claim, keyed by (dnum, dlevel): a tile pays once per +// episode, so revisiting cannot farm it. static float nethack_tile_claim(Nethack* env, long dn, long dl, long px, long py) { if (px < 0 || px >= NH_COLS || py < 0 || py >= NH_ROWS) return 0.0f; if (dn < 0 || dn > 15 || dl < 1 || dl > 64) return 0.0f; @@ -699,26 +777,15 @@ static float nethack_tile_claim(Nethack* env, long dn, long dl, long px, long py env->stats.visited_key[d] = key; } int idx = (int)py * NH_COLS + (int)px; - unsigned char prev = env->stats.visited[d][idx]; - if (env->scout_ready <= 0.0f) { // plain first-visit - if (prev) return 0.0f; - env->stats.visited[d][idx] = 1; - return 1.0f; - } - int depth = (int)env->blstats[NLE_BL_DEPTH]; - if (depth < 1) depth = 1; - int req = (int)((float)depth * env->scout_ready + 0.5f); - if (req < 1) req = 1; - if (req > 255) req = 255; - int cap = env->stats.max_xp < req ? env->stats.max_xp : req; // max_xp is monotonic - if (cap <= (int)prev) return 0.0f; - env->stats.visited[d][idx] = (unsigned char)cap; - return (float)(cap - (int)prev) / (float)req; + if (env->stats.visited[d][idx]) return 0.0f; + env->stats.visited[d][idx] = 1; + return 1.0f; } static float nethack_reward(Nethack* env) { // death payout if (env->obs.done) return env->death_penalty; + long score_before = env->prev_score; // update_stats overwrites prev_score below nethack_update_stats(env); int depth = (int)env->blstats[NLE_BL_DEPTH]; @@ -763,6 +830,9 @@ static float nethack_reward(Nethack* env) { env->stats.max_xp = xp; } + // in-game score delta (the eval metric itself); coef 0 adds exactly 0.0f + r += env->score_coef * (float)(env->blstats[NLE_BL_SCORE] - score_before); + // scout: pay every tile walked this step; a rush resolves many moves in // one nle_step, so drain the engine's path rather than crediting only // where the hero stopped. A mid-step level change (trapdoor, hole) @@ -788,30 +858,22 @@ static float nethack_reward(Nethack* env) { touched += (c > 0.0f); } - if (n && fresh < (float)n - 1e-6f) env->stats.scout_held++; if (fresh > 0.0f) { r += env->scout_coef * fresh; env->stats.new_tiles += touched; } - // ac: delta reward through a conservation ledger -- at most +-ac_coef - // pays per step and the remainder carries, so telescoping stays exact - // under the clamp and churn nets zero. ac_nospell is the unpaid fraction - // of protection-spell AC (1 = durable AC only; kills cast-cycle arbitrage). long ac = env->blstats[NLE_BL_AC]; env->stats.last_ac = (int)ac; if ((int)ac < env->stats.min_ac) env->stats.min_ac = (int)ac; - - float ac_led = env->ac_nospell != 0.0f - ? (float)ac + env->ac_nospell * (float)nle_spellprot(env->ctx) : (float)ac; - env->ac_account += env->ac_coef * (env->prev_ac_led - ac_led); - env->prev_ac_led = ac_led; - - float cap = env->ac_coef; - float pay = env->ac_account > cap ? cap - : (env->ac_account < -cap ? -cap : env->ac_account); - env->ac_account -= pay; - r += pay; + if (!env->obs.done) { // death-step blstats are torn down + env->stats.last_gold = env->blstats[NLE_BL_GOLD]; + env->stats.last_exp = env->blstats[NLE_BL_EXP]; + env->stats.last_xlvl = (int)env->blstats[NLE_BL_XP]; + env->stats.last_hp = (int)env->blstats[NLE_BL_HP]; + env->stats.last_hpmax = (int)env->blstats[NLE_BL_HPMAX]; + env->stats.last_depth = (int)env->blstats[NLE_BL_DEPTH]; + } return r; } @@ -1043,13 +1105,11 @@ void puf_init(Env* env, Dict* kwargs) { env->agents[0].policy = 0; init(env); env->gold_coef = dict_get(kwargs, "gold_coef"); + env->score_coef = dict_get(kwargs, "score_coef"); env->exp_coef = dict_get(kwargs, "exp_coef"); env->descent_coef = dict_get(kwargs, "descent_coef"); env->floor_coef = dict_get(kwargs, "floor_coef"); env->scout_coef = dict_get(kwargs, "scout_coef"); - env->ac_coef = dict_get(kwargs, "ac_coef"); - env->scout_ready = dict_get(kwargs, "scout_ready"); - env->ac_nospell = dict_get(kwargs, "ac_nospell"); env->xp_coef = dict_get(kwargs, "xp_coef"); env->death_penalty = dict_get(kwargs, "death_penalty"); env->mask_search20 = dict_get(kwargs, "mask_search20"); @@ -1096,10 +1156,10 @@ void puf_log(Log* log, Dict* out) { dict_set(out, "burdened_frac", log->burdened_frac); dict_set(out, "game_time", log->game_time); dict_set(out, "max_xp_level", log->max_xp_level); + dict_set(out, "gold", log->gold); + dict_set(out, "exp_points", log->exp_points); dict_set(out, "floors", log->floors); - dict_set(out, "scout_held", log->scout_held); dict_set(out, "truncated", log->truncated); - dict_set(out, "n", log->n); } // Per-(verb,head) consumption map for PPO consumed-head gating (weak symbol diff --git a/ocean/nethack/netlib.h b/ocean/nethack/netlib.h index 7dfd17b200..be347a505b 100644 --- a/ocean/nethack/netlib.h +++ b/ocean/nethack/netlib.h @@ -78,8 +78,19 @@ static const signed char nh_obj_armcat[NH_NUM_OBJECTS] = { // encumbrance percent (unclipped past 100) + raw carry capacity #define NETHACK_SPELL_SLOTS 8 #define NETHACK_OFF_EXTRA (NETHACK_OFF_BLSTATS + NLE_BLSTATS_SIZE * 4) -#define NETHACK_EXTRA_INTS (2 + NETHACK_NUM_OCLASSES + 2 + 1 + 4 * NETHACK_SPELL_SLOTS + 2 + 13 + 5 + 2) +// intrinsics: 8 player-knowable intrinsic bits (poison/fire/cold/sleep/shock +// res, telepathy, see-invis, fast) -- 7% of deaths are poison instadeath at +// full HP and the policy has no persistent representation of resistance +// typed token streams (encoder side nethack.cu); implies +// the lean token lists. +// flag, default off since the aux-heads score verdict). +// aux target count: 10 = original probe set, 26 = mega-spec (NH_LAB) +#define NETHACK_EXTRA_INTS (2 + NETHACK_NUM_OCLASSES + 2 + 1 + 4 * NETHACK_SPELL_SLOTS + 2 + 13 + 5 + 2 \ + + 1) #define NETHACK_EXTRA_ROLEOH (2 + NETHACK_NUM_OCLASSES + 2 + 1 + 4 * NETHACK_SPELL_SLOTS + 2) +#define NETHACK_EXTRA_INTRINS (NETHACK_EXTRA_ROLEOH + 20) +// threat block: [dist, bearing, adj_cnt, near3_cnt, vis_cnt, diff, speed] +#define NETHACK_EXTRA_THREAT (NETHACK_EXTRA_INTRINS + 1) #define NETHACK_EXTRA_SHOP (2 + NETHACK_NUM_OCLASSES) #define NETHACK_EXTRA_SPELL (NETHACK_EXTRA_SHOP + 2) #define NETHACK_EXTRA_WEIGHT (NETHACK_EXTRA_SPELL + 1 + 4 * NETHACK_SPELL_SLOTS) @@ -95,7 +106,40 @@ static const signed char nh_obj_armcat[NH_NUM_OBJECTS] = { // raw topline chars, null-padded; must match NH_MSG_LEN in ocean/nethack/nethack.cu #define NETHACK_OFF_MSG (NETHACK_OFF_INVTRUE + NETHACK_INV_SLOTS * 2) #define NETHACK_MSG_LEN 128 -#define NETHACK_OBS_SIZE (NETHACK_OFF_MSG + NETHACK_MSG_LEN) +// v2 split planes: remembered-terrain byte map (0 = never seen, else +// cmap_index+1) and a K-nearest visible entity list. Player-fair: terrain +// memory only records cmap glyphs the display has shown; the entity list is a +// re-indexing of the visible glyph map plus the farlook-public attitude bit. +#define NETHACK_GLYPH_CMAP_OFF 2359 +#define NETHACK_ENT_K 32 +#define NETHACK_ENT_FIELDS 6 // glyph u16 | dx i8 | dy i8 | flags u8 | pad +#define NETHACK_OFF_TERR (NETHACK_OFF_MSG + NETHACK_MSG_LEN) +#define NETHACK_OFF_ENTS (NETHACK_OFF_TERR + NH_GRID) +// v3 typed-level planes: remembered terrain (as split), remembered floor +// objects (u16 item-row+1 per cell, 0 = none seen; kept while occluded by a +// monster, cleared when the bare floor is shown), and two typed K-nearest +// token lists. Requires NH_THREAT_OBS=1 (set by the build). +#define NETHACK_GLYPH_PET_OFF 381 +#define NETHACK_GLYPH_DET_OFF 762 +#define NETHACK_GLYPH_BODY_HI (NETHACK_GLYPH_BODY_OFF + NETHACK_NUMMONS) +// item rows: objects (1906..2358) -> 1..453, bodies (1144..1524) -> 454..834 +#define NETHACK_ITEM_ROWS 840 +#define NETHACK_MON_ROWS 384 +#ifndef NETHACK_V3_K +#define NETHACK_V3_K 16 // nearest-token cap per typed list (K=32 arm: -DNETHACK_V3_K=32) +#endif +#define NETHACK_V3_MONF 8 // species u16 | dx i8 | dy i8 | flags u8 | diff u8 | speed u8 | pad +#define NETHACK_V3_ITEMF 8 // itemrow u16 | dx i8 | dy i8 | oclass u8 | flags u8 | pad | pad +#define NETHACK_OFF_OBJM (NETHACK_OFF_TERR + NH_GRID) +#define NETHACK_OFF_VMON (NETHACK_OFF_OBJM + NH_GRID * 2) +#define NETHACK_OFF_VITEM (NETHACK_OFF_VMON + NETHACK_V3_K * NETHACK_V3_MONF) +// lean token block (NH_TOK_OBS): the two typed K-nearest lists WITHOUT the +// dense memory planes (obs H2D dominates rollout; the planes are 5KB/step). +// terr_mem/obj_mem stay env-internal so item tokens keep occlusion memory. +#define NETHACK_OFF_TOKM (NETHACK_OFF_MSG + NETHACK_MSG_LEN) +#define NETHACK_OFF_TOKI (NETHACK_OFF_TOKM + NETHACK_V3_K * NETHACK_V3_MONF) +#define NETHACK_OBS_SIZE (NETHACK_OFF_MSG + NETHACK_MSG_LEN \ + + NETHACK_V3_K * (NETHACK_V3_MONF + NETHACK_V3_ITEMF)) // engine state @@ -273,7 +317,6 @@ typedef struct Log { float depth_5, depth_10, depth_15; // fraction of episodes with max_depth >= N float mines_depth; // unique Gnomish Mines floors visited (0 = never entered) float sokoban_depth; // unique Sokoban floors visited (4 = reached the top) - float scout_held; // steps where scout_ready withheld a tile claim float enhances; // #enhance presses (skill advancement claims) float floor_eats; // eats that accepted a floor "eat it?" offer float reads_scroll; @@ -286,6 +329,8 @@ typedef struct Log { float min_ac; // best (lowest) AC reached this episode float game_time; // NetHack turns survived float max_xp_level; + float gold; // net gold on the last living obs (score's netgold term) + float exp_points; // experience points on the last living obs (score's urexp term) // episode end reason (game_end_types in hack.h); other = traps/wrath/poison/... float death_combat; float death_starved; @@ -303,11 +348,14 @@ typedef struct Stats { long valid_moves; long illegal_actions; long new_tiles; - long scout_held; long enhances; long burdened_steps; int min_ac; int last_ac; // AC on the last living obs (death-step blstats are torn down) + long last_gold; // gold on the last living obs (same teardown) + long last_exp; // experience points on the last living obs + int last_xlvl; // xp level on the last living obs + int last_hp, last_hpmax, last_depth; // vitals on the last living obs long floor_eats; long reads_scroll; long reads_book; @@ -327,3 +375,69 @@ typedef struct Stats { unsigned short visited_key[NETHACK_MAX_DEPTH]; // dnum << 8 | dlevel per slot int n_visited_floors; } Stats; + +// per-species engine-static tables, generated from libnethack mons[] +// (tools/interp): difficulty, speed, hazard bits by monster number (= glyph). +#define NH_MONS_STATIC_N 381 +static const unsigned char NH_MON_DIFF[NH_MONS_STATIC_N] = { + 4, 5, 6, 6, 6, 12, 2, 6, 8, 7, 8, 8, 1, 1, 2, 4, 3, 5, 5, 7, + 6, 7, 7, 8, 9, 9, 14, 2, 3, 8, 8, 8, 3, 5, 6, 7, 7, 7, 8, 8, + 8, 11, 2, 4, 5, 6, 8, 13, 19, 3, 3, 4, 5, 7, 7, 5, 6, 8, 1, 2, + 3, 4, 4, 8, 9, 11, 5, 5, 5, 1, 3, 3, 4, 5, 5, 5, 7, 4, 6, 9, + 4, 7, 8, 9, 13, 15, 22, 1, 2, 4, 4, 4, 4, 3, 4, 7, 8, 12, 14, 4, + 6, 6, 6, 7, 9, 4, 6, 7, 9, 9, 10, 6, 9, 10, 17, 1, 9, 5, 7, 11, + 11, 12, 19, 21, 26, 2, 3, 6, 7, 6, 8, 9, 13, 13, 13, 13, 13, 13, 13, 13, + 13, 20, 20, 20, 20, 20, 20, 20, 20, 20, 9, 10, 10, 10, 10, 1, 2, 2, 2, 2, + 2, 5, 3, 4, 5, 6, 8, 8, 10, 11, 13, 13, 19, 20, 17, 18, 3, 4, 5, 6, + 14, 18, 21, 29, 4, 5, 6, 6, 7, 7, 8, 10, 4, 4, 4, 4, 8, 10, 13, 16, + 7, 9, 11, 4, 6, 8, 12, 9, 8, 14, 3, 6, 7, 8, 9, 10, 9, 12, 12, 13, + 16, 12, 12, 14, 32, 7, 8, 17, 11, 4, 6, 7, 7, 8, 9, 1, 2, 3, 3, 4, + 5, 7, 5, 9, 14, 4, 4, 6, 6, 7, 8, 10, 12, 15, 18, 22, 2, 3, 3, 6, + 12, 6, 7, 8, 11, 11, 11, 15, 14, 14, 13, 15, 30, 8, 10, 13, 12, 14, 8, 12, + 25, 34, 22, 12, 14, 11, 8, 9, 8, 10, 10, 11, 11, 12, 13, 14, 15, 16, 15, 20, + 26, 31, 36, 36, 40, 45, 53, 57, 34, 34, 34, 8, 5, 6, 9, 7, 10, 22, 1, 2, + 3, 4, 6, 7, 7, 12, 1, 12, 12, 12, 12, 12, 12, 11, 12, 12, 12, 12, 12, 12, + 12, 12, 22, 22, 22, 22, 23, 30, 30, 22, 24, 23, 22, 23, 23, 23, 22, 23, 23, 22, + 31, 23, 17, 20, 19, 19, 20, 7, 7, 7, 7, 7, 8, 8, 7, 7, 7, 7, 8, 7, + 8}; +static const unsigned char NH_MON_SPEED[NH_MONS_STATIC_N] = { + 18, 18, 18, 18, 6, 24, 3, 1, 6, 4, 6, 6, 12, 15, 12, 12, 18, 16, 16, 15, + 12, 12, 12, 12, 12, 12, 14, 3, 1, 13, 13, 13, 18, 16, 15, 15, 15, 15, 12, 12, + 10, 15, 9, 6, 9, 6, 6, 12, 12, 3, 12, 12, 3, 15, 13, 0, 0, 3, 6, 6, + 6, 6, 15, 3, 3, 3, 12, 12, 12, 6, 9, 9, 9, 5, 7, 9, 5, 1, 1, 1, + 9, 9, 18, 3, 12, 12, 12, 12, 10, 12, 12, 3, 3, 12, 4, 15, 15, 3, 3, 16, + 24, 24, 24, 20, 24, 1, 20, 20, 20, 22, 22, 3, 3, 3, 9, 12, 18, 15, 15, 8, + 10, 8, 10, 18, 16, 22, 22, 20, 20, 18, 18, 20, 9, 9, 9, 9, 9, 9, 9, 9, + 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 12, 36, 12, 6, 6, 1, 0, 0, 0, 0, + 1, 1, 6, 8, 10, 10, 6, 6, 10, 12, 12, 12, 12, 18, 15, 12, 6, 8, 10, 12, + 6, 9, 9, 9, 8, 10, 10, 10, 12, 12, 12, 14, 10, 10, 10, 10, 12, 14, 14, 16, + 10, 12, 14, 1, 3, 6, 6, 12, 18, 12, 8, 15, 15, 3, 15, 18, 12, 10, 12, 14, + 12, 6, 12, 14, 26, 12, 12, 12, 9, 12, 12, 12, 15, 12, 15, 6, 6, 6, 6, 6, + 6, 8, 6, 8, 8, 12, 12, 9, 9, 6, 3, 8, 7, 6, 6, 6, 12, 12, 12, 12, + 12, 12, 12, 12, 12, 12, 12, 18, 12, 12, 0, 12, 15, 10, 10, 6, 10, 10, 10, 10, + 12, 12, 15, 3, 10, 12, 12, 9, 12, 12, 12, 12, 12, 6, 15, 6, 9, 6, 12, 5, + 3, 18, 9, 3, 15, 9, 12, 15, 12, 12, 12, 12, 3, 12, 12, 9, 10, 3, 6, 6, + 6, 6, 6, 5, 9, 12, 0, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, + 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, + 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, + 12}; +// hazard bits per species (generated via nle_mon_haz from mons[]): +// bit0 passive-counterattack, bit1 engulf, bit2 explosive, bit3 poisonous +static const unsigned char NH_MON_HAZ[NH_MONS_STATIC_N] = { +0,8,8,0,0,8,1,0,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,4,1,4,4,4,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,0,0,8,0,1,1,3,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,8,0,0,0,0,8,8, +8,2,2,0,0,0,0,0,0,2,2,2,3,2,3,0,0,0,2,0,0,4,4,0, +8,0,0,0,0,0,0,0,8,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,0,8,0,0,2,1,0,0,0,1,1,1,1,0,0,0,0,0,0,0,0, +0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,0,0,0,0,8,0,0,0,0,0,0,0,0,0,0,0,8,8,0,8,8, +0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,8,0,0,0,0,0,0,0,0, +0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,8,0,0,0,0,0,0,0, +0,8,0,0,0,0,8,0,0,0,0,0,2,0,8,8,0,8,0,0,0,0,0,0, +8,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +0,0,0,8,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, +}; diff --git a/profile.sh b/profile.sh index 1815128538..a5a27c4e56 100755 --- a/profile.sh +++ b/profile.sh @@ -33,8 +33,8 @@ nsys profile \ --trace=cuda,nvtx \ -o "$OUT" \ ./puffer train \ - "base.profile=1" \ - "train.total_timesteps=$STEPS" \ + "--base.profile=1" \ + "--train.total_timesteps=$STEPS" \ "${EXTRA[@]}" REP="${OUT}.nsys-rep" diff --git a/resources/nethack/msgcode128.bin b/resources/nethack/msgcode128.bin new file mode 100644 index 0000000000..383986837a Binary files /dev/null and b/resources/nethack/msgcode128.bin differ diff --git a/resources/nethack/nethack_score_weights.bin b/resources/nethack/nethack_score_weights.bin index 1f4c1564c6..9aee34c9dc 100644 Binary files a/resources/nethack/nethack_score_weights.bin and b/resources/nethack/nethack_score_weights.bin differ diff --git a/scripts/wandb_sync.py b/scripts/wandb_sync.py new file mode 100755 index 0000000000..2f50d713f7 --- /dev/null +++ b/scripts/wandb_sync.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +"""Stream a puffer run's metrics jsonl (base.wandb=True) to wandb. + +Wrapper (live): wandb_sync.py --project P [--group G] [--name N] -- ./puffer train ENV ... + injects base.wandb=True base.run_id=, follows the file +Import (post): wandb_sync.py --project P --file logs/ENV/RUN.jsonl [--name N] +Sweep (live): wandb_sync.py --project P [--group G] --watch logs/ENV + one wandb run per new *.jsonl; a run finishes when its file + goes idle for --idle seconds +""" +import argparse +import configparser +import json +import os +import subprocess +import sys +import time + + +def read_config(ini_path): + cp = configparser.ConfigParser(strict=False) + try: + cp.read(ini_path) + return {f"{s}.{k}": v for s in cp.sections() for k, v in cp[s].items()} + except Exception: + return {} + + +def base_key(configs, key, default): + cp = configparser.ConfigParser(strict=False) + cp.read(configs) + return cp.get("base", key, fallback=default) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--project", required=True) + ap.add_argument("--entity", default=None) + ap.add_argument("--group", default=None) + ap.add_argument("--name", default=None) + ap.add_argument("--env", default=None, help="env name for env-compiled binaries") + ap.add_argument("--tags", default=None, help="comma-separated") + ap.add_argument("--file", default=None, help="import an existing jsonl") + ap.add_argument("--watch", default=None, help="dir: one wandb run per new jsonl") + ap.add_argument("--idle", type=float, default=0.0, + help="file mode: keep following until idle this long (0 = stop at EOF)") + ap.add_argument("cmd", nargs="*", help="-- ./puffer train ENV [overrides]") + args = ap.parse_args() + + if args.watch: + seen = {} + try: + while True: + for fn in sorted(os.listdir(args.watch)): + if fn.endswith(".jsonl") and fn not in seen: + fp = os.path.join(args.watch, fn) + cmd = [sys.executable, __file__, "--project", args.project, + "--file", fp, "--idle", str(args.idle or 180.0)] + for flag in ("entity", "group", "tags"): + v = getattr(args, flag) + if v: + cmd += [f"--{flag}", v] + print(f"wandb_sync: new trial {fn}") + seen[fn] = subprocess.Popen(cmd) + time.sleep(2.0) + except KeyboardInterrupt: + for c in seen.values(): + c.wait() + return + + child = None + if args.cmd: + env_name = args.env or next( + (a for a in args.cmd[1:] if not a.startswith("-") and "=" not in a + and a not in ("train", "eval", "sweep", "match")), None) + assert env_name, "could not find ENV in command" + name = args.name or f"{env_name}_{int(time.time())}" + log_dir = base_key( + ["config/default.ini", f"config/{env_name}.ini"], "log_dir", "logs") + path = os.path.join(log_dir, env_name, f"{name}.jsonl") + if os.path.exists(path): + os.remove(path) + child = subprocess.Popen(args.cmd + ["--base.wandb=1", f"--base.run_id={name}"]) + else: + assert args.file, "need --file or a command after --" + path = args.file + name = args.name or os.path.basename(path)[:-len(".jsonl")] + + import wandb + run = wandb.init( + project=args.project, entity=args.entity, group=args.group, name=name, + tags=args.tags.split(",") if args.tags else None) + wandb.define_metric("agent_steps") + wandb.define_metric("*", step_metric="agent_steps") + + f = None + config_sent = False + last = None + rows = 0 + last_data = time.time() + while True: + if f is None: + if os.path.exists(path): + f = open(path) + elif child is not None and child.poll() is not None: + break + else: + time.sleep(0.5) + continue + line = f.readline() + if not line: + if child is not None: + if child.poll() is not None: + break + elif args.idle <= 0 or time.time() - last_data > args.idle: + break + time.sleep(1.0) + continue + if not line.endswith("\n"): # partial write; rewind and retry + f.seek(f.tell() - len(line)) + time.sleep(0.2) + continue + if not config_sent: + cfg = read_config(path[:-len(".jsonl")] + ".ini") + if cfg: + run.config.update(cfg) + config_sent = True + try: + last = json.loads(line) + except json.JSONDecodeError: + continue + wandb.log(last) + rows += 1 + last_data = time.time() + + if last: + run.summary.update(last) + rc = child.wait() if child is not None else 0 + print(f"wandb_sync: {rows} rows -> {run.url}") + run.finish(exit_code=rc) + sys.exit(rc) + + +if __name__ == "__main__": + main() diff --git a/tests/test_nethack_cuda.cu b/tests/test_nethack_cuda.cu index ecb2151085..7c7de520fd 100644 --- a/tests/test_nethack_cuda.cu +++ b/tests/test_nethack_cuda.cu @@ -71,6 +71,26 @@ int nh_dec_pad() { return NH_DEC_PAD; } int nh_num_actions() { return NH_ACTIONS; } int nh_heads() { return NH_HEADS; } +#if 0 +void nh_get_gA(void* dst, int B) { + cudaMemcpy(dst, g_a->gA.data, (size_t)B * NH_INV * NH_GK * sizeof(float), cudaMemcpyDeviceToDevice); +} +void nh_get_gS(void* dst, int B) { + cudaMemcpy(dst, g_a->gS.data, (size_t)B * NH_INV * NH_GK * sizeof(float), cudaMemcpyDeviceToDevice); +} +#endif +void nh_get_invout(void* dst, int B) { + cudaMemcpy(dst, g_a->inv_out.data, (size_t)B * NH_INV * 16 * sizeof(float), cudaMemcpyDeviceToDevice); +} +void nh_get_sfeat(void* dst, int B) { + cudaMemcpy(dst, g_a->inv_sfeat.data, (size_t)B * NH_INV * 24 * sizeof(float), cudaMemcpyDeviceToDevice); +} +// per-entity last-layer outputs (post-activation) for dead-unit analysis +void nh_get_mvv(void* dst, int B) { cudaMemcpy(dst, g_a->mvv.data, (size_t)B * NH_INV * NH_MV * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } +void nh_get_mvm(void* dst, int B) { cudaMemcpy(dst, g_a->mvm.data, (size_t)B * NH_LABK * NH_MV * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } +void nh_get_mvi(void* dst, int B) { cudaMemcpy(dst, g_a->mvi.data, (size_t)B * NH_LABK * NH_MV * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } +void nh_get_spv(void* dst, int B) { cudaMemcpy(dst, g_a->spv.data, (size_t)B * NH_SPELL_SLOTS * NH_SPM * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } +int nh_inv_n() { return NH_INV; } int nh_labk() { return NH_LABK; } int nh_mv() { return NH_MV; } int nh_spell_slots() { return NH_SPELL_SLOTS; } int nh_spm() { return NH_SPM; } void nh_forward(void* out, void* obs, int B) { Prec in = {.data = (precision_t*)obs, .shape = {B, NH_OBS_SIZE}}; Prec r = g_enc.forward(g_w, g_a, in, 0); @@ -98,26 +118,59 @@ TENSOR_ACC(proj_w, proj_w) TENSOR_ACC(proj_b, proj_b) TENSOR_ACC(loc_w, loc_w) TENSOR_ACC(loc_b, loc_b) -TENSOR_ACC(glb1_w, glb1_w) -TENSOR_ACC(glb1_xy, glb1_xy) -TENSOR_ACC(glb1_b, glb1_b) -TENSOR_ACC(glb2_w, glb2_w) -TENSOR_ACC(glb2_b, glb2_b) +TENSOR_ACC(terr1_w, terr1_w) +TENSOR_ACC(terr1_b, terr1_b) +TENSOR_ACC(terr2_w, terr2_w) +TENSOR_ACC(terr2_b, terr2_b) +TENSOR_ACC(locc_w, locc_w) TENSOR_ACC(inv1_w, inv1_w) TENSOR_ACC(inv1_b, inv1_b) TENSOR_ACC(inv1s_w, inv1s_w) TENSOR_ACC(invt_w, invt_w) -TENSOR_ACC(inv2_w, inv2_w) -TENSOR_ACC(inv2_b, inv2_b) TENSOR_ACC(msg_w, msg_w) TENSOR_ACC(spk_w, spk_w) -TENSOR_ACC(spk2_w, spk2_w) -TENSOR_ACC(spk2_b, spk2_b) -#if NH_ID_EMBED +TENSOR_ACC(spm1_w, spm1_w) +TENSOR_ACC(spm1_b, spm1_b) +TENSOR_ACC(spm2_w, spm2_w) +TENSOR_ACC(spm2_b, spm2_b) +TENSOR_ACC(loc2_w, loc2_w) +TENSOR_ACC(loc2_b, loc2_b) TENSOR_ACC(ide_role_w, ide_role_w) TENSOR_ACC(ide_race_w, ide_race_w) TENSOR_ACC(ide_gend_w, ide_gend_w) TENSOR_ACC(ide_algn_w, ide_algn_w) +#if NH_FILM +TENSOR_ACC(film_g_w, film_g_w) +TENSOR_ACC(film_b_w, film_b_w) +#endif +#if 0 +TENSOR_ACC(gln_g, gln_g) +TENSOR_ACC(gln_b, gln_b) +TENSOR_ACC(gnv_w, gnv_w) +TENSOR_ACC(gnv_b, gnv_b) +TENSOR_ACC(gns_w, gns_w) +#endif +TENSOR_ACC(mv1_w, mv1_w) +TENSOR_ACC(mv1_b, mv1_b) +TENSOR_ACC(mv2_w, mv2_w) +TENSOR_ACC(mv2_b, mv2_b) +TENSOR_ACC(mr_w, mr_w) +TENSOR_ACC(mr_b, mr_b) +TENSOR_ACC(mm1_w, mm1_w) +TENSOR_ACC(mm1_b, mm1_b) +TENSOR_ACC(mm2_w, mm2_w) +TENSOR_ACC(mm2_b, mm2_b) +TENSOR_ACC(ir_w, ir_w) +TENSOR_ACC(ir_b, ir_b) +TENSOR_ACC(im1_w, im1_w) +TENSOR_ACC(im1_b, im1_b) +TENSOR_ACC(im2_w, im2_w) +TENSOR_ACC(im2_b, im2_b) +#if 0 +TENSOR_ACC(gws_w, gws_w) +TENSOR_ACC(gv_w, gv_w) +TENSOR_ACC(gv_b, gv_b) +TENSOR_ACC(gtau, gtau) #endif #define GRAD_ACC(name, field) \ @@ -131,28 +184,68 @@ GRAD_ACC(proj_w, proj_wgrad) GRAD_ACC(proj_b, proj_bgrad) GRAD_ACC(loc_w, loc_wgrad) GRAD_ACC(loc_b, loc_bgrad) -GRAD_ACC(glb1_w, glb1_wgrad) -GRAD_ACC(glb1_xy, glb1_xygrad) -GRAD_ACC(glb1_b, glb1_bgrad) -GRAD_ACC(glb2_w, glb2_wgrad) -GRAD_ACC(glb2_b, glb2_bgrad) +GRAD_ACC(terr1_w, terr1_wgrad) +GRAD_ACC(terr1_b, terr1_bgrad) +GRAD_ACC(terr2_w, terr2_wgrad) +GRAD_ACC(terr2_b, terr2_bgrad) +GRAD_ACC(locc_w, locc_wgrad) GRAD_ACC(inv1_w, inv1_wgrad) GRAD_ACC(inv1_b, inv1_bgrad) GRAD_ACC(inv1s_w, inv1s_wgrad) GRAD_ACC(invt_w, invt_wgrad) -GRAD_ACC(inv2_w, inv2_wgrad) -GRAD_ACC(inv2_b, inv2_bgrad) +void nh_get_locc_lut(void* dst) { cudaMemcpy(dst, nh_locc_lut_dev, NH_GLYPH_VOCAB, cudaMemcpyDeviceToHost); cudaDeviceSynchronize(); } +void nh_get_terrc_lut(void* dst) { cudaMemcpy(dst, nh_terrc_lut_dev, NH_GLYPH_VOCAB, cudaMemcpyDeviceToHost); cudaDeviceSynchronize(); } +void nh_get_terr_tf(void* dst, int B) { cudaMemcpy(dst, g_a->terr_tf.data, (size_t)B * NH_TERRF * sizeof(float), cudaMemcpyDeviceToHost); cudaDeviceSynchronize(); } +int nh_terrf() { return NH_TERRF; } +void nh_get_haz_lut(void* dst) { cudaMemcpy(dst, nh_haz_lut_dev, 381, cudaMemcpyDeviceToHost); cudaDeviceSynchronize(); } GRAD_ACC(msg_w, msg_wgrad) GRAD_ACC(spk_w, spk_wgrad) -GRAD_ACC(spk2_w, spk2_wgrad) -GRAD_ACC(spk2_b, spk2_bgrad) -#if NH_ID_EMBED +GRAD_ACC(spm1_w, spm1_wgrad) +GRAD_ACC(spm1_b, spm1_bgrad) +GRAD_ACC(spm2_w, spm2_wgrad) +GRAD_ACC(spm2_b, spm2_bgrad) +GRAD_ACC(loc2_w, loc2_wgrad) +GRAD_ACC(loc2_b, loc2_bgrad) GRAD_ACC(ide_role_w, ide_role_wgrad) GRAD_ACC(ide_race_w, ide_race_wgrad) GRAD_ACC(ide_gend_w, ide_gend_wgrad) GRAD_ACC(ide_algn_w, ide_algn_wgrad) +#if NH_FILM +GRAD_ACC(film_g_w, film_g_wgrad) +GRAD_ACC(film_b_w, film_b_wgrad) +#endif +#if 0 +GRAD_ACC(gln_g, gln_ggrad) +GRAD_ACC(gln_b, gln_bgrad) +GRAD_ACC(gnv_w, gnv_wgrad) +GRAD_ACC(gnv_b, gnv_bgrad) +GRAD_ACC(gns_w, gns_wgrad) +#endif +GRAD_ACC(mv1_w, mv1_wgrad) +GRAD_ACC(mv1_b, mv1_bgrad) +GRAD_ACC(mv2_w, mv2_wgrad) +GRAD_ACC(mv2_b, mv2_bgrad) +GRAD_ACC(mr_w, mr_wgrad) +GRAD_ACC(mr_b, mr_bgrad) +GRAD_ACC(mm1_w, mm1_wgrad) +GRAD_ACC(mm1_b, mm1_bgrad) +GRAD_ACC(mm2_w, mm2_wgrad) +GRAD_ACC(mm2_b, mm2_bgrad) +GRAD_ACC(ir_w, ir_wgrad) +GRAD_ACC(ir_b, ir_bgrad) +GRAD_ACC(im1_w, im1_wgrad) +GRAD_ACC(im1_b, im1_bgrad) +GRAD_ACC(im2_w, im2_wgrad) +GRAD_ACC(im2_b, im2_bgrad) +#if 0 +GRAD_ACC(gws_w, gws_wgrad) +GRAD_ACC(gv_w, gv_wgrad) +GRAD_ACC(gv_b, gv_bgrad) +GRAD_ACC(gtau, gtau_grad) #endif -int nh_id_embed() { return NH_ID_EMBED; } +int nh_id_embed() { return 1; } +void nh_get_eeff(void* dst) { cudaMemcpy(dst, g_a->e_eff.data, (size_t)NH_GLYPH_VOCAB * NH_EMBED_DIM * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } +void nh_get_concat(void* dst, int B) { cudaMemcpy(dst, g_a->concat.data, (size_t)B * NH_CONCAT * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } // ---- pointer decoder (fed by the encoder's inv_out keys) ---- // forward: encoder -> decoder directly (no mingru in the harness); the diff --git a/tests/test_nethack_encoder.py b/tests/test_nethack_encoder.py index e9b30b95c8..d96f4d511e 100644 --- a/tests/test_nethack_encoder.py +++ b/tests/test_nethack_encoder.py @@ -32,10 +32,134 @@ ] # NH_TEST_IDEMB=1: build and check the identity-embedding arm (NH_ID_EMBED) -IDEMB = bool(os.environ.get("NH_TEST_IDEMB")) +IDEMB = True IDE_NAMES = ["ide_role_w", "ide_race_w", "ide_gend_w", "ide_algn_w"] if IDEMB: WEIGHT_NAMES += IDE_NAMES +# NH_TEST_FILM=1: identity FiLM on the encoder output (implies IDEMB) +FILM = False +# NH_TEST_GMEAN=1: token-mean of the patch reps appended to the concat tail +GMEAN = False +# NH_TEST_LOC2=1: two-layer local branch (loc_w -> 512 hidden -> loc2_w -> 256) +LOC2 = True +# NH_TEST_TOPK=1: top-K salient patch tokens with coords appended to the concat +TOPK = False +# NH_TEST_GLBSG=1: no gradient from the global patch path into the embedding table +GLBSG = False +# NH_TEST_DIRPTR=1: direction heads get a pointer term over the 80 patch tokens +DIRPTR = False +# NH_TEST_INTRINS=1: 8 intrinsic bits appended to the blstats features +INTRINS = True +# NH_TEST_ATTNPOOL=1: softmax-attention pooled token + coords appended to the concat +ATTNPOOL = False +# NH_TEST_EFACTOR=1: E_res frozen (zero grad); NH_TEST_GSCALE=: global dE scaled +EFACTOR = False +GSCALE = 1.0 +# NH_TEST_SPLIT=1: terrain-global (own 128-row table) + entity-token branch +SPLIT = False +SPLITR = False +if SPLITR: SPLIT = True +# NH_TEST_THREAT=1: 16 nearest-hostile/threat features appended to blstats +THREAT = False +# NH_TEST_GLBPOS=1: each global channel's argmax patch (dx,dy) appended (256 dims) +GLBPOS = False +# NH_TEST_INVATTN=1: 4-query attention tail over the 55 slot reps +INVATTN = True +# NH_TEST_ENTMAX=1: entity pool = channel-wise max (attn query unused) +ENTMAX = False +# NH_TEST_V3=1: full typed-level encoder (implies SPLIT machinery for terrain) +V3 = False +if V3: + SPLIT = True; SPLITR = True; THREAT = True +# NH_TEST_LAB=1: lab arm — typed token streams (deep values, 8-head pools, +# rank) + invattn8 + aux mega-spec heads (aux MSE folded into the FD loss). +# Leave-one-out toggles: NH_TEST_LAB_TOK/AUX/IVA=0 disable a component. +LAB = True +LAB_TOK = True +LAB_AUX = False +LAB_IVA = True +# NH_TEST_V5=1: finalized v5 encoder (ENCODER_V5.md) — class-crop local (LUT), +# terrain featurizer replacing patch, hazard token bits, intrinsics, inv2 +# (wield readout + masked-sum channel replacing the inv max-pool). Aux OFF. +V5 = True +# NH_TEST_SPELL2=1: v5.1 spell fix — sum channel over RAW slot inputs + 4 +# doorstep scalars replaces the spk2 max-pool. Implies V5. +SPELL2 = True + +GMP = os.environ.get("NH_TEST_GMP") == "1" +GEN = os.environ.get("NH_TEST_GEN") == "1" +MIN = os.environ.get("NH_TEST_MIN", "1") == "1" # min is the encoder now +APANEL = "-DNH_NO_APANEL" not in os.environ.get("NH_TEST_DEFS", "") # champion default on +LINPOOL = "-DNH_ENT_LINPOOL" in os.environ.get("NH_TEST_DEFS", "") # linear last entity layer before sum|max +def act_last(x): + import torch as _t + return x if LINPOOL else _t.relu(x) +def pool_max(v, keep, dim): + import torch as _t + m = v.masked_fill(keep < .5, -1e9).max(dim=dim).values + return _t.where((keep > .5).any(dim=dim), m, _t.zeros_like(m)) if LINPOOL else _t.relu(m) +ACCOBS = os.environ.get("NH_TEST_ACC", "1") == "1" # worn rings/amulets/eyewear + armor in inventory (exercises the panels) +import re as _re0 +_ARMOR_OTYPS = np.flatnonzero(np.array([int(x) for x in _re0.findall(r"-?\d+", _re0.search( + r"nh_obj_armcat\[NH_NUM_OBJECTS\] = \{(.*?)\};", open("ocean/nethack/netlib.h").read(), _re0.S).group(1))]) >= 0) +MSGH = int(os.environ.get("NH_TEST_MSGH", "256")) # champion msg width +KT = int(os.environ.get("NH_TEST_K", "16")) # nearest-token cap (NETHACK_V3_K) +def entact(x): + import torch as _t + return _t.relu(x) +if SPELL2: + V5 = True +if V5: + LAB = True; LAB_TOK = True; LAB_IVA = True; LAB_AUX = False + INTRINS = True; LOC2 = True +if LAB and LAB_IVA: + INVATTN = True +if V5: + for _nm in ["glb1_w", "glb1_xy", "glb1_b", "glb2_w", "glb2_b", "inv2_w", "inv2_b"]: + WEIGHT_NAMES.remove(_nm) + WEIGHT_NAMES += ["terr1_w", "terr1_b", "terr2_w", "terr2_b", "locc_w"] +if GMP: + for _nm in ["isum_w", "isum_b", "iaq_w"]: + if _nm in WEIGHT_NAMES: WEIGHT_NAMES.remove(_nm) + WEIGHT_NAMES += ["gws_w", "gv_w", "gv_b", "gtau"] +if GEN: + for _nm in ["isum_w", "isum_b", "iaq_w"]: + if _nm in WEIGHT_NAMES: WEIGHT_NAMES.remove(_nm) + WEIGHT_NAMES += ["gln_g", "gln_b", "gnv_w", "gnv_b", "gns_w"] +if MIN: + for _nm in ["isum_w", "isum_b", "iaq_w"]: + if _nm in WEIGHT_NAMES: WEIGHT_NAMES.remove(_nm) + WEIGHT_NAMES += ["mv1_w", "mv1_b", "mv2_w", "mv2_b"] +if SPELL2: + WEIGHT_NAMES.remove("spk2_w"); WEIGHT_NAMES.remove("spk2_b") + WEIGHT_NAMES += ["spm1_w", "spm1_b", "spm2_w", "spm2_b"] +if SPLIT: + WEIGHT_NAMES += ["eterr_w", "tglb1_w", "tglb1_xy", "tglb1_b", "tglb2_w", "tglb2_b"] + if not V3: # mixed-ent branch is disabled under v3 (typed lists replace it) + WEIGHT_NAMES += ["ent1_w", "ent1_b", "entq_w", "entb"] +if ATTNPOOL: + WEIGHT_NAMES += ["apq_w", "apb"] +if INVATTN and not (MIN or GMP or GEN): + WEIGHT_NAMES += ["iaq_w"] +if V3: + WEIGHT_NAMES += ["emon_w", "eitem_w", "eterrc_w", "mon1_w", "mon1_b", "monq_w", "monb"] +LAB_NAMES = ((["mr_w", "mr_b", "mm1_w", "mm1_b", "mm2_w", "mm2_b", + "ir_w", "ir_b", "im1_w", "im1_b", "im2_w", "im2_b"] if MIN + else ["lm1_w", "lm1_b", "lm2_w", "lm2_b", "lma_w", "lma_b", + "li1_w", "li1_b", "li2_w", "li2_b", "lia_w", "lia_b"]) if LAB_TOK else []) \ + + (["aux_w"] if LAB_AUX else []) +if LAB: + WEIGHT_NAMES += LAB_NAMES +if DIRPTR: + WEIGHT_NAMES += ["dec_qd_w", "dec_kd_w", "dec_taud"] +_G = {} +if TOPK: + WEIGHT_NAMES += ["sal_w", "sal_b"] +if LOC2: + WEIGHT_NAMES += ["loc2_w", "loc2_b"] +if FILM: + assert IDEMB, "NH_TEST_FILM needs NH_TEST_IDEMB=1" + WEIGHT_NAMES += ["film_g_w", "film_b_w"] def build(): @@ -56,12 +180,15 @@ def build(): "-I" + os.path.join(raylib, "include"), '-DENV_HEADER="ocean/nethack/nethack.h"', "-DPUFFER_NETHACK", "-DENV_NAME=nethack", '-DPUFFER_ENV_NAME="nethack"', + ] + [d for d in os.environ.get("NH_TEST_DEFS", "").split() if d] + [ "-Xcompiler=-DPLATFORM_DESKTOP", "-Xcompiler=-fPIC", "-Xcompiler=-fopenmp", "-O2", "-L" + os.path.join(root, "vendor", "fast-nle", "build"), "-lnethack", "-Xlinker", "-rpath", "-Xlinker", os.path.join(root, "vendor", "fast-nle", "build"), - "-lcublas", "-lcusolver", "-lcurand", "-lnvidia-ml", "-lcudart", + "-L" + os.path.join(raylib, "lib"), "-lraylib", + "-Xlinker", "-rpath", "-Xlinker", os.path.join(raylib, "lib"), + "-lcublas", "-lcublasLt", "-lcusolver", "-lcurand", "-lnvidia-ml", "-lcudart", ] try: import nvidia.nccl @@ -70,9 +197,10 @@ def build(): "-L" + os.path.join(nccl, "lib"), "-lnccl"] except ImportError: cmd += ["-lnccl"] - cmd.append(f"-DNH_ID_EMBED={int(IDEMB)}") # default is 1; arm A needs explicit 0 print("building:", " ".join(cmd)) subprocess.run(cmd, check=True) + if os.environ.get("NH_TEST_BUILD_ONLY"): + print("BUILD ONLY OK"); sys.exit(0) def load(): @@ -97,6 +225,8 @@ def load(): lib.nh_dec_forward.argtypes = [VP, VP, ctypes.c_int] lib.nh_dec_backward.argtypes = [VP, VP, VP, ctypes.c_int] lib.nh_dec_keygrad.argtypes = [VP, ctypes.c_int] + if DIRPTR: + lib.nh_dec_tokgrad.argtypes = [VP, ctypes.c_int] return lib @@ -150,16 +280,60 @@ def make_obs(B, obs_size, grid, max_glyph_used): rng.integers(0, 320, size=(B, 1)), # encumbrance percent (unclipped) rng.integers(50, 1001, size=(B, 1)), # carry capacity oh, - ], axis=1).astype(np.int64).astype(np.uint32) + ] + ([rng.integers(0, 256, size=(B, 1))] if INTRINS else []) + + ([np.concatenate([rng.integers(0, 16, size=(B, 1)), # dist (15 = none) + rng.integers(0, 9, size=(B, 1)), # bearing (8 = none) + rng.integers(0, 9, size=(B, 1)), # adjacent count + rng.integers(0, 9, size=(B, 1)), # within-3 count + rng.integers(0, 30, size=(B, 1)), # visible count + rng.integers(0, 58, size=(B, 1)), # difficulty + rng.integers(0, 37, size=(B, 1))], axis=1)] if THREAT else []) + + ([np.concatenate([rng.integers(0, 2, size=(B, 1)), # aux: ds_seen + rng.integers(-39, 40, size=(B, 1)), # ds_dx + rng.integers(-10, 11, size=(B, 1)), # ds_dy + rng.integers(0, 16, size=(B, 1)), # nh_d + rng.integers(-39, 40, size=(B, 1)), # nh_dx + rng.integers(-10, 11, size=(B, 1)), # nh_dy + rng.integers(0, 30, size=(B, 1)), # nh vis + rng.integers(0, 26, size=(B, 1)), # n_items + rng.integers(0, 2, size=(B, 1)), # has_food + rng.integers(-1, 18, size=(B, 1)), # wield_class + rng.integers(-39, 40, size=(B, 1)), # nm1_dx + rng.integers(-10, 11, size=(B, 1)), # nm1_dy + np.where(rng.random((B, 1)) < 0.2, 127, + rng.integers(0, 40, size=(B, 1))), # nm1_d (127 = none) + rng.integers(-39, 40, size=(B, 1)), # nm2_dx + rng.integers(-10, 11, size=(B, 1)), # nm2_dy + np.where(rng.random((B, 1)) < 0.3, 127, + rng.integers(0, 40, size=(B, 1))), # nm2_d + rng.integers(0, 30, size=(B, 1)), # mcnt + rng.integers(0, 12, size=(B, 1)), # qNW + rng.integers(0, 12, size=(B, 1)), # qNE + rng.integers(0, 12, size=(B, 1)), # qSW + rng.integers(0, 12, size=(B, 1)), # qSE + rng.integers(-39, 40, size=(B, 1)), # ni_dx + rng.integers(-10, 11, size=(B, 1)), # ni_dy + np.where(rng.random((B, 1)) < 0.3, 127, + rng.integers(0, 40, size=(B, 1))), # ni_d + rng.integers(0, 30, size=(B, 1)), # icnt + rng.integers(0, 8, size=(B, 1))], axis=1)] if LAB and LAB_AUX else []), # lnc bits + axis=1).astype(np.int64).astype(np.uint32) + NEX = 27 + ex.shape[1] for k in range(4): - obs[:, bl_off + k::4][:, 27:104] = ((ex >> (8 * k)) & 0xFF).astype(np.float32) + obs[:, bl_off + k::4][:, 27:NEX] = ((ex >> (8 * k)) & 0xFF).astype(np.float32) # inventory entities: 55 slot glyphs int16 LE, tail padded (5976) - inv_off = bl_off + 104 * 4 + inv_off = bl_off + NEX * 4 inv = rng.integers(0, max_glyph_used, size=(B, 55)).astype(np.int32) inv[:, ::2] = rng.integers(1906, 2359, size=(B, 28)) # object glyphs: armcat coverage n_items = rng.integers(3, 12, size=B) for b in range(B): inv[b, n_items[b]:] = 5976 + if ACCOBS: # accessory otyps (rings 150-177, amulets 178-188, eyewear 207-209) in some slots + for k in range(n_items[b]): + if rng.random() < 0.25: + inv[b, k] = 1906 + int(rng.choice(np.r_[150:189, 207:210])) + elif rng.random() < 0.25: # armor otyps (any subclass in the baked armcat table) + inv[b, k] = 1906 + int(rng.choice(_ARMOR_OTYPS)) obs[:, inv_off + 0::2][:, :55] = (inv & 0xFF).astype(np.float32) obs[:, inv_off + 1::2][:, :55] = ((inv >> 8) & 0xFF).astype(np.float32) # per-slot item state @ +55*2: 8 int8 fields, incl. the -128 spe sentinel @@ -187,7 +361,72 @@ def make_obs(B, obs_size, grid, max_glyph_used): # trigram message @ msg_off: raw topline chars (null-padded). Random # lowercase words so the char-trigram bag hits many buckets. msg_off = itr_off + 55 * 2 - msg_len = obs_size - msg_off + msg_len = 128 # fixed NETHACK_MSG_LEN; the split planes sit after it + if V3: + # planes: terr (GRID) | objm (GRID u16) | vmon 16x8 | vitem 16x8 + terr_off = obs_size - (21 * 79 * 3 + 2 * KT * 8) + tb = rng.integers(0, 97, size=(B, 21 * 79)) + obs[:, terr_off:terr_off + 21 * 79] = tb.astype(np.float32) + om_off = terr_off + 21 * 79 + objm = np.where(rng.random(size=(B, 21 * 79)) < 0.03, + rng.integers(1, 835, size=(B, 21 * 79)), 0).astype(np.int64) + obs[:, om_off + 0:om_off + 2 * 21 * 79:2] = (objm & 0xFF).astype(np.float32) + obs[:, om_off + 1:om_off + 2 * 21 * 79:2] = ((objm >> 8) & 0xFF).astype(np.float32) + vm_off = om_off + 2 * 21 * 79 + vmon = np.zeros((B, KT, 8), dtype=np.int64) + vitem = np.zeros((B, KT, 8), dtype=np.int64) + for b in range(B): + for k in range(int(rng.integers(0, KT + 1))): + sp = int(rng.integers(1, 382)) + dx = int(rng.integers(-39, 40)); dy = int(rng.integers(-10, 11)) + fl = int(rng.integers(0, 16)) + vmon[b, k] = [sp & 0xFF, (sp >> 8) & 0xFF, dx & 0xFF, dy & 0xFF, fl, + int(rng.integers(0, 58)), int(rng.integers(0, 37)), 0] + for k in range(int(rng.integers(0, KT + 1))): + row = int(rng.integers(1, 835)) + dx = int(rng.integers(-39, 40)); dy = int(rng.integers(-10, 11)) + vitem[b, k] = [row & 0xFF, (row >> 8) & 0xFF, dx & 0xFF, dy & 0xFF, 0, + int(rng.integers(0, 4)), 0, 0] + obs[:, vm_off:vm_off + KT * 8] = vmon.reshape(B, -1).astype(np.float32) + obs[:, vm_off + KT * 8:vm_off + 2 * KT * 8] = vitem.reshape(B, -1).astype(np.float32) + make_obs.terr = tb; make_obs.objm = objm + make_obs.vmon = vmon; make_obs.vitem = vitem + elif SPLIT: + terr_off = obs_size - (21 * 79 + 32 * 6) + tb = rng.integers(0, 97, size=(B, 21 * 79)) + obs[:, terr_off:terr_off + 21 * 79] = tb.astype(np.float32) + eoff = terr_off + 21 * 79 + ents = np.zeros((B, 32, 6), dtype=np.int64) + for b in range(B): + ne = rng.integers(0, 33) + for k in range(ne): + g = int(rng.integers(1, 381)) if rng.random() < 0.6 else int(rng.integers(1906, 2359)) + dx = int(rng.integers(-39, 40)); dy = int(rng.integers(-10, 11)) + fl = int(rng.integers(0, 8)) + ents[b, k] = [g & 0xFF, (g >> 8) & 0xFF, dx & 0xFF, dy & 0xFF, fl, 0] + obs[:, eoff:eoff + 32 * 6] = ents.reshape(B, -1).astype(np.float32) + make_obs.terr = tb; make_obs.ents = ents + if LAB: + # lean token lists at the obs tail: vmon 16x8 | vitem 16x8 + tm_off = obs_size - 2 * KT * 8 + labm = np.zeros((B, KT, 8), dtype=np.int64) + labi = np.zeros((B, KT, 8), dtype=np.int64) + DXR, DYR = (61, 16) if V5 else (40, 11) # V5 widened: exercise the clamps + for b in range(B): + for k in range(int(rng.integers(0, KT + 1))): + sp = int(rng.integers(1, 382)) + dx = int(rng.integers(-DXR + 1, DXR)); dy = int(rng.integers(-DYR + 1, DYR)) + fl = int(rng.integers(0, 16)) + labm[b, k] = [sp & 0xFF, (sp >> 8) & 0xFF, dx & 0xFF, dy & 0xFF, fl, + int(rng.integers(0, 58)), int(rng.integers(0, 37)), 0] + for k in range(int(rng.integers(0, KT + 1))): + row = int(rng.integers(1, 835)) + dx = int(rng.integers(-DXR + 1, DXR)); dy = int(rng.integers(-DYR + 1, DYR)) + labi[b, k] = [row & 0xFF, (row >> 8) & 0xFF, dx & 0xFF, dy & 0xFF, 0, + int(rng.integers(0, 4)), 0, 0] + obs[:, tm_off:tm_off + KT * 8] = labm.reshape(B, -1).astype(np.float32) + obs[:, tm_off + KT * 8:tm_off + 2 * KT * 8] = labi.reshape(B, -1).astype(np.float32) + make_obs.labm = labm; make_obs.labi = labi msg = np.zeros((B, msg_len), dtype=np.int64) alpha = np.frombuffer(b"abcdefghijklmnopqrstuvwxyz ", dtype=np.uint8).astype(np.int64) for b in range(B): @@ -256,32 +495,62 @@ def torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, ms BL_ISLOG = np.array([0,0,0,0,0,0,0,0,0,1,0,0,0,1,0,0,0,0,0,1,1,0,0,0,0,0,0]) w = {} - E_res = w["embed_w"] = getw(lib, "embed_w", (5977, 32)) - K_w = w["ekind_w"] = getw(lib, "ekind_w", (14, 32)) - S_w = w["esub_w"] = getw(lib, "esub_w", (getattr(lib, "nh_numel_esub_w")() // 32, 32)) + # widths from the built lib, so width arms (loc/glb/inv/bl/embed) check as-is + D = lib.nh_embed_dim() + LH = lib.nh_numel_loc_b() + if V5: + P1 = 16; GH = lib.nh_numel_terr2_b(); IP = 0 + else: + P1 = lib.nh_numel_glb1_b(); GH = lib.nh_numel_glb2_b(); IP = lib.nh_numel_inv2_b() + IH = lib.nh_numel_inv1_b(); BH = lib.nh_numel_bl_b() + SK = 16; SI = lib.nh_numel_spk_w() // SK + E_res = w["embed_w"] = getw(lib, "embed_w", (5977, D)) + K_w = w["ekind_w"] = getw(lib, "ekind_w", (14, D)) + S_w = w["esub_w"] = getw(lib, "esub_w", (getattr(lib, "nh_numel_esub_w")() // D, D)) kind_map, sub_map = glyph_map() + if EFACTOR: + E_res = E_res.detach() # residual frozen E = E_res + K_w[torch.tensor(kind_map)] + S_w[torch.tensor(sub_map)] # E_eff - loc_w = w["loc_w"] = getw(lib, "loc_w", (256, CROP * CROP * 32)) - loc_b = w["loc_b"] = getw(lib, "loc_b", (256,)) - g1_w = w["glb1_w"] = getw(lib, "glb1_w", (16, PW * PH * 32)) - g1_b = w["glb1_b"] = getw(lib, "glb1_b", (16,)) - g1_xy = w["glb1_xy"] = getw(lib, "glb1_xy", (16, 2)) - g2_w = w["glb2_w"] = getw(lib, "glb2_w", (128, 16)) - g2_b = w["glb2_b"] = getw(lib, "glb2_b", (128,)) - inv1_w = w["inv1_w"] = getw(lib, "inv1_w", (16, 32)) - inv1_b = w["inv1_b"] = getw(lib, "inv1_b", (16,)) - inv1s_w = w["inv1s_w"] = getw(lib, "inv1s_w", (16, 24)) - invt_w = w["invt_w"] = getw(lib, "invt_w", (16, 32)) - inv2_w = w["inv2_w"] = getw(lib, "inv2_w", (128, 16)) - inv2_b = w["inv2_b"] = getw(lib, "inv2_b", (128,)) - bl_w = w["bl_w"] = getw(lib, "bl_w", (64, lib.nh_bl_feat())) - bl_b = w["bl_b"] = getw(lib, "bl_b", (64,)) + loc_w = w["loc_w"] = getw(lib, "loc_w", (LH, CROP * CROP * 8 + 10 * D if V3 + else CROP * CROP * 8 if V5 else CROP * CROP * D)) + loc_b = w["loc_b"] = getw(lib, "loc_b", (LH,)) + if not V5: + g1_w = w["glb1_w"] = getw(lib, "glb1_w", (P1, PW * PH * D)) + g1_b = w["glb1_b"] = getw(lib, "glb1_b", (P1,)) + g1_xy = w["glb1_xy"] = getw(lib, "glb1_xy", (P1, 2)) + g2_w = w["glb2_w"] = getw(lib, "glb2_w", (GH, P1)) + g2_b = w["glb2_b"] = getw(lib, "glb2_b", (GH,)) + else: + TH1 = lib.nh_numel_terr1_b() + t1_w = w["terr1_w"] = getw(lib, "terr1_w", (TH1, lib.nh_terrf())) + t1_b = w["terr1_b"] = getw(lib, "terr1_b", (TH1,)) + t2_w = w["terr2_w"] = getw(lib, "terr2_w", (GH, TH1)) + t2_b = w["terr2_b"] = getw(lib, "terr2_b", (GH,)) + locc = w["locc_w"] = getw(lib, "locc_w", (lib.nh_numel_locc_w() // 8, 8)) + inv1_w = w["inv1_w"] = getw(lib, "inv1_w", (IH, D)) + inv1_b = w["inv1_b"] = getw(lib, "inv1_b", (IH,)) + inv1s_w = w["inv1s_w"] = getw(lib, "inv1s_w", (IH, 24)) + invt_w = w["invt_w"] = getw(lib, "invt_w", (IH, D)) + if not V5: + inv2_w = w["inv2_w"] = getw(lib, "inv2_w", (IP, IH)) + inv2_b = w["inv2_b"] = getw(lib, "inv2_b", (IP,)) + elif not MIN: + isum_w = w["isum_w"] = getw(lib, "isum_w", (64, IH)) + isum_b = w["isum_b"] = getw(lib, "isum_b", (64,)) + bl_w = w["bl_w"] = getw(lib, "bl_w", (BH, lib.nh_bl_feat())) + bl_b = w["bl_b"] = getw(lib, "bl_b", (BH,)) proj_w = w["proj_w"] = getw(lib, "proj_w", (H, lib.nh_concat())) proj_b = w["proj_b"] = getw(lib, "proj_b", (H,)) - msg_w = w["msg_w"] = getw(lib, "msg_w", (lib.nh_numel_msg_w() // 32, 32)) - spk_w = w["spk_w"] = getw(lib, "spk_w", (16, 36)) - spk2_w = w["spk2_w"] = getw(lib, "spk2_w", (16, 16)) - spk2_b = w["spk2_b"] = getw(lib, "spk2_b", (16,)) + msg_w = w["msg_w"] = getw(lib, "msg_w", (lib.nh_numel_msg_w() // MSGH, MSGH)) + spk_w = w["spk_w"] = getw(lib, "spk_w", (SK, SI)) + if not SPELL2: + spk2_w = w["spk2_w"] = getw(lib, "spk2_w", (SK, SK)) + spk2_b = w["spk2_b"] = getw(lib, "spk2_b", (SK,)) + else: + spm1_w = w["spm1_w"] = getw(lib, "spm1_w", (64, 16)) + spm1_b = w["spm1_b"] = getw(lib, "spm1_b", (64,)) + spm2_w = w["spm2_w"] = getw(lib, "spm2_w", (64, 64)) + spm2_b = w["spm2_b"] = getw(lib, "spm2_b", (64,)) # local: crop glyph ids with pad off-map hx, hy = bl_vals[:, 0], bl_vals[:, 1] @@ -291,23 +560,123 @@ def torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, ms r, c = hy[b] - 4 + p // CROP, hx[b] - 4 + p % CROP if 0 <= r < ROWS and 0 <= c < COLS: crop_idx[b, p] = glyphs[b, r * COLS + c] - x_local = E[torch.tensor(crop_idx)].reshape(B, -1) + if V3: # v3.2 local: terrain field (81x8) + adjacency ring (9x32) + underfoot item + emon_res = w["emon_w"] = getw(lib, "emon_w", (384, D)) + eitem_res = w["eitem_w"] = getw(lib, "eitem_w", (840, D)) + eterrc = w["eterrc_w"] = getw(lib, "eterrc_w", (128, 8)) + # factorized eff tables (champion scheme, pad row hard zero) + mon_g = np.arange(384) - 1; mon_g[0] = 0 + item_g = np.zeros(840, dtype=np.int64) + item_g[1:454] = 1906 + np.arange(453); item_g[454:835] = 1144 + np.arange(381) + def eff(res, gmap): + e = res + K_w[torch.tensor(kind_map)[torch.tensor(gmap)]] \ + + S_w[torch.tensor(sub_map)[torch.tensor(gmap)]] + return torch.cat([torch.zeros(1, D, dtype=torch.float64), e[1:]], dim=0) + emon = eff(emon_res, mon_g) + eitem = eff(eitem_res, item_g) + w["_emon_eff"] = emon; w["_eitem_eff"] = eitem + ct = np.full((B, CROP * CROP), 127, dtype=np.int64) + cm = np.zeros((B, 9), dtype=np.int64) + ci = np.zeros((B,), dtype=np.int64) + for b in range(B): + for p in range(CROP * CROP): + r, c = hy[b] - 4 + p // CROP, hx[b] - 4 + p % CROP + if 0 <= r < ROWS and 0 <= c < COLS: + cell = r * COLS + c + ct[b, p] = make_obs.terr[b, cell] + rr, rc = p // CROP - 4, p % CROP - 4 + if abs(rr) <= 1 and abs(rc) <= 1: + g = glyphs[b, cell] + m = 0 + if 0 <= g < 381: m = g + 1 + elif 381 <= g < 762: m = g - 381 + 1 + elif 762 <= g < 1144: m = (g - 762) % 381 + 1 + cm[b, (rr + 1) * 3 + (rc + 1)] = m + if rr == 0 and rc == 0: + ci[b] = make_obs.objm[b, cell] + ring = torch.where(torch.tensor(cm > 0)[:, :, None], emon[torch.tensor(cm)], + torch.zeros(1, dtype=torch.float64)) + uf = torch.where(torch.tensor(ci > 0)[:, None], eitem[torch.tensor(ci)], + torch.zeros(1, dtype=torch.float64)) + x_local = torch.cat([eterrc[torch.tensor(ct)].reshape(B, -1), + ring.reshape(B, -1), uf], dim=1) + elif V5: # class crop: LUT glyph -> 9 classes, center forced to class 7 + lib.nh_get_locc_lut.argtypes = [VP] + loclut = np.empty(5977, dtype=np.uint8) + lib.nh_get_locc_lut(loclut.ctypes.data_as(VP)) + cls = loclut[crop_idx].astype(np.int64) + cls[:, (CROP * CROP) // 2] = 7 + x_local = locc[torch.tensor(cls)].reshape(B, -1) + else: + x_local = E[torch.tensor(crop_idx)].reshape(B, -1) loc = torch.relu(x_local @ loc_w.T + loc_b) - # global: per-patch flatten -> 16 -> 128, max over tokens - pat_idx = np.full((B, PX * PY, PW * PH), PAD, dtype=np.int64) - dxy = np.zeros((B, PX * PY, 2), dtype=np.float64) - for tk in range(PX * PY): - r0, c0 = (tk // PX) * PH, (tk % PX) * PW - dxy[:, tk, 0] = (c0 + 0.5 * (PW - 1) - hx) / COLS - dxy[:, tk, 1] = (r0 + 0.5 * (PH - 1) - hy) / ROWS - for pos in range(PW * PH): - r, c = r0 + pos // PW, c0 + pos % PW - if r < ROWS and c < COLS: - pat_idx[:, tk, pos] = glyphs[:, r * COLS + c] - xp = E[torch.tensor(pat_idx)].reshape(B, PX * PY, -1) - t16 = torch.relu(xp @ g1_w.T + torch.tensor(dxy) @ g1_xy.T + g1_b) - t128 = t16 @ g2_w.T - glb = torch.relu(t128.max(dim=1).values + g2_b) + if LOC2: + w["loc2_w"] = getw(lib, "loc2_w", (lib.nh_numel_loc2_b(), LH)) + w["loc2_b"] = getw(lib, "loc2_b", (lib.nh_numel_loc2_b(),)) + loc = torch.relu(loc @ w["loc2_w"].T + w["loc2_b"]) + if V5: + # terrain branch: featurize in numpy (mirrors nh_terr_feat_kernel — + # integer-exact octants), parity-check against the kernel's terr_tf, + # then the 592 -> 256 -> 128 MLP into the glb slot + lib.nh_get_terrc_lut.argtypes = [VP] + lib.nh_get_terr_tf.argtypes = [VP, ctypes.c_int] + tlut = np.empty(5977, dtype=np.uint8) + lib.nh_get_terrc_lut(tlut.ctypes.data_as(VP)) + tf = np.zeros((B, 592)) + for b in range(B): + lmd = np.full(12, 1 << 30); lmt = np.zeros(48); sec = np.zeros((8, 4, 17)) + hcell = hy[b] * COLS + hx[b] + for cell in range(ROWS * COLS): + tc = 13 if cell == hcell else int(tlut[glyphs[b, cell]]) + if tc == 255: + continue + dyc, dxc = cell // COLS - hy[b], cell % COLS - hx[b] + adx, ady = abs(dxc), abs(dyc); cheb = max(adx, ady) + if tc < 12 and cheb < lmd[tc]: + lmd[tc] = cheb + lmt[tc*4:tc*4+4] = [1.0, dxc / 78.0, dyc / 20.0, min(cheb, 30) / 30.0] + a = np.float32(np.arctan2(np.float32(dyc), np.float32(dxc))) + np.float32(3.14159265358979) + s = int(a / np.float32(0.78539816339745)) & 7 + band = 0 if cheb < 3 else 1 if cheb < 7 else 2 if cheb < 15 else 3 + sec[s, band, tc] += 1.0 + tf[b] = np.concatenate([lmt, np.log1p(sec.reshape(-1)) / np.log(1660.0)]) + tf_cuda = np.empty(B * 592, dtype=np.float32) + lib.nh_get_terr_tf(tf_cuda.ctypes.data_as(VP), B) + diff = np.abs(tf - tf_cuda.reshape(B, 592).astype(np.float64)) + lm_err = diff[:, :48].max() # landmark block: exact + sec_bad = (diff[:, 48:] > 2e-4).mean() # sectors: atan2f boundary ULPs may + # flip exact-diagonal cells between adjacent sectors — tolerate rare flips + ok_par = lm_err < 2e-4 and sec_bad < 0.02 + print(f" [{'OK ' if ok_par else 'FAIL'}] terr featurize parity lm_max={lm_err:.2e} sector_flip_frac={sec_bad:.4f}") + assert ok_par, "terrain featurize parity" + terr_h = torch.relu(torch.tensor(tf) @ t1_w.T + t1_b) + glb = torch.relu(terr_h @ t2_w.T + t2_b) + dxy = None + else: + # global: per-patch flatten -> 16 -> 128, max over tokens + pat_idx = np.full((B, PX * PY, PW * PH), PAD, dtype=np.int64) + dxy = np.zeros((B, PX * PY, 2), dtype=np.float64) + for tk in range(PX * PY): + r0, c0 = (tk // PX) * PH, (tk % PX) * PW + dxy[:, tk, 0] = (c0 + 0.5 * (PW - 1) - hx) / COLS + dxy[:, tk, 1] = (r0 + 0.5 * (PH - 1) - hy) / ROWS + for pos in range(PW * PH): + r, c = r0 + pos // PW, c0 + pos % PW + if r < ROWS and c < COLS: + pat_idx[:, tk, pos] = glyphs[:, r * COLS + c] + xp = E[torch.tensor(pat_idx)].reshape(B, PX * PY, -1) + if GLBSG: # global path reads E without gradient + xp = xp.detach() + elif GSCALE != 1.0: # global path's gradient into E scaled + xp = GSCALE * xp + (1.0 - GSCALE) * xp.detach() + t16 = torch.relu(xp @ g1_w.T + torch.tensor(dxy) @ g1_xy.T + g1_b) + if SPLITR: # old global silenced: slice zeroed, no grad through it + t16 = t16.detach() + _G['t16'] = t16; _G['dxy'] = dxy + t128 = t16 @ g2_w.T + glb = torch.relu(t128.max(dim=1).values + g2_b) + if SPLITR: + glb = torch.zeros_like(glb).detach() # blstats features f = np.zeros((B, lib.nh_bl_feat()), dtype=np.float64) j = 0 @@ -351,6 +720,24 @@ def torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, ms if not IDEMB: # identity one-hots (dead features under NH_ID_EMBED) f[:, j:j + 20] = ex_vals[:, 57:77] j += 20 + if INTRINS: + for k in range(8): + f[:, j] = (ex_vals[:, 77].astype(np.int64) >> k) & 1; j += 1 + if THREAT: + tx = 77 + (1 if INTRINS else 0) + dist = ex_vals[:, tx].astype(np.int64) + near = dist < 15 + f[:, j] = near; j += 1 + f[:, j] = np.where(near, 1.0 - dist * np.float64(np.float32(1.0 / 15.0)), 0.0); j += 1 + for k in range(8): + f[:, j] = (ex_vals[:, tx + 1].astype(np.int64) == k); j += 1 + f[:, j] = ex_vals[:, tx + 2] * 0.125; j += 1 + f[:, j] = ex_vals[:, tx + 3] * 0.125; j += 1 + f[:, j] = ex_vals[:, tx + 4] * 0.0625; j += 1 + f[:, j] = ex_vals[:, tx + 5] * np.float64(np.float32(0.04)); j += 1 + f[:, j] = np.where(near, (ex_vals[:, tx + 5].astype(np.float64) + - bl_vals[:, 18]) * np.float64(np.float32(0.1)), 0.0); j += 1 + f[:, j] = ex_vals[:, tx + 6] * np.float64(np.float32(1.0 / 24.0)); j += 1 f = np.clip(f, -1.0, 1.0) # strict clamp, mirrors the kernel fb = torch.tensor(f) blh = torch.relu(fb @ bl_w.T + bl_b) @@ -361,8 +748,12 @@ def torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, ms sf[:, :, c] = (st_vals[:, :, 0] == c) sk = st_vals[:, :, 1] != -128 sf[:, :, 4] = sk - sf[:, :, 5] = np.where(sk, st_vals[:, :, 1] * np.float64(np.float32(0.1)), 0.0) - sf[:, :, 6] = np.log1p(np.maximum(st_vals[:, :, 2], 0)) * 0.5 + if V5: # strict [-1,1]: spe clamp/7, quantity linear capped at 30 + sf[:, :, 5] = np.where(sk, np.clip(st_vals[:, :, 1], -7, 7) / 7.0, 0.0) + sf[:, :, 6] = np.minimum(np.maximum(st_vals[:, :, 2], 0), 30) / 30.0 + else: + sf[:, :, 5] = np.where(sk, st_vals[:, :, 1] * np.float64(np.float32(0.1)), 0.0) + sf[:, :, 6] = np.log1p(np.maximum(st_vals[:, :, 2], 0)) * 0.5 sf[:, :, 7] = st_vals[:, :, 3] / 3.0 sf[:, :, 8] = st_vals[:, :, 4] / 3.0 for c in range(7): @@ -377,13 +768,25 @@ def torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, ms cat = np.where((ot >= 0) & (ot < len(_tbl)), _tbl[np.clip(ot, 0, len(_tbl) - 1)], -1) for c in range(7): sf[:, :, 17 + c] = (cat == c) - xi = E[torch.tensor(inv_vals.astype(np.int64))] - # discovered-type channel: pad (5976) slots contribute hard zero - xt = E[torch.tensor(itr_vals.astype(np.int64))] - kt = torch.tensor((itr_vals != 5976).astype(np.float64))[:, :, None] + def _itemrow(v): + v = v.astype(np.int64) + row = np.zeros_like(v) + o = (v >= 1906) & (v < 2359); row[o] = v[o] - 1906 + 1 + bdy = (v >= 1144) & (v < 1525); row[bdy] = v[bdy] - 1144 + 454 + return row + if V3: + xi = w["_eitem_eff"][torch.tensor(_itemrow(inv_vals))] + xt = w["_eitem_eff"][torch.tensor(_itemrow(itr_vals))] + kt = torch.tensor((_itemrow(itr_vals) != 0).astype(np.float64))[:, :, None] + else: + xi = E[torch.tensor(inv_vals.astype(np.int64))] + # discovered-type channel: pad (5976) slots contribute hard zero + xt = E[torch.tensor(itr_vals.astype(np.int64))] + kt = torch.tensor((itr_vals != 5976).astype(np.float64))[:, :, None] invh = torch.relu(xi @ inv1_w.T + kt * (xt @ invt_w.T) + torch.tensor(sf) @ inv1s_w.T + inv1_b) # (B,55,16) - invp = torch.relu((invh @ inv2_w.T).max(dim=1).values + inv2_b) # (B,128) + if not V5: + invp = torch.relu((invh @ inv2_w.T).max(dim=1).values + inv2_b) # (B,128) # trigram message bag: hash char-trigrams (matching nh_msg_hash), sum the # embed rows, scale by 1/sqrt(count+1). Concatenated raw (no relu). def _lc(c): return c + 32 if 65 <= c <= 90 else c @@ -396,30 +799,57 @@ def _lc(c): return c + 32 if 65 <= c <= 90 else c if c0 == 0 or c1 == 0 or c2 == 0: break key = (_lc(c0) << 16) | (_lc(c1) << 8) | _lc(c2) - ids.append(((key * 2654435761) & 0xFFFFFFFF) >> (32 - 12)) + ids.append(((key * 2654435761) & 0xFFFFFFFF) >> (32 - (msg_w.shape[0].bit_length() - 1))) cnt = len(ids) - s = msg_w[torch.tensor(ids, dtype=torch.long)].sum(dim=0) if cnt else torch.zeros(32, dtype=torch.float64) + s = msg_w[torch.tensor(ids, dtype=torch.long)].sum(dim=0) if cnt else torch.zeros(MSGH, dtype=torch.float64) rows.append(s / np.sqrt(cnt + 1)) msg_sum = torch.stack(rows, dim=0) # (B, 32); grad flows to msg_w # spell-key path: per slot, key = spk_w . [e_eff(book glyph) | known, # lev/7, fail/100, know/20000]; sum-pool feeds the trunk spk_w = w["spk_w"] - spkeys = [] + spkeys, spxs, spocc = [], [], [] for s in range(8): c = 23 + 4 * s sid = torch.tensor(ex_vals[:, c].astype(np.int64)) - sg = torch.clamp(sid + 1906, max=5975) - emb = torch.where((sid > 0)[:, None], E[sg], torch.zeros_like(E[sg])) + if V3: + sg = torch.clamp(sid + 1, max=839) + emb = torch.where((sid > 0)[:, None], w["_eitem_eff"][sg], torch.zeros_like(w["_eitem_eff"][sg])) + else: + sg = torch.clamp(sid + 1906, max=5975) + emb = torch.where((sid > 0)[:, None], E[sg], torch.zeros_like(E[sg])) sc = torch.stack([ (sid > 0).double(), torch.clamp(torch.tensor(ex_vals[:, c + 1]) * 0.142857, max=1.0), torch.clamp(torch.tensor(ex_vals[:, c + 2]) * 0.01, max=1.0), torch.clamp(torch.tensor(ex_vals[:, c + 3]) * 0.00005, max=1.0), ], dim=1) - spkeys.append(torch.relu(torch.cat([emb, sc], dim=1) @ spk_w.T)) # (B,16) + xs = torch.cat([emb, sc], dim=1) # (B,36) + spxs.append(xs); spocc.append((sid > 0).double()) + spkeys.append(torch.relu(xs @ spk_w.T)) # (B,16) sk = torch.stack(spkeys, dim=1) # (B,8,16) - spool = torch.relu((sk @ spk2_w.T).max(dim=1).values + spk2_b) - parts = [loc, glb, invp, blh, fb, msg_sum, spool] + if SPELL2: + occ = torch.stack(spocc, dim=1) # (B,8) + sph1 = entact(sk @ spm1_w.T + spm1_b) + spv = act_last(sph1 @ spm2_w.T + spm2_b) # (B,8,64) + ssum = 0.25 * (spv * occ[:, :, None]).sum(dim=1) # (B,64) + smax = pool_max(spv, occ[:, :, None], 1) + # doorstep scalars (exact; empty book -> [1,0,0,1]) + sidm = ex_vals[:, 23:23+32:4].astype(np.int64) + levm = ex_vals[:, 24:24+32:4].astype(np.int64) + failm = ex_vals[:, 25:25+32:4].astype(np.int64) + knowm = ex_vals[:, 26:26+32:4].astype(np.int64) + occn = sidm > 0 + mf = np.where(occn.any(1), np.min(np.where(occn, failm, 999), 1), 100) + ml = np.where(occn.any(1), np.max(np.where(occn, levm, 0), 1), 0) + nn = occn.sum(1) + mr = np.where(occn.any(1), np.min(np.where(occn, knowm, 99999), 1), 20000) + eng = np.stack([np.clip(mf * 0.01, 0, 1), np.clip(ml / 7.0, 0, 1), + np.clip(np.minimum(nn, 8) * 0.125, 0, 1), + np.clip(mr * 0.00005, 0, 1)], 1) + spool = torch.cat([ssum, smax, torch.tensor(eng, dtype=torch.float64)], dim=1) # (B,132) + else: + spool = torch.relu((sk @ spk2_w.T).max(dim=1).values + spk2_b) + parts = [loc, glb] + ([] if V5 else [invp]) + [blh, fb, msg_sum, spool] if IDEMB: # identity embeddings: direct table rows, indices per the kernel for nm, rows, dims in [("ide_role_w", 13, 16), ("ide_race_w", 5, 8), @@ -433,8 +863,375 @@ def _lc(c): return c + 32 if 65 <= c <= 90 else c w["ide_race_w"][torch.tensor(race)], w["ide_gend_w"][torch.tensor(gend)], w["ide_algn_w"][torch.tensor(al)]] + if SPLIT: + et_w = w["eterr_w"] = getw(lib, "eterr_w", (128, 32)) + tg1_w = w["tglb1_w"] = getw(lib, "tglb1_w", (16, PW * PH * 32)) + tg1_xy = w["tglb1_xy"] = getw(lib, "tglb1_xy", (16, 2)) + tg1_b = w["tglb1_b"] = getw(lib, "tglb1_b", (16,)) + tg2_w = w["tglb2_w"] = getw(lib, "tglb2_w", (128, 16)) + tg2_b = w["tglb2_b"] = getw(lib, "tglb2_b", (128,)) + tarr = make_obs.terr + tpat = np.full((B, PX * PY, PW * PH), 127, dtype=np.int64) + for tk in range(PX * PY): + r0, c0 = (tk // PX) * PH, (tk % PX) * PW + for pos in range(PW * PH): + r, c = r0 + pos // PW, c0 + pos % PW + if r < ROWS and c < COLS: + tpat[:, tk, pos] = tarr[:, r * COLS + c] + txp = et_w[torch.tensor(tpat)].reshape(B, PX * PY, -1) + tt16 = torch.relu(txp @ tg1_w.T + torch.tensor(dxy) @ tg1_xy.T + tg1_b) + tglb = torch.relu((tt16 @ tg2_w.T).max(dim=1).values + tg2_b) + if GMEAN: # token-mean of the relu'd patch reps, after the identity tail + parts.append(t16.mean(dim=1)) + if ATTNPOOL: # softmax over tokens of beta*(qa.t16), pooled [t16 | dxy] + w["apq_w"] = getw(lib, "apq_w", (P1,)); w["apb"] = getw(lib, "apb", (8,)) + att = torch.softmax(w["apb"][0] * (t16 @ w["apq_w"]), dim=1) # (B,80) + vtok = torch.cat([t16, torch.tensor(dxy)], dim=2) # (B,80,18) + parts.append(torch.einsum('bt,btk->bk', att, vtok)) + if TOPK: # K most salient tokens (sigmoid gate), gated rep + (dx,dy); stable top-k by index + w["sal_w"] = getw(lib, "sal_w", (P1,)); w["sal_b"] = getw(lib, "sal_b", (8,)) # sal_w stored (1,P1) + sc = torch.sigmoid(t16 @ w["sal_w"] + w["sal_b"][0]) # (B,80) + K = 8; rows = [] + for b in range(B): + order = sorted(range(t16.shape[1]), key=lambda tk: (-float(sc[b, tk]), tk))[:K] + rows.append(torch.cat([torch.cat([sc[b, tk] * t16[b, tk], torch.tensor(dxy[b, tk])]) for tk in order])) + parts.append(torch.stack(rows, dim=0)) + if V3: + parts.append(tglb) + for nm, lst, ismon in [("mon", make_obs.vmon, True)]: + tbl = w["_emon_eff"] + w1 = w[nm + "1_w"] = getw(lib, nm + "1_w", (32, 40)) + b1v = w[nm + "1_b"] = getw(lib, nm + "1_b", (32,)) + qv = w[nm + "q_w"] = getw(lib, nm + "q_w", (32,)) + bv = w[nm + ("b" if ismon else "b")] = getw(lib, nm + "b", (8,)) + rows = [] + for b in range(B): + toks, valid = [], [] + cnt, ndist, ndx, ndy = 0, 1.0, 0.0, 0.0 + for k in range(16): + row = int(lst[b, k, 0]) | (int(lst[b, k, 1]) << 8) + dx = int(lst[b, k, 2]); dx = dx - 256 if dx >= 128 else dx + dy = int(lst[b, k, 3]); dy = dy - 256 if dy >= 128 else dy + f4, f5, f6 = int(lst[b, k, 4]), int(lst[b, k, 5]), int(lst[b, k, 6]) + if row > 0: + if ismon: + tail8 = [dx / 40.0, dy / 11.0, 1.0 if f4 & 1 else 0.0, + 1.0 if f4 & 4 else 0.0, min(f5 * np.float64(np.float32(0.04)), 1.0), + min(f6 * np.float64(np.float32(1.0 / 24.0)), 1.0), + 1.0 if abs(dx) <= 1 and abs(dy) <= 1 else 0.0, + 1.0 if f4 & 8 else 0.0] + else: + tail8 = [dx / 40.0, dy / 11.0, 1.0 if f5 & 1 else 0.0, + 1.0 if f5 & 2 else 0.0, 0.0, 0.0, 0.0, 0.0] + tk = torch.cat([tbl[row], torch.tensor(tail8, dtype=torch.float64)]) + valid.append(True) + hit = (f4 & 1) if ismon else 1 + if hit: + cnt += 1 + dd = max(abs(dx / 40.0) * 40.0, abs(dy / 11.0) * 11.0) * 0.125 + if dd < ndist: + ndist = dd; ndx = dx / 40.0; ndy = dy / 11.0 + else: + tk = torch.zeros(40, dtype=torch.float64); valid.append(False) + toks.append(tk) + tokm = torch.stack(toks) + h = torch.relu(tokm @ w1.T + b1v) + if ENTMAX: # channel-wise max; (dx,dy) = nearest counted entity, as data + if any(valid): + hm = h + torch.tensor([0.0 if v else -1e30 for v in valid], + dtype=torch.float64).unsqueeze(1) + pooled = hm.max(dim=0).values + else: + pooled = torch.zeros(32, dtype=torch.float64) + rows.append(torch.cat([pooled, torch.tensor([ndx, ndy], dtype=torch.float64), + torch.tensor([min(cnt * 0.125, 1.0), min(ndist, 1.0)], dtype=torch.float64)])) + else: + sc = bv[0] * (h @ qv) + mask = torch.tensor([0.0 if v else -1e30 for v in valid], dtype=torch.float64) + att = torch.softmax(sc + mask, dim=0) if any(valid) else torch.zeros(16, dtype=torch.float64) + pooled = att @ h + pdx = (att * tokm[:, 32]).sum(); pdy = (att * tokm[:, 33]).sum() + rows.append(torch.cat([pooled, pdx.reshape(1), pdy.reshape(1), + torch.tensor([min(cnt * 0.125, 1.0), min(ndist, 1.0)], dtype=torch.float64)])) + parts.append(torch.stack(rows)) + elif SPLIT: + parts.append(tglb) + e1_w = w["ent1_w"] = getw(lib, "ent1_w", (32, 40)) + e1_b = w["ent1_b"] = getw(lib, "ent1_b", (32,)) + eq_w = w["entq_w"] = getw(lib, "entq_w", (32,)) + eb = w["entb"] = getw(lib, "entb", (8,)) + ents = make_obs.ents + rows = [] + for b in range(B): + toks, valid, hostd = [], [], [] + nhost, ndist, ndx, ndy = 0, 1.0, 0.0, 0.0 + for k in range(32): + g = int(ents[b, k, 0]) | (int(ents[b, k, 1]) << 8) + dx = int(ents[b, k, 2]); dx = dx - 256 if dx >= 128 else dx + dy = int(ents[b, k, 3]); dy = dy - 256 if dy >= 128 else dy + fl = int(ents[b, k, 4]) + if g > 0: + t = torch.cat([E[min(g, 5976)], torch.tensor([dx / 40.0, dy / 11.0, + 1.0 if fl & 1 else 0.0, 1.0 if fl & 2 else 0.0, 1.0 if fl & 4 else 0.0, 0.0, 0.0, 0.0], dtype=torch.float64)]) + valid.append(True) + if (fl & 1) and not (fl & 2): + nhost += 1 + dd = max(abs(dx), abs(dy)) * 0.125 + if dd < ndist: + ndist = dd; ndx = dx / 40.0; ndy = dy / 11.0 + else: + t = torch.zeros(40, dtype=torch.float64); valid.append(False) + toks.append(t) + tokm = torch.stack(toks) # (32, 40) + h = torch.relu(tokm @ e1_w.T + e1_b) # (32, 32) + if ENTMAX: # channel-wise max over valid tokens; (dx,dy) = nearest hostile, as data + if any(valid): + hm = h + torch.tensor([0.0 if v else -1e30 for v in valid], + dtype=torch.float64).unsqueeze(1) + pooled = hm.max(dim=0).values + else: + pooled = torch.zeros(32, dtype=torch.float64) + tail = torch.cat([pooled, torch.tensor([ndx, ndy], dtype=torch.float64), + torch.tensor([min(nhost * 0.125, 1.0), min(ndist, 1.0)], dtype=torch.float64)]) + else: + sc = eb[0] * (h @ eq_w) + mask = torch.tensor([0.0 if v else -1e30 for v in valid], dtype=torch.float64) + att = torch.softmax(sc + mask, dim=0) if any(valid) else torch.zeros(32, dtype=torch.float64) + pooled = att @ h + pdx = (att * tokm[:, 32]).sum(); pdy = (att * tokm[:, 33]).sum() + tail = torch.cat([pooled, pdx.reshape(1), pdy.reshape(1), + torch.tensor([min(nhost * 0.125, 1.0), min(ndist, 1.0)], dtype=torch.float64)]) + rows.append(tail) + parts.append(torch.stack(rows)) + if GLBPOS: # argmax patch (dx,dy) per global channel; data selection, no grad + idxs = (t16 @ g2_w.T).argmax(dim=1) # (B,128) + dxyt = torch.tensor(dxy) # (B,80,2) + parts.append(dxyt[torch.arange(idxs.shape[0])[:, None], idxs].reshape(idxs.shape[0], -1)) + if MIN: # V6-min: per-slot deep MLP -> masked sum|max + pass-throughs; no attention + mv1 = w["mv1_w"] = getw(lib, "mv1_w", (64, IH)) + mv1b = w["mv1_b"] = getw(lib, "mv1_b", (64,)) + mv2 = w["mv2_w"] = getw(lib, "mv2_w", (64, 64)) + mv2b = w["mv2_b"] = getw(lib, "mv2_b", (64,)) + occb = torch.tensor((inv_vals.astype(np.int64) != 5976).astype(np.float64)) + h1 = entact(invh @ mv1.T + mv1b) + v = act_last(h1 @ mv2.T + mv2b) # (B,55,64) + sm = 0.2 * (v * occb[:, :, None]).sum(dim=1) + mx = pool_max(v, occb[:, :, None], 1) + wldb = torch.tensor(sf[:, :, 10]) * occb + qvb = torch.tensor(sf[:, :, 12]) * occb + wrnb = torch.tensor(sf[:, :, 9]) * occb + sft = torch.tensor(sf) + pw = torch.cat([(invh * wldb[:, :, None]).sum(1), (sft * wldb[:, :, None]).sum(1)], 1) + pq = torch.cat([(invh * qvb[:, :, None]).sum(1), (sft * qvb[:, :, None]).sum(1)], 1) + nw = wrnb.sum(1).clamp(min=1.0) + pworn = (invh * wrnb[:, :, None]).sum(1) / nw[:, None] + if APANEL: # 4 single-occupant accessory slots [amulet | ring A | ring B | eyewear], first owner wins + ot = inv_vals.astype(np.int64) - 1906 + worn = (sf[:, :, 9] > 0.5) & (occb.numpy() > 0.5) + isam = worn & (ot >= 178) & (ot <= 188); isey = worn & (ot >= 207) & (ot <= 209) + isr = worn & (ot >= 150) & (ot <= 177); rr = np.cumsum(isr, 1) - 1 + masks = [isam & (np.cumsum(isam, 1) == 1), isr & (rr == 0), isr & (rr == 1), isey & (np.cumsum(isey, 1) == 1)] + if os.environ.get("NH_TEST_ACC_CANARY") == "1": masks[1], masks[2] = masks[2], masks[1] # must FAIL if exercised + pacc = torch.cat([torch.cat([(invh * m[:, :, None]).sum(1), (sft * m[:, :, None]).sum(1)], 1) + for m in [torch.tensor(mm.astype(np.float64)) for mm in masks]], 1) + parts.append(torch.cat([sm, mx, pw, pq, pworn, pacc], 1)) + else: + parts.append(torch.cat([sm, mx, pw, pq, pworn], 1)) + elif GEN: # final arch: LN + 1/sqrt(d) unit (attn slices|cnt|sum|max) + pass-throughs + gln_g = w["gln_g"] = getw(lib, "gln_g", (IH,)) + gln_b = w["gln_b"] = getw(lib, "gln_b", (IH,)) + gnv_w = w["gnv_w"] = getw(lib, "gnv_w", (64, IH)) + gnv_b = w["gnv_b"] = getw(lib, "gnv_b", (64,)) + gns_w = w["gns_w"] = getw(lib, "gns_w", (8, IH)) + occb = torch.tensor((inv_vals.astype(np.int64) != 5976).astype(np.float64)) + mu = invh.mean(dim=2, keepdim=True) + var = invh.var(dim=2, unbiased=False, keepdim=True) + lnr = (invh - mu) / torch.sqrt(var + 1e-5) * gln_g + gln_b + S = (lnr @ gns_w.T) # (B,55,8) + maskb = torch.where(occb > 0, 0.0, float("-inf"))[:, :, None] + A = torch.softmax(0.25 * S + maskb, dim=1) + A = torch.where(torch.isnan(A), torch.zeros_like(A), A) + v = torch.relu(invh @ gnv_w.T + gnv_b) # (B,55,64) + Bn = invh.shape[0] + att = torch.zeros(Bn, 64, dtype=torch.float64) + cnt = torch.zeros(Bn, 64, dtype=torch.float64) + for hh in range(8): + sl = v[:, :, hh*8:(hh+1)*8] + att[:, hh*8:(hh+1)*8] = torch.einsum('bs,bsd->bd', A[:, :, hh], sl) + cnt[:, hh*8:(hh+1)*8] = 0.2 * torch.einsum('bs,bsd->bd', + torch.sigmoid(S[:, :, hh]) * occb, sl) + vm = v * occb[:, :, None] + sm = 0.2 * vm.sum(dim=1) + mx = torch.relu(v.masked_fill(occb[:, :, None] < .5, -1e9).max(dim=1).values) + wldb = torch.tensor(sf[:, :, 10]) * occb + qvb = torch.tensor(sf[:, :, 12]) * occb + wrnb = torch.tensor(sf[:, :, 9]) * occb + sft = torch.tensor(sf) + NOPASS = int(os.environ.get("NH_GEN_NOPASS", "0")) + pw = torch.cat([(invh * wldb[:, :, None]).sum(1), (sft * wldb[:, :, None]).sum(1)], 1) + if NOPASS & 1: pw = pw * 0 + pq = torch.cat([(invh * qvb[:, :, None]).sum(1), (sft * qvb[:, :, None]).sum(1)], 1) + if NOPASS & 2: pq = pq * 0 + nw = wrnb.sum(1).clamp(min=1.0) + pworn = (invh * wrnb[:, :, None]).sum(1) / nw[:, None] + if NOPASS & 4: pworn = pworn * 0 + parts.append(torch.cat([att, cnt, sm, mx, pw, pq, pworn], 1)) + elif GMP: # generalized pooling: K heads x (softmax(tau S) att | 0.2 sigmoid(S) count) + gws = w["gws_w"] = getw(lib, "gws_w", (8, IH + 24)) + gv_w = w["gv_w"] = getw(lib, "gv_w", (16, IH)) + gv_b = w["gv_b"] = getw(lib, "gv_b", (16,)) + gtau = w["gtau"] = getw(lib, "gtau", (8,)) + occb = torch.tensor((inv_vals.astype(np.int64) != 5976).astype(np.float64)) + xs = torch.cat([invh, torch.tensor(sf)], dim=2) # (B,55,40) + S = xs @ gws.T # (B,55,8) + maskb = torch.where(occb > 0, 0.0, float("-inf"))[:, :, None] + A = torch.softmax(torch.exp(gtau)[None, None, :] * S + maskb, dim=1) + A = torch.where(torch.isnan(A), torch.zeros_like(A), A) # empty inv rows + v = torch.relu(invh @ gv_w.T + gv_b) # (B,55,16) + att = torch.einsum('bsk,bsd->bkd', A, v) # (B,8,16) + sig = torch.einsum('bsk,bsd->bkd', torch.sigmoid(S) * occb[:, :, None], v) * 0.2 + parts.append(torch.cat([att, sig], dim=2).reshape(invh.shape[0], -1)) + elif INVATTN: # M-query softmax attention over the 55 post-relu slot reps + M = lib.nh_numel_iaq_w() // IH + iaq = w["iaq_w"] = getw(lib, "iaq_w", (M, IH)) + att = torch.softmax(invh @ iaq.T, dim=1) # (B,55,M) + parts.append(torch.einsum('bsm,bsk->bmk', att, invh).reshape(invh.shape[0], -1)) + if LAB and LAB_TOK: # typed streams: 48-dim tokens -> deep values -> 8-head pools + if V5: + lib.nh_get_haz_lut.argtypes = [VP] + hazlut = np.empty(381, dtype=np.uint8) + lib.nh_get_haz_lut(hazlut.ctypes.data_as(VP)) + for nm, lst, ismon in [("lm", make_obs.labm, True), ("li", make_obs.labi, False)]: + if MIN: + rn = "mr" if ismon else "ir"; mn = "mm" if ismon else "im" + mrw = w[rn + "_w"] = getw(lib, rn + "_w", (16, 48)) + mrb = w[rn + "_b"] = getw(lib, rn + "_b", (16,)) + m1w = w[mn + "1_w"] = getw(lib, mn + "1_w", (64, 16)) + m1b = w[mn + "1_b"] = getw(lib, mn + "1_b", (64,)) + m2w = w[mn + "2_w"] = getw(lib, mn + "2_w", (64, 64)) + m2b = w[mn + "2_b"] = getw(lib, mn + "2_b", (64,)) + else: + w1 = w[nm + "1_w"] = getw(lib, nm + "1_w", (64, 48)) + b1v = w[nm + "1_b"] = getw(lib, nm + "1_b", (64,)) + w2 = w[nm + "2_w"] = getw(lib, nm + "2_w", (64, 64)) + b2v = w[nm + "2_b"] = getw(lib, nm + "2_b", (64,)) + aw = w[nm + "a_w"] = getw(lib, nm + "a_w", (8, 48)) + abv = w[nm + "a_b"] = getw(lib, nm + "a_b", (8,)) + rows = [] + for b in range(B): + toks, valid = [], [] + for k in range(KT): + row = int(lst[b, k, 0]) | (int(lst[b, k, 1]) << 8) + dx = int(lst[b, k, 2]); dx = dx - 256 if dx >= 128 else dx + dy = int(lst[b, k, 3]); dy = dy - 256 if dy >= 128 else dy + f4, f5, f6 = int(lst[b, k, 4]), int(lst[b, k, 5]), int(lst[b, k, 6]) + if row > 0: + g = row - 1 if ismon else (1906 + row - 1 if row < 454 else 1144 + row - 454) + cheb = max(abs(dx), abs(dy)) + if V5: # strict [-1,1]: rare geometry tails clamped + base = [np.clip(dx / 40.0, -1.0, 1.0), np.clip(dy / 11.0, -1.0, 1.0), + min(cheb, 15) * np.float64(np.float32(1.0 / 15.0)), + k * np.float64(np.float32(1.0 / 15.0))] + else: + base = [dx / 40.0, dy / 11.0, + min(cheb, 15) * np.float64(np.float32(1.0 / 15.0)), + k * np.float64(np.float32(1.0 / 15.0))] + if ismon: + fl = [1.0 if f4 & 1 else 0.0, 1.0 if f4 & 8 else 0.0, + 1.0 if f4 & 4 else 0.0, 1.0 if cheb <= 1 else 0.0, + min(f5 * np.float64(np.float32(0.04)), 1.0), + min(f6 * np.float64(np.float32(1.0 / 24.0)), 1.0)] + else: + fl = [1.0 if f5 & 1 else 0.0, 1.0 if f5 & 2 else 0.0, + 0.0, 0.0, 0.0, 0.0] + if V5 and ismon: # hazard LUT bits at dims 42-45 + hb = int(hazlut[(row - 1) % 381]) + tail6 = [float((hb >> j) & 1) for j in range(4)] + [0.0, 0.0] + else: + tail6 = [0.0] * 6 + tk = torch.cat([E[g], torch.tensor(base + fl + tail6, + dtype=torch.float64)]) + valid.append(True) + else: + tk = torch.zeros(48, dtype=torch.float64); valid.append(False) + toks.append(tk) + tokm = torch.stack(toks) # (16, 48) + if MIN: + vb = torch.tensor([1.0 if v else 0.0 for v in valid], + dtype=torch.float64) + rp = torch.relu(tokm @ mrw.T + mrb) # (16, 16) + hv1 = entact(rp @ m1w.T + m1b) + hv = act_last(hv1 @ m2w.T + m2b) + sm = 0.25 * (hv * vb[:, None]).sum(0) + mx = pool_max(hv, vb[:, None], 0) + g0 = rp[0] * vb[0] + if ismon: + rows.append(torch.cat([sm, mx, g0])) + else: + uf = (rp * (vb * (tokm[:, 36] > 0.5).double())[:, None]).sum(0) + rows.append(torch.cat([sm, mx, g0, uf])) + continue + h2 = torch.relu(torch.relu(tokm @ w1.T + b1v) @ w2.T + b2v) # (16, 64) + sc = tokm @ aw.T + abv # (16, 8) + mask = torch.tensor([0.0 if v else -1e30 for v in valid], + dtype=torch.float64)[:, None] + if any(valid): + att = torch.softmax(sc + mask, dim=0) # (16, 8) + pooled = torch.cat([att[:, hh] @ h2[:, hh * 8:(hh + 1) * 8] + for hh in range(8)]) + else: + pooled = torch.zeros(64, dtype=torch.float64) + rows.append(pooled) + parts.append(torch.stack(rows)) + if V5 and not GMP and not GEN and not MIN: # inv2 tail: parameterless wield readout + masked-sum channel + wldg = torch.tensor(sf[:, :, 10]) # wielded state bit + parts.append((invh * wldg[:, :, None]).sum(dim=1)) # (B,16) + ih = torch.relu(invh @ isum_w.T + isum_b) # (B,55,64) + occ = torch.tensor((inv_vals.astype(np.int64) != 5976).astype(np.float64)) + parts.append(0.2 * (ih * occ[:, :, None]).sum(dim=1)) # (B,64) concat = torch.cat(parts, dim=1) + if LAB and LAB_AUX: # aux heads: linear preds + targets/masks mirroring the kernel + aux_w = w["aux_w"] = getw(lib, "aux_w", (32, lib.nh_concat())) + ap = concat @ aux_w.T # (B, 32) + ax0 = ex_vals.shape[1] - 26 + av = ex_vals[:, ax0:].astype(np.int64) + AT = np.zeros((B, 32)); AM = np.zeros((B, 32)) + dsm = (av[:, 0] != 0).astype(np.float64) + AT[:, 0] = av[:, 1] / 20.0; AM[:, 0] = dsm + AT[:, 1] = av[:, 2] / 10.0; AM[:, 1] = dsm + nhm = (av[:, 6] > 0).astype(np.float64) + AT[:, 2] = av[:, 4] / 20.0; AM[:, 2] = nhm + AT[:, 3] = av[:, 5] / 10.0; AM[:, 3] = nhm + AT[:, 4] = av[:, 3] / 15.0; AM[:, 4] = nhm + AT[:, 5] = np.minimum(av[:, 6], 16) / 16.0; AM[:, 5] = 1 + m1m = (av[:, 12] < 127).astype(np.float64) + AT[:, 6] = av[:, 10] / 20.0; AM[:, 6] = m1m + AT[:, 7] = av[:, 11] / 10.0; AM[:, 7] = m1m + AT[:, 8] = np.minimum(av[:, 12], 15) / 15.0; AM[:, 8] = m1m + m2m = (av[:, 15] < 127).astype(np.float64) + AT[:, 9] = av[:, 13] / 20.0; AM[:, 9] = m2m + AT[:, 10] = av[:, 14] / 10.0; AM[:, 10] = m2m + AT[:, 11] = np.minimum(av[:, 15], 15) / 15.0; AM[:, 11] = m2m + AT[:, 12] = np.minimum(av[:, 16], 16) / 16.0; AM[:, 12] = 1 + for qq in range(4): + AT[:, 13 + qq] = np.minimum(av[:, 17 + qq], 8) * 0.125; AM[:, 13 + qq] = 1 + i1m = (av[:, 23] < 127).astype(np.float64) + AT[:, 17] = av[:, 21] / 20.0; AM[:, 17] = i1m + AT[:, 18] = av[:, 22] / 10.0; AM[:, 18] = i1m + AT[:, 19] = np.minimum(av[:, 23], 15) / 15.0; AM[:, 19] = i1m + AT[:, 20] = np.minimum(av[:, 24], 16) / 16.0; AM[:, 20] = 1 + AT[:, 21] = np.minimum(av[:, 7], 20) / 20.0; AM[:, 21] = 1 + AT[:, 22] = (av[:, 8] != 0); AM[:, 22] = 1 + AT[:, 23] = (av[:, 9] >= 0); AM[:, 23] = 1 + AT[:, 24] = ((av[:, 25] & 1) != 0); AM[:, 24] = 1 + AT[:, 25] = ((av[:, 25] & 4) != 0); AM[:, 25] = 1 + w["_aux"] = (ap, torch.tensor(AT), torch.tensor(AM)) out = torch.relu(concat @ proj_w.T + proj_b) + if FILM: + ide = concat[:, -40:] # the identity tail nh_idemb_kernel wrote + w["film_g_w"] = getw(lib, "film_g_w", (40, H)) + w["film_b_w"] = getw(lib, "film_b_w", (40, H)) + out = out * (1 + ide @ w["film_g_w"]) + ide @ w["film_b_w"] return out, invh, w, torch.stack(spkeys, dim=1) @@ -447,14 +1244,28 @@ def run(lib): print(f"obs_size={obs_size} grid={grid} vocab={vocab} " f"bl_feat={lib.nh_bl_feat()} concat={lib.nh_concat()} dec_od={lib.nh_dec_od()}") - # glb1_xy zero-inits; randomize it so a broken dx,dy forward term is visible - wxy = np.random.default_rng(5).standard_normal(lib.nh_numel_glb1_xy()).astype(np.float32) - lib.nh_set_glb1_xy(wxy.ctypes.data_as(VP)) + if not V5: + # glb1_xy zero-inits; randomize it so a broken dx,dy forward term is visible + wxy = np.random.default_rng(5).standard_normal(lib.nh_numel_glb1_xy()).astype(np.float32) + lib.nh_set_glb1_xy(wxy.ctypes.data_as(VP)) if IDEMB: # same idiom: zero-init tables would hide forward bugs for i, nm in enumerate(IDE_NAMES): wv = np.random.default_rng(6 + i).standard_normal( getattr(lib, f"nh_numel_{nm}")()).astype(np.float32) getattr(lib, f"nh_set_{nm}")(wv.ctypes.data_as(VP)) + if FILM: # zero-init scale/shift tables would make FiLM an invisible no-op + for i, nm in enumerate(["film_g_w", "film_b_w"]): + wv = (0.3 * np.random.default_rng(12 + i).standard_normal( + getattr(lib, f"nh_numel_{nm}")())).astype(np.float32) + getattr(lib, f"nh_set_{nm}")(wv.ctypes.data_as(VP)) + if LAB and LAB_AUX: # zero-init aux heads would leave the concat injection untested + lib.nh_auxh.restype = ctypes.c_int + lib.nh_aux_coef.restype = ctypes.c_float + for fn in ["nh_get_auxp", "nh_get_auxt", "nh_get_auxm"]: + getattr(lib, fn).argtypes = [VP, ctypes.c_int] + wv = (0.1 * np.random.default_rng(21).standard_normal( + lib.nh_numel_aux_w())).astype(np.float32) + lib.nh_set_aux_w(wv.ctypes.data_as(VP)) max_glyph_used = 40 # keep embedding usage dense & checkable obs, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg = make_obs(B, obs_size, grid, max_glyph_used) @@ -468,18 +1279,56 @@ def run(lib): def forward_loss(): lib.nh_forward(out_d, obs_d, B) out = d2h(out_d, B * hidden).reshape(B, hidden) - return float((out * g_out).sum()) + L = float((out * g_out).sum()) + if LAB and LAB_AUX: # total loss includes the aux MSE the backward differentiates + AH = lib.nh_auxh() + bufs = [] + for fn in ["nh_get_auxp", "nh_get_auxt", "nh_get_auxm"]: + d = dev(B * AH * 4) + getattr(lib, fn)(d, B) + bufs.append(d2h(d, B * AH).reshape(B, AH).astype(np.float64)) + p, t, m = bufs + L += 0.5 * float(lib.nh_aux_coef()) / B * float((m * (p - t) ** 2).sum()) + return L # Analytic grads: forward then backward with dL/dout = g_out. L0 = forward_loss() grad_d, _ = h2d(g_out) # backward mutates grad in place lib.nh_backward(grad_d, B) - enc_names = ["ekind_w", "esub_w", "proj_w", "proj_b", "bl_w", "bl_b", "loc_w", "loc_b", - "glb1_w", "glb1_xy", "glb1_b", "glb2_w", "glb2_b", - "inv1_w", "inv1_b", "inv1s_w", "invt_w", "inv2_w", "inv2_b", "embed_w", "msg_w"] + enc_names = ["ekind_w", "esub_w", "proj_w", "proj_b", "bl_w", "bl_b", "loc_w", "loc_b"] \ + + (["terr1_w", "terr1_b", "terr2_w", "terr2_b", "locc_w"] if V5 + else ["glb1_w", "glb1_xy", "glb1_b", "glb2_w", "glb2_b"]) \ + + ["inv1_w", "inv1_b", "inv1s_w", "invt_w"] \ + + ([] if (GMP or GEN or MIN) else (["isum_w", "isum_b"] if V5 else ["inv2_w", "inv2_b"])) \ + + ["embed_w", "msg_w"] if IDEMB: enc_names += IDE_NAMES + if FILM: + enc_names += ["film_g_w", "film_b_w"] + if LOC2: + enc_names += ["loc2_w", "loc2_b"] + if TOPK: + enc_names += ["sal_w", "sal_b"] + if ATTNPOOL: + enc_names += ["apq_w", "apb"] + if SPLIT: + enc_names += ["eterr_w", "tglb2_w", "tglb2_b"] + if not V3: + enc_names += ["ent1_w", "ent1_b", "entq_w", "entb"] + if V3: + enc_names += ["emon_w", "eitem_w", "eterrc_w", "mon1_w", "mon1_b", "monq_w", "monb"] + if INVATTN: + enc_names += (["mv1_w", "mv1_b", "mv2_w", "mv2_b"] if MIN + else ["gws_w", "gv_w", "gv_b", "gtau"] if GMP + else (["gln_g", "gln_b", "gnv_w", "gnv_b", "gns_w"] if GEN else ["iaq_w"])) + if LAB: + enc_names += LAB_NAMES + enc_names += ["spk_w"] if "spk_w" not in enc_names else [] + if SPELL2: + enc_names += ["spm1_w", "spm1_b", "spm2_w", "spm2_b"] + else: + enc_names += ["spk2_w", "spk2_b"] # Central finite differences of L = sum(out*g_out). The encoder ends in a # ReLU (and the blstats branch has its own), so a perturbation that flips a @@ -495,7 +1344,8 @@ def forward_loss(): # perturbations flip near-tied argmax winners under the kink detector's # radar and bias the quotient (worse at the 16-dim inv bottleneck, where # ties are denser). The exact float64 torch reference covers them. - fd_skip = {"glb1_w", "glb1_xy", "glb1_b", "inv2_w", "inv1_w", "inv1_b", "invt_w"} + # v5 removes both max-pools (patch, inv) — everything is FD-checkable + fd_skip = set() if V5 else {"glb1_w", "glb1_xy", "glb1_b", "inv2_w", "inv1_w", "inv1_b", "invt_w"} all_ok = True for name in enc_names: if name in fd_skip: @@ -535,7 +1385,8 @@ def forward_loss(): rel = abs(gnum - ga[i]) / max(1.0, abs(gnum), abs(ga[i])) max_rel = max(max_rel, rel) checked += 1 - ok = checked >= 3 and max_rel < rel_tol + # padded tensors (e.g. sal_b: 1 live of 8) can't yield 3 checks; require what exists + ok = checked >= min(3, len(cand)) and checked >= 1 and max_rel < rel_tol all_ok = all_ok and ok print(f" [{'OK ' if ok else 'FAIL'}] {name:8s} n={n:8d} " f"checked={checked} kink_skipped={skipped} " @@ -553,14 +1404,21 @@ def forward_loss(): def torch_check(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg, g_out, enc_names, H): import torch out, _, w, _ = torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg, H) - (out * torch.tensor(g_out.astype(np.float64))).sum().backward() + L = (out * torch.tensor(g_out.astype(np.float64))).sum() + if LAB and LAB_AUX: # fold the aux MSE in so aux_w and the concat injection are checked + ap, at, am = w["_aux"] + lib.nh_aux_coef.restype = ctypes.c_float + coef = float(lib.nh_aux_coef()) + L = L + 0.5 * coef / out.shape[0] * (am * (ap - at) ** 2).sum() + L.backward() ok = True for name in enc_names: n = getattr(lib, f"nh_numel_{name}")() ga = np.empty(n, dtype=np.float32) getattr(lib, f"nh_grad_{name}")(ga.ctypes.data_as(VP)) - gt = w[name].grad.numpy().reshape(-1) + gt = (w[name].grad.numpy().reshape(-1) if w[name].grad is not None + else np.zeros(w[name].numel())) # frozen/detached weight: reference grad is zero denom = max(1.0, np.abs(gt).max()) rel = np.abs(ga - gt).max() / denom good = rel < 1e-3 @@ -591,7 +1449,8 @@ def dec_check(lib, obs_d, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, di_d = dev(B * H * 4) lib.nh_dec_backward(gl_d, gv_d, di_d, B) di_cuda = d2h(di_d, B * H).reshape(B, H) - kg_cuda = np.empty(B * 55 * 16, dtype=np.float32) + IH = lib.nh_numel_inv1_b(); SK = 16 + kg_cuda = np.empty(B * 55 * IH, dtype=np.float32) lib.nh_dec_keygrad(kg_cuda.ctypes.data_as(VP), B) # torch replica: hidden state and keys detached so grads are decoder-local, @@ -602,23 +1461,48 @@ def dec_check(lib, obs_d, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, s_k = invh.detach().clone().requires_grad_(True) sp_k = spkeys.detach().clone().requires_grad_(True) # (B,8,16) lin_w = getw(lib, "dec_lin_w", (PAD, H)) - q_w = getw(lib, "dec_q_w", ((HEADS + 1) * 16, H)) - k_w = getw(lib, "dec_k_w", (16, 16)) + q_w = getw(lib, "dec_q_w", ((HEADS + 1) * IH, H)) + k_w = getw(lib, "dec_k_w", (IH, IH)) tau = getw(lib, "dec_tau", (lib.nh_numel_dec_tau(),)) # padded; first HEADS live tmp = h_in @ lin_w.T # (B,PAD), rows N_ACT+48+1 used - qall = (h_in @ q_w.T).reshape(B, HEADS + 1, 16) + qall = (h_in @ q_w.T).reshape(B, HEADS + 1, IH) q = qall[:, :HEADS] kmat = s_k @ k_w.T # (B,55,16) qn = q.norm(dim=2) + 1e-6 kn = kmat.norm(dim=2) + 1e-6 - cos = torch.einsum('bhk,bik->bhi', q, kmat) / (qn[:, :, None] * kn[:, None, :]) + cos = torch.einsum('bhk,bik->bhi', q, kmat) / qn[:, :, None] slot = torch.exp(tau[:HEADS])[None, :, None] * cos # (B,HEADS,55) log-tau # spell head: dot(q_spell, key_s) / 4 (dot-product pointer, no tau) - spell = torch.einsum('bk,bsk->bs', qall[:, HEADS], sp_k) * 0.25 + spell = torch.einsum('bk,bsk->bs', qall[:, HEADS, :SK], sp_k) * 0.25 # spell query uses the first SK dims N_DIRS_LIN = 48 + dirlin = tmp[:, N_ACT:N_ACT+N_DIRS_LIN] + if DIRPTR: + t_k = _G['t16'].detach().clone().requires_grad_(True) # (B,80,16) token keys, leaf + dxy = _G['dxy'] + qd_w = getw(lib, "dec_qd_w", (6 * 16, H)); kd_w = getw(lib, "dec_kd_w", (16, 16)); taud = getw(lib, "dec_taud", (8,)) + qd = (h_in @ qd_w.T).reshape(B, 6, 16); kd = t_k @ kd_w.T + qdn = qd.norm(dim=2) + 1e-6; kdn = kd.norm(dim=2) + 1e-6 + cosd = torch.einsum('bhk,btk->bht', qd, kd) / (qdn[:, :, None] * kdn[:, None, :]) # (B,6,80) + # octant per token, mirroring nh_tok_octant (E=0..., mapped to N,S,W,E,NW,NE,SW,SE) + octs = np.full((B, 80), -1) + mp = [3, 7, 1, 6, 2, 4, 0, 5] + for b in range(B): + for tk in range(80): + cx, cy = dxy[b, tk, 0] * 79, dxy[b, tk, 1] * 21 + if abs(cx) < 2.6 and abs(cy) < 2.6: continue + a = np.degrees(np.arctan2(cy, cx)); a = a + 360 if a < 0 else a + octs[b, tk] = mp[int(np.floor((a + 22.5) / 45.0)) % 8] + extra = torch.zeros(B, 6, 8, dtype=torch.float64) + for b in range(B): + for h in range(6): + for d in range(8): + idx = [tk for tk in range(80) if octs[b, tk] == d] + if idx: + extra[b, h, d] = torch.exp(taud[h]) * cosd[b, h, idx].max() + dirlin = dirlin + extra.reshape(B, 48) out = torch.cat([tmp[:, :N_ACT], slot.reshape(B, HEADS * 55), - tmp[:, N_ACT:N_ACT+N_DIRS_LIN], spell, + dirlin, spell, tmp[:, N_ACT+N_DIRS_LIN:N_ACT+N_DIRS_LIN+1]], dim=1) (out * torch.tensor(g.astype(np.float64))).sum().backward() @@ -628,8 +1512,10 @@ def dec_check(lib, obs_d, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, ok = ok and good print(f" [{'OK ' if good else 'FAIL'}] torch dec_out max_rel_err={rel:.2e}") - for name, ref in [("dec_lin_w", lin_w), ("dec_q_w", q_w), - ("dec_k_w", k_w), ("dec_tau", tau)]: + dec_pairs = [("dec_lin_w", lin_w), ("dec_q_w", q_w), ("dec_k_w", k_w), ("dec_tau", tau)] + if DIRPTR: + dec_pairs += [("dec_qd_w", qd_w), ("dec_kd_w", kd_w), ("dec_taud", taud)] + for name, ref in dec_pairs: n = getattr(lib, f"nh_numel_{name}")() ga = np.empty(n, dtype=np.float32) getattr(lib, f"nh_grad_{name}")(ga.ctypes.data_as(VP)) @@ -640,7 +1526,11 @@ def dec_check(lib, obs_d, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, ok = ok and good print(f" [{'OK ' if good else 'FAIL'}] torch {name:9s} max_rel_err={rel:.2e}") - for name, cuda_g, ref in [("keygrad", kg_cuda, s_k), ("grad_input", di_cuda, h_in)]: + extra_pairs = [("keygrad", kg_cuda, s_k), ("grad_input", di_cuda, h_in)] + if DIRPTR: + tg_cuda = np.empty(B * 80 * 16, dtype=np.float32); lib.nh_dec_tokgrad(tg_cuda.ctypes.data_as(VP), B) + extra_pairs.append(("tokgrad", tg_cuda, t_k)) + for name, cuda_g, ref in extra_pairs: gt = ref.grad.numpy().reshape(-1) denom = max(1.0, np.abs(gt).max()) rel = np.abs(cuda_g.reshape(-1) - gt).max() / denom diff --git a/vendor/nh_web_shell.html b/vendor/nh_web_shell.html new file mode 100644 index 0000000000..5b66080a49 --- /dev/null +++ b/vendor/nh_web_shell.html @@ -0,0 +1,73 @@ + + + + +NetHack — PufferLib + + + + +
loading engine + weights…
+ +{{{ SCRIPT }}} + +