diff --git a/config/default.ini b/config/default.ini index 62d82194e1..b94499c24b 100644 --- a/config/default.ini +++ b/config/default.ini @@ -107,6 +107,11 @@ vtrace_c_clip = 1.0 prio_alpha = 0.8 prio_beta0 = 0.2 +# Verb-eps exploration floor (experimental, Nethack; see src/algo.cu). 0 = off. +verb_eps = 0.0 +verb_eps_anneal_start = 0.4 +verb_eps_anneal_end = 1.0 + [sweep] metric = score metric_distribution = linear diff --git a/config/nethack.ini b/config/nethack.ini index 3c85a29b08..a4593e3faf 100644 --- a/config/nethack.ini +++ b/config/nethack.ini @@ -1,55 +1,63 @@ [base] env_name = nethack reset_every_horizon = 0 -async = 1 +async = 0 [vec] -total_agents = 256 +total_agents = 512 num_buffers = 4 -#num_threads = 8 num_threads = 32 -action_mask_size = 690 +action_mask_size = 734 [env] -gold_coef = 1 -exp_coef = 0.0222526733 -descent_coef = 4.57267237 -xp_coef = 0.648124218 -scout_coef = 0.1 -hp_coef = 0.00617194036 -hunger_coef = 0.01 -illegal_penalty = -0.01 -death_penalty = -0.112676442 +gold_coef = 0.021659979757862095 +exp_coef = 0.08341832601491733 +descent_coef = 1.0 +xp_coef = 0.5926087261188399 +scout_coef = 0.025 ac_coef = 1.0 -heal_coef = 0.001 -status_coef = 0.0116204191 +death_penalty = -0.8209825710553373 + +# flags +# scout reward scales with xp level relative to depth; <= 0 disables +scout_ready = 1.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 [policy] hidden_size = 512 num_layers = 3 [train] -#total_timesteps = 166275056 -total_timesteps = 4_000_000_000 -minibatch_size = 16384 -horizon = 512 -learning_rate = 0.00175767997 -gamma = 0.993521929 -gae_lambda = 0.611256123 -replay_ratio = 2.42882919 -clip_coef = 0.478926629 -vf_coef = 2.29309034 -vf_clip_coef = 3.19970345 -max_grad_norm = 1.43588173 -ent_coef = 0.0192690529 -beta1 = 0.9694917707305044 -beta2 = 0.9475724955192064 -eps = 0.0001 -prio_alpha = 0.931300879 -prio_beta0 = 0.885229468 +total_timesteps = 247765094 +minibatch_size = 32768 +horizon = 64 +learning_rate = 0.011854702184553156 +momentum = 0.805399553172738 +verb_eps = 0.01 +verb_eps_anneal_start = 0.4 +verb_eps_anneal_end = 1.0 +gamma = 0.9973714278257158 +gae_lambda = 0.7339946788354172 +replay_ratio = 1.036717410123433 +clip_coef = 1.0 +vf_coef = 1.8486010633248169 +vf_clip_coef = 4.624607731543118 +max_grad_norm = 0.902165688849037 +ent_coef = 0.008389210353641681 +beta1 = 0.805399553172738 +beta2 = 0.999139492311758 +eps = 1e-14 +vtrace_rho_clip = 2.770648671516198 +vtrace_c_clip = 4.5026502564349595 +prio_alpha = 0.6754855924309153 +prio_beta0 = 0.8672463883417074 [sweep] -metric = max_depth +metric = score [sweep.train.total_timesteps] distribution = log_normal @@ -57,74 +65,50 @@ min = 1.5e8 max = 4e8 scale = time -[sweep.env.gold_coef] +[sweep.env.exp_coef] distribution = log_normal -min = 0.001 -max = 1.0 +min = 0.01 +max = 0.1 scale = auto -[sweep.env.exp_coef] +[sweep.env.gold_coef] distribution = log_normal -min = 0.01 -max = 10.0 +min = 0.001 +max = 0.1 scale = auto [sweep.env.descent_coef] -distribution = log_normal -min = 0.5 -max = 10.0 +distribution = uniform +min = 0.0 +max = 1.0 scale = auto [sweep.env.xp_coef] -distribution = log_normal -min = 0.05 +distribution = uniform +min = 0.0 max = 1.0 scale = auto [sweep.env.scout_coef] distribution = log_normal -min = 0.0001 +min = 0.001 max = 0.1 scale = auto -[sweep.env.hp_coef] -distribution = log_normal -min = 0.0005 -max = 0.06 -scale = auto - -[sweep.env.hunger_coef] -distribution = log_normal -min = 0.01 -max = 0.25 -scale = auto - -[sweep.env.illegal_penalty] -distribution = uniform -min = -0.01 -max = 0.0 -scale = auto - [sweep.env.death_penalty] distribution = uniform min = -1.0 max = 0.0 scale = auto -[sweep.env.heal_coef] -distribution = log_normal -min = 0.001 -max = 0.1 -scale = auto - -[sweep.env.status_coef] -distribution = log_normal -min = 0.01 -max = 1.0 -scale = auto - -[sweep.env.ac_coef] -distribution = log_normal -min = 0.01 -max = 0.5 -scale = auto +#[sweep.train.verb_eps] +#distribution = log_normal +#min = 0.0005 +#max = 0.02 +#scale = auto + +#[sweep.train.verb_eps_anneal_start] +#distribution = uniform +#min = 0.2 +#max = 0.8 +#scale = auto diff --git a/ocean/nethack/README.md b/ocean/nethack/README.md index 69e9a15427..03b6e1b78d 100644 --- a/ocean/nethack/README.md +++ b/ocean/nethack/README.md @@ -1,8 +1,8 @@ # NetHack PufferLib environment for NetHack 3.6.6 over -[fast-nle](https://github.com/FinlaySanders/fast-nle): 22-verb factored -action space (verb, item slot, direction), legality masking, +[fast-nle](https://github.com/FinlaySanders/fast-nle): 26-verb factored +action space (verb, item slot, direction, spell slot), legality masking, decomposed-score reward, custom CUDA encoder/decoder (`ocean/nethack/nethack.cu`). ## Setup @@ -39,11 +39,8 @@ Interactive controls (default when stdin is a TTY): | `Shift+Space` | hold advances at **20 Hz** (fallback: hold `S`) | | `q` / `Esc` | quit | -Weight resolution: `NH_WEIGHTS` if set, else the highest-step checkpoint in -the **most recently written** run under `checkpoints/nethack/`, else -`resources/nethack/nethack_weights.bin`. (Global max-step is wrong after a -role switch — a mature monk 2B run outranks a fresh valk 200M run.) Set -`NH_SEED` to replay a seed. +Weight resolution: `NH_WEIGHTS` if set, else the checked-in +`resources/nethack/nethack_weights.bin`. Set `NH_SEED` to replay a seed. ```bash ./nethack 10000 0 # headless 10k steps (prints avg_score / avg_max_depth) diff --git a/ocean/nethack/depth.ini b/ocean/nethack/depth.ini new file mode 100644 index 0000000000..b69c701b31 --- /dev/null +++ b/ocean/nethack/depth.ini @@ -0,0 +1,64 @@ +# async depth config: sweep-derived depth-seeking rewards (descent-dominated, +# soft death penalty) on the async trainer. To use: copy to config/depth.ini, +# then ./puffer train depth +[base] +env_name = nethack +reset_every_horizon = 0 +async = 1 + +[vec] +total_agents = 256 +num_buffers = 4 +num_threads = 32 +action_mask_size = 734 + +[env] +gold_coef = 1.0 +exp_coef = 0.0222526733 +descent_coef = 4.57267237 +xp_coef = 0.648124218 +scout_coef = 0.1 +ac_coef = 0.0 +death_penalty = -0.112676442 + +# 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 + +[policy] +hidden_size = 512 +num_layers = 3 + +[train] +total_timesteps = 247765094 +minibatch_size = 16384 +horizon = 512 +learning_rate = 0.00175767997 +anneal_lr = 1 +min_lr_ratio = 0.0 +momentum = 0.9694917707305044 +verb_eps = 0.005 +verb_eps_anneal_start = 0.4 +verb_eps_anneal_end = 1.0 +gamma = 0.993521929 +gae_lambda = 0.611256123 +replay_ratio = 2.42882919 +clip_coef = 0.478926629 +vf_coef = 2.29309034 +vf_clip_coef = 3.19970345 +max_grad_norm = 1.43588173 +ent_coef = 0.0192690529 +anneal_ent_coef = 0 +min_ent_coef_ratio = 0.1 +beta1 = 0.9694917707305044 +beta2 = 0.9475724955192064 +eps = 0.0001 +vtrace_rho_clip = 1.0 +vtrace_c_clip = 1.0 +prio_alpha = 0.931300879 +prio_beta0 = 0.885229468 diff --git a/ocean/nethack/fs.h b/ocean/nethack/fs.h index 843ebfd211..8c42819438 100644 --- a/ocean/nethack/fs.h +++ b/ocean/nethack/fs.h @@ -19,7 +19,7 @@ static void nethack_touch(const char* path) { } static void nethack_rm_rf(const char* path, int depth) { - if (depth > 3) return; // vardir trees are at most base/env/save/files + if (depth > 3) return; // vardir trees are at most base/env/save/files DIR* d = opendir(path); if (d) { struct dirent* e; @@ -82,7 +82,7 @@ static int nethack_make_vardir(const char* source_hackdir, char* out_buf, size_t // fail fast on a dangling nhdat symlink: symlink(2) would succeed and the // error surface later as a cryptic init_dungeons panic - char resolved[4096]; // realpath(3) requires a PATH_MAX buffer + char resolved[4096]; // realpath(3) requires a PATH_MAX buffer if (realpath(src, resolved) == NULL || access(resolved, R_OK) != 0) { fprintf(stderr, "nethack: NETHACKDIR misconfigured — no readable nhdat at %s (%s).\n" diff --git a/ocean/nethack/macros.h b/ocean/nethack/macros.h index 82db8e046e..4e46eb4d5c 100644 --- a/ocean/nethack/macros.h +++ b/ocean/nethack/macros.h @@ -34,15 +34,15 @@ static int nethack_parse_candidates(const Nethack* env, char* cand, int cap) { int n = 0; for (i++; i < NLE_MESSAGE_SIZE && m[i] && n < cap; i++) { unsigned char c = m[i]; - if (n == 0 && (c == '-' || c == ' ' || c == '$')) continue; // leading "- " (allownone) / "$" (gold) - if (c == '-' && i + 1 < NLE_MESSAGE_SIZE) { // compactified run + if (n == 0 && (c == '-' || c == ' ' || c == '$')) continue; // leading "- " (allownone) / "$" (gold) + if (c == '-' && i + 1 < NLE_MESSAGE_SIZE) { // compactified run for (char x = cand[n-1] + 1; x <= (char)m[i+1] && n < cap; x++) cand[n++] = x; i++; continue; } if ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')) cand[n++] = (char)c; - else break; // ' ' before "or ?*", ']', '#', ...: end of the letter list + else break; // ' ' before "or ?*", ']', '#', ...: end of the letter list } return n; } @@ -57,7 +57,9 @@ static void nethack_drain_prompts(Nethack* env) { if (nethack_msg_contains(env, "more confident in your")) env->enh_ready = 1; int yn = env->misc[NETHACK_MISC_YN]; if (!yn && !env->misc[NETHACK_MISC_GETLIN] && !env->misc[NETHACK_MISC_XWAIT]) break; - env->obs.action = yn ? 27 : '\r'; + // alternate return/ESC: getpos-class pickers eat '\r' silently and + // only ESC exits them (the t-frozen wedge class) + env->obs.action = yn ? 27 : (i & 1 ? 27 : '\r'); env->ctx = nle_step(env->ctx, &env->obs); } } @@ -87,12 +89,16 @@ static int nethack_handle_prompts(Nethack* env) { if (!yn && !env->misc[NETHACK_MISC_GETLIN] && !env->misc[NETHACK_MISC_XWAIT] && !nethack_msg_is_prompt(env)) break; int ring = yn && nethack_msg_contains(env, "ight or Left"); - // commit 'y' ONLY to prompts rendering y/n choices ("[yn"): getobj also - // polls through yn_function but wants an item LETTER — an auto-'y' - // there reads as slot y and re-prompts forever (zero-turn loop) + // commit 'y' ONLY to rendered y/n choices ("[yn"): getobj also polls + // through yn_function but wants a LETTER — auto-'y' loops forever int commit = yn && nethack_msg_contains(env, "[yn") && !nethack_msg_contains(env, "no return") && !nethack_msg_contains(env, "eally attack"); + // shopkeeper's " offers N gold pieces for your X. Sell it?" is + // ynaq-rendered, so commit already accepts it — count the conversion + if (commit && nethack_msg_contains(env, "gold piece") + && nethack_msg_contains(env, "Sell")) + env->stats.sells++; env->obs.action = ring ? 'r' : (commit ? 'y' : 27); env->ctx = nle_step(env->ctx, &env->obs); } @@ -144,10 +150,16 @@ static int nethack_item_use(Nethack* env, int cmd, const char* gate, int n = nethack_parse_candidates(env, cand, (int)sizeof(cand)); char want = (char)env->inv_letters[slot]; int ok = 0; - for (int j = 0; j < n; j++) - if (cand[j] == want) { ok = 1; break; } + for (int j = 0; j < n; j++) { + if (cand[j] != want) continue; + ok = 1; + break; + } nethack_send_key(env, ok ? want : 27); - if (!ok) { if (bad_pick) *bad_pick = 1; return 0; } + if (!ok) { + if (bad_pick) *bad_pick = 1; + return 0; + } if (stat) (*stat)++; return 1; } @@ -168,7 +180,6 @@ static void nethack_wear_takeoff_conflict(Nethack* env, int slot) { int cat_i = (gi >= 0 && gi < NH_NUM_OBJECTS) ? nh_obj_armcat[gi] : -1; if (cat_i == cat_new && i != slot) { nethack_item_use(env, 'T', "take off", NULL, i, NULL, NULL); - env->stats.armor_swaps++; return; } } @@ -210,6 +221,39 @@ static void nethack_do_elbereth(Nethack* env) { } } +// engrave-test the first unidentified wand (E, wand letter, ride the +// dialogue): fire/lightning/digging formally identify (+10 score), the rest +// print their tell. Throwaway text 'x'; same prompt loop as Elbereth. +static void nethack_do_engrave_id(Nethack* env) { + int slot = -1; + for (int i = 0; i < NETHACK_INV_SLOTS && env->inv_letters[i]; i++) { + if (env->inv_oclasses[i] != 11 + || env->inv_state[i * NLE_INV_STATE_FIELDS + 6] != 0) continue; + int lb = nethack_letter_bit(env->inv_letters[i]); + if (lb >= 0 && (env->engid_tested & (1ULL << lb))) continue; + slot = i; + break; + } + if (slot < 0) return; + int lb = nethack_letter_bit(env->inv_letters[slot]); + if (lb >= 0) env->engid_tested |= 1ULL << lb; + env->obs.action = 'E'; + env->ctx = nle_step(env->ctx, &env->obs); + if (env->obs.done || !env->misc[NETHACK_MISC_YN] + || !nethack_msg_contains(env, "write with")) return; + env->obs.action = env->inv_letters[slot]; + env->ctx = nle_step(env->ctx, &env->obs); + const char* c = "x\r"; + for (int i = 0; i < NETHACK_AUTODISMISS_MAX && !env->obs.done && *c; i++) { + if (env->misc[NETHACK_MISC_XWAIT]) env->obs.action = ' '; + else if (env->misc[NETHACK_MISC_YN] + && nethack_msg_contains(env, "current engraving")) env->obs.action = 'n'; + else if (env->misc[NETHACK_MISC_GETLIN]) env->obs.action = (unsigned char)*c++; + else break; + env->ctx = nle_step(env->ctx, &env->obs); + } +} + // claim a banked skill advance the moment its notice appears (zero-turn, // unconditionally good): #enhance, then 'a' = the first advanceable skill static void nethack_auto_enhance(Nethack* env) { @@ -224,9 +268,8 @@ static void nethack_auto_enhance(Nethack* env) { if (!env->obs.done && nethack_msg_contains(env, "# ")) { for (const char* c = "enhance\r"; !env->obs.done && *c; c++) nethack_send_key(env, (unsigned char)*c); - if (!env->obs.done && env->misc[NETHACK_MISC_XWAIT]) { + if (!env->obs.done && env->misc[NETHACK_MISC_XWAIT]) nethack_send_key(env, 'a'); - } } nethack_drain_prompts(env); if (!env->obs.done) nle_obs_refresh(env->ctx, &env->obs); diff --git a/ocean/nethack/nethack.c b/ocean/nethack/nethack.c index 25af5e7ea3..151ef7b72e 100644 --- a/ocean/nethack/nethack.c +++ b/ocean/nethack/nethack.c @@ -3,9 +3,6 @@ #include #include #include -#include -#include -#include #include #include "nethack.h" #include "../../src/puffercpu.h" @@ -14,56 +11,71 @@ // single-agent env, reset immediately (training's puf_reset is lazy) static void env_open(Nethack* env) { memset(env, 0, sizeof(*env)); + // dungeon variety: rng feeds init()'s seed; srand() runs before env_open + // in both demo modes, so NH_SEED replays exactly and no-seed varies by time + env->rng = (unsigned)rand(); env->num_agents = 1; - env->observations = (unsigned char*)calloc(NETHACK_OBS_SIZE, 1); - env->actions = (float*)calloc(14, sizeof(float)); // {verb, 12 per-verb slots, direction} - env->action_mask = (unsigned char*)calloc(NETHACK_NUM_ACTIONS - + 12 * NETHACK_INV_SLOTS + NETHACK_NUM_DIRS, 1); - env->rewards = (float*)calloc(1, sizeof(float)); - env->terminals = (float*)calloc(1, sizeof(float)); + Agent* a = &env->agents[0]; + a->observations = (unsigned char*)calloc(NETHACK_OBS_SIZE, 1); + a->actions = (float*)calloc(20, sizeof(float)); // {verb, 12 per-verb slots, 6 per-verb dirs, spell slot} + a->action_mask = (unsigned char*)calloc(NETHACK_NUM_ACTIONS + + 12 * NETHACK_INV_SLOTS + NETHACK_DIR_HEADS * NETHACK_NUM_DIRS + + NETHACK_SPELL_SLOTS, 1); + a->rewards = (float*)calloc(1, sizeof(float)); + a->terminals = (float*)calloc(1, sizeof(float)); init(env); + nethack_sync_buffers(env); // flat mask pointer, written by compute_mask nethack_do_reset(env); } static void env_close(Nethack* env) { puf_close(env); - free(env->observations); free(env->actions); free(env->rewards); free(env->terminals); - free(env->action_mask); + Agent* a = &env->agents[0]; + free(a->observations); + free(a->actions); + free(a->rewards); + free(a->terminals); + free(a->action_mask); } // CPU port of the CUDA encoder (ocean/nethack/nethack.cu) + puffernet MinGRU/decoder; // weight order matches param registration: encoder, decoder, mingru -#define DEMO_VOCAB 5977 -#define DEMO_EMBED 32 -#define DEMO_BL_FEAT (25 + 7 + 13 + NETHACK_NUM_ACTIONS + NETHACK_NUM_OCLASSES + 2 + 8 + 2) -#define DEMO_INV_HID 16 // 16-dim slot rep: pool bottleneck + decoder key (unified) +#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) +#define DEMO_SPKEY 16 +#define DEMO_SPIN (DEMO_EMBED + 4) +#define DEMO_INV_HID 16 // 16-dim slot rep: pool bottleneck + decoder key (unified) #define DEMO_INV_FLAT (NETHACK_INV_SLOTS * DEMO_INV_HID) #define DEMO_INV_POOL 128 -#define DEMO_SFEAT 24 // buc4 + known+spe + quan + ero2 + flags7 + tk + armcat7 -#define DEMO_OD (NETHACK_NUM_ACTIONS + 12 * NETHACK_INV_SLOTS + NETHACK_NUM_DIRS) -#define DEMO_NUM_HEADS 14 +#define DEMO_SFEAT 24 // buc4 + known+spe + quan + ero2 + flags7 + tk + armcat7 +#define DEMO_OD (NETHACK_NUM_ACTIONS + 12 * NETHACK_INV_SLOTS + NETHACK_DIR_HEADS * NETHACK_NUM_DIRS \ + + NETHACK_SPELL_SLOTS) +#define DEMO_NUM_HEADS 20 #define DEMO_PTR_HEADS 12 -#define DEMO_QDIM (DEMO_PTR_HEADS * DEMO_INV_HID) -#define DEMO_DEC_PAD 32 -#define DEMO_DEC_LIN (NETHACK_NUM_ACTIONS + NETHACK_NUM_DIRS + 1) -#define DEMO_LOC_IN (NETHACK_CROP_GRID * DEMO_EMBED) // 9x9 crop, per-cell embeds +#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_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_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 // trigram message branch, mirroring NH_MSG_* in ocean/nethack/nethack.cu -#define DEMO_MSG_LEN NETHACK_MSG_LEN +#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_HID 32 #define DEMO_MSG_CONCAT_OFF (DEMO_LOC_HID + DEMO_GLB_HID + DEMO_INV_POOL + 64 + DEMO_BL_FEAT) -#define DEMO_CONCAT (DEMO_MSG_CONCAT_OFF + DEMO_MSG_HID) +#define DEMO_SPELL_CONCAT_OFF (DEMO_MSG_CONCAT_OFF + DEMO_MSG_HID) +#define DEMO_CONCAT (DEMO_SPELL_CONCAT_OFF + DEMO_SPKEY) // per-blstat normalization, mirroring NH_BL_SCALE / NH_BL_ISLOG in ocean/nethack/nethack.cu static const float DEMO_BL_SCALE[27] = { @@ -71,39 +83,43 @@ static const float DEMO_BL_SCALE[27] = { 1.f/25, 1.f/125, 1.f/25, 1.f/25, 1.f/25, 1.f/25, 1.f/25, 0.1f, 1.f/200, 1.f/200, 1.f/50, 0.1f, 1.f/100, 1.f/100, 1.f/10, 1.f/10, 1.f/30, - 0.1f, 0.1f, 0.f, 1.f/4, 0.f, 1.f/50, 0.f, 1.f, // dnum one-hot (scale dead) + 0.1f, 0.1f, 0.f, 1.f/4, 0.f, 1.f/50, 0.f, 1.f, // dnum one-hot (scale dead) }; static const int DEMO_BL_ISLOG[27] = {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}; 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 *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 *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 *inv2_w, *inv2_b; // (128, 16), (128): pooled trunk summary (max over slots) - 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 *dec_lin; // (32, H) bias-free; rows [22 verb | 8 dir | value], 31 used - float *dec_q; // (192, H): twelve stacked 16-dim query projections - float *dec_k; // (16, 16): key projection over slot features - float *dec_tau; // (12,): per-head log cosine temperature + float *g2_w, *g2_b; // (128, 16), (128): 16 -> 128, maxed over tokens + 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 *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 *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) + float *dec_k; // (16, 16): key projection over slot features + float *dec_tau; // (12,): per-head log cosine temperature MinGRU* mingru; Multidiscrete* md; int hidden_size, num_layers, num_actions; - float x[DEMO_LOC_IN]; // crop cell embeds, flattened - float px[DEMO_GLB_IN]; // one patch's cell embeds, flattened + 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 slots[DEMO_INV_FLAT]; // per-slot post-relu features (decoder keys) - float concat[DEMO_CONCAT]; // [local hid | global hid | inv pool | bl hidden | bl feats | msg] - float logits[DEMO_OD + 1]; // assembled decoder output; last entry is value - float* hidden; // (hidden_size) + 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] + float logits[DEMO_OD + 1]; // assembled decoder output; last entry is value + float* hidden; // (hidden_size) } NethackNet; // (hidden, layers) from the checkpoint float count: @@ -116,10 +132,12 @@ typedef struct { + DEMO_GLB_HID*DEMO_P1 + DEMO_GLB_HID \ + 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_MSG_VOCAB*DEMO_MSG_HID) -#define DEMO_DEC_FIXED (DEMO_INV_HID*DEMO_INV_HID + 16) // k_w + tau padded 12->16 + + DEMO_MSG_VOCAB*DEMO_MSG_HID \ + + DEMO_SPKEY*DEMO_SPIN + DEMO_SPKEY*DEMO_SPKEY + DEMO_SPKEY) +#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 best_l = 1 << 30; @@ -130,7 +148,12 @@ static int demo_infer_arch(int total, int* hidden, int* layers, int* actions) { if (rem <= 0) break; if (rem % per_layer) continue; long L = rem / per_layer; - if (L >= 1 && L < best_l) { best_l = (int)L; *hidden = H; *layers = (int)L; *actions = NETHACK_NUM_ACTIONS; } + if (L >= 1 && L < best_l) { + best_l = (int)L; + *hidden = H; + *layers = (int)L; + *actions = NETHACK_NUM_ACTIONS; + } } return best_l == 1 << 30 ? -1 : 0; } @@ -146,36 +169,41 @@ static NethackNet* make_nethack_net(Weights* w) { 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->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->inv1_w = get_weights_aligned(w, DEMO_INV_HID * DEMO_EMBED); - net->inv1_b = get_weights_aligned(w, DEMO_INV_HID); + 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->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->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 * 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->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 * 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); 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); + 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); net->dec_tau = get_weights_aligned(w, DEMO_PTR_HEADS); - net->mingru = make_mingru(w, 1, net->hidden_size, net->num_layers); + net->mingru = make_mingru(w, 1, net->hidden_size, net->num_layers); static int logit_sizes[DEMO_NUM_HEADS] = { NETHACK_NUM_ACTIONS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, - NETHACK_INV_SLOTS, NETHACK_NUM_DIRS}; + NETHACK_INV_SLOTS, NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, + NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, NETHACK_SPELL_SLOTS}; net->md = make_multidiscrete(1, logit_sizes, DEMO_NUM_HEADS); assert(w->idx == w->size - 7); // materialize the residual-factorized embedding once (host, load time) @@ -189,7 +217,7 @@ static NethackNet* make_nethack_net(Weights* w) { } static inline int demo_msg_lc(int c) { - return (c >= 'A' && c <= 'Z') ? c + 32 : c; // lowercase; keep spaces/punct + return (c >= 'A' && c <= 'Z') ? c + 32 : c; // lowercase; keep spaces/punct } static inline int demo_msg_hash(int c0, int c1, int c2) { unsigned key = ((unsigned)c0 << 16) | ((unsigned)c1 << 8) | (unsigned)c2; @@ -211,6 +239,7 @@ static void demo_msg_pool(NethackNet* net, const unsigned char* obs, float* out) } float scale = 1.0f / sqrtf((float)count + 1.0f); for (int d = 0; d < DEMO_MSG_HID; d++) out[d] *= scale; + } // blstats/extra live at unaligned byte offsets: assemble, don't cast @@ -228,7 +257,7 @@ static int demo_glyph_at(const int16_t* glyphs, int r, int c) { return g; } -static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // fills decoder->output +static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // fills decoder->output const int16_t* glyphs = (const int16_t*)(obs + NETHACK_OFF_GLYPHS); const unsigned char* bl = obs + NETHACK_OFF_BLSTATS; @@ -273,6 +302,7 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // // as the pointer decoder's keys), then 32 -> 128 with max over slots for // the trunk (matches the CUDA fused pool) const int16_t* inv = (const int16_t*)(obs + NETHACK_OFF_INV); + const int16_t* invt = (const int16_t*)(obs + NETHACK_OFF_INVTRUE); const signed char* invst = (const signed char*)(obs + NETHACK_OFF_INVST); for (int slot = 0; slot < NETHACK_INV_SLOTS; slot++) { int g = inv[slot]; @@ -289,7 +319,7 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // sf[8] = (float)st[4] * (1.0f / 3.0f); for (int c = 0; c < 7; c++) sf[9 + c] = (float)((st[5] >> c) & 1); sf[16] = (float)st[6]; - int ot = inv[slot] - NH_GLYPH_OBJ_OFF; // armor slot category one-hot + 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; float* h32 = net->slots + slot * DEMO_INV_HID; @@ -298,6 +328,12 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // for (int k = 0; k < DEMO_INV_HID; k++) for (int j = 0; j < DEMO_SFEAT; j++) h32[k] += net->inv1s_w[k * DEMO_SFEAT + j] * sf[j]; + int gt = invt[slot]; // discovered-type channel; pad = unknown + if (gt >= 0 && gt < DEMO_VOCAB - 1) + for (int k = 0; k < DEMO_INV_HID; k++) + for (int d = 0; d < DEMO_EMBED; d++) + h32[k] += net->invt_w[k * DEMO_EMBED + d] + * net->e_eff[gt * DEMO_EMBED + d]; _relu(h32, h32, DEMO_INV_HID); } float* invp = net->concat + DEMO_LOC_HID + DEMO_GLB_HID; @@ -317,7 +353,7 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // float* f = net->concat + DEMO_LOC_HID + DEMO_GLB_HID + DEMO_INV_POOL + 64; int j = 0; for (int i = 0; i < 27; i++) { - if (i == 21 || i == 25) continue; // hunger, condition: expanded below + if (i == 21 || i == 25) continue; // hunger, condition: expanded below float v = (float)demo_i32(bl + 4*i); f[j++] = DEMO_BL_ISLOG[i] ? log1pf(fmaxf(v, 0.f)) * DEMO_BL_SCALE[i] : v * DEMO_BL_SCALE[i]; @@ -329,7 +365,7 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // const unsigned char* ex = obs + NETHACK_OFF_EXTRA; for (int h = 0; h < NETHACK_NUM_ACTIONS; h++) f[j++] = (h == demo_i32(ex + 4)) ? 1.f : 0.f; for (int k = 0; k < NETHACK_NUM_OCLASSES; k++) f[j++] = (float)demo_i32(ex + 4*(2 + k)) * 0.125f; - for (int p = 0; p < 2; p++) { // hp_frac, ene_frac + for (int p = 0; p < 2; p++) { // hp_frac, ene_frac int cur = demo_i32(bl + 4*(p ? 14 : 10)), mx = demo_i32(bl + 4*(p ? 15 : 11)); f[j++] = fminf(fmaxf((float)cur / (float)(mx > 1 ? mx : 1), 0.f), 1.f); } @@ -337,8 +373,16 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // int dnum = d23 < 0 ? 0 : (d23 > 7 ? 7 : d23); for (int d = 0; d < 8; d++) f[j++] = (d == dnum) ? 1.f : 0.f; int engr = demo_i32(ex); - f[j++] = engr >= 1 ? 1.f : 0.f; // any engraving underfoot - f[j++] = engr >= 2 ? 1.f : 0.f; // active Elbereth + f[j++] = engr >= 1 ? 1.f : 0.f; // any engraving underfoot + f[j++] = engr >= 2 ? 1.f : 0.f; // active Elbereth + f[j++] = (float)demo_i32(ex + 4*NETHACK_EXTRA_SHOP); // in shop + f[j++] = (float)demo_i32(ex + 4*(NETHACK_EXTRA_SHOP+1)) * 0.01f; // affordability + // spell scalar; mirrors NH_F_SPELL: known count/8 only + f[j++] = (float)demo_i32(ex + 4*NETHACK_EXTRA_SPELL) * 0.125f; + { // encumbrance pair; mirrors NH_F_WEIGHT in nethack.cu + 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 < 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; @@ -347,6 +391,43 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // demo_msg_pool(net, obs, net->concat + DEMO_MSG_CONCAT_OFF); + { // spell-key path; mirrors nh_spkey_kernel + nh_sppool_kernel + 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); + float in[DEMO_SPIN]; + if (sid > 0) { + int g = sid + 1906; if (g > 5975) g = 5975; + for (int d = 0; d < DEMO_EMBED; d++) in[d] = net->e_eff[g * DEMO_EMBED + d]; + } else + for (int d = 0; d < DEMO_EMBED; d++) in[d] = 0.f; + in[DEMO_EMBED + 0] = sid > 0 ? 1.f : 0.f; + float lv = (float)demo_i32(q + 4) * 0.142857f; + float fl = (float)demo_i32(q + 8) * 0.01f; + float kn = (float)demo_i32(q + 12) * 0.00005f; + in[DEMO_EMBED + 1] = lv > 1.f ? 1.f : lv; + in[DEMO_EMBED + 2] = fl > 1.f ? 1.f : fl; + in[DEMO_EMBED + 3] = kn > 1.f ? 1.f : kn; + for (int r = 0; r < DEMO_SPKEY; r++) { + float acc = 0.f; + for (int c = 0; c < DEMO_SPIN; c++) acc += net->spk_w[r * DEMO_SPIN + c] * in[c]; + net->spkeys[s * DEMO_SPKEY + r] = acc > 0.f ? acc : 0.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; + } + float v = best + net->spk2_b[d]; + sp[d] = v > 0.f ? v : 0.f; + } + } _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); @@ -395,13 +476,23 @@ static int nethack_net_forward(NethackNet* net, const unsigned char* obs) { // expf(net->dec_tau[h]) * dot / (nq * kn[i]); } } - for (int d = 0; d <= NETHACK_NUM_DIRS; d++) // 8 dirs + value + for (int d = 0; d < NETHACK_DIR_HEADS * NETHACK_NUM_DIRS; d++) // 48 dir rows net->logits[NETHACK_NUM_ACTIONS + DEMO_PTR_HEADS * NETHACK_INV_SLOTS + d] = tmp[NETHACK_NUM_ACTIONS + d]; + { // spell head: dot(q_spell, rep_s) / 4 (dot-product pointer) + const float* qs = q + DEMO_PTR_HEADS * DEMO_INV_HID; + for (int sp = 0; sp < NETHACK_SPELL_SLOTS; sp++) { + float dot = 0.f; + for (int k = 0; k < DEMO_SPKEY; k++) dot += qs[k] * net->spkeys[sp * DEMO_SPKEY + k]; + net->logits[NETHACK_NUM_ACTIONS + DEMO_PTR_HEADS * NETHACK_INV_SLOTS + + NETHACK_DIR_HEADS * NETHACK_NUM_DIRS + sp] = dot * 0.25f; + } + } + net->logits[DEMO_OD] = tmp[NETHACK_NUM_ACTIONS + NETHACK_DIR_HEADS * NETHACK_NUM_DIRS]; // value return 0; } -// ---- interactive TTY demo ------------------------------------------------- +// interactive TTY demo // Space: one step on press; hold advances at 5 Hz. Shift+Space (or hold S): 20 Hz. // Terminals that support xterm modifyOtherKeys report Shift+Space as a CSI // sequence; 'S' is the fallback for everything else. q / Esc quits. @@ -411,8 +502,8 @@ static int g_term_raw = 0; static void demo_restore_term(void) { if (!g_term_raw) return; - printf("\x1b[>4;0m"); // disable modifyOtherKeys - printf("\x1b[?25h"); // show cursor + printf("\x1b[>4;0m"); // disable modifyOtherKeys + printf("\x1b[?25h"); // show cursor tcsetattr(STDIN_FILENO, TCSAFLUSH, &g_term_orig); g_term_raw = 0; fflush(stdout); @@ -447,98 +538,18 @@ static double demo_now(void) { return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9; } -// Weight resolution (eval the policy you just trained, not a stale fallback): -// 1. NH_WEIGHTS if set -// 2. Most recently written run under checkpoints/nethack/*/, highest-step -// .bin in that run (matches the intent of puffer eval --load-model-path=latest -// without being poisoned by an older multi-B-step run for a different role) -// 3. resources/nethack/nethack_weights.bin (checked-in demo fallback only) -// Global highest-step is wrong after a role switch (e.g. monk 2B vs valk 200M): -// the bigger step number is a different character. Resources-first is worse: -// a committed demo bin silently masks every local training run. +// NH_WEIGHTS overrides the checked-in demo weights static const char* demo_find_weights(void) { const char* envw = getenv("NH_WEIGHTS"); if (envw && envw[0]) return envw; - - static const char* resources = "resources/nethack/nethack_weights.bin"; - static char best_path[PATH_MAX]; - best_path[0] = '\0'; - - // Pass 1: find the run directory with the newest .bin mtime. - char best_run[PATH_MAX]; - best_run[0] = '\0'; - time_t best_run_mt = 0; - DIR* root = opendir("checkpoints/nethack"); - if (root) { - struct dirent* run; - while ((run = readdir(root)) != NULL) { - if (run->d_name[0] == '.') continue; - char rundir[PATH_MAX]; - int n = snprintf(rundir, sizeof(rundir), "checkpoints/nethack/%s", run->d_name); - if (n < 0 || n >= (int)sizeof(rundir)) continue; - DIR* rd = opendir(rundir); - if (!rd) continue; - time_t run_mt = 0; - struct dirent* f; - while ((f = readdir(rd)) != NULL) { - size_t len = strlen(f->d_name); - if (len < 5 || strcmp(f->d_name + len - 4, ".bin") != 0) continue; - char path[PATH_MAX]; - n = snprintf(path, sizeof(path), "%s/%s", rundir, f->d_name); - if (n < 0 || n >= (int)sizeof(path)) continue; - struct stat st; - if (stat(path, &st) != 0 || !S_ISREG(st.st_mode)) continue; - if (st.st_mtime >= run_mt) run_mt = st.st_mtime; - } - closedir(rd); - if (run_mt > best_run_mt) { - best_run_mt = run_mt; - snprintf(best_run, sizeof(best_run), "%s", rundir); - } - } - closedir(root); - } - - // Pass 2: highest zero-padded step within that run. - if (best_run[0]) { - long long best_step = -1; - time_t best_mt = 0; - DIR* rd = opendir(best_run); - if (rd) { - struct dirent* f; - while ((f = readdir(rd)) != NULL) { - size_t len = strlen(f->d_name); - if (len < 5 || strcmp(f->d_name + len - 4, ".bin") != 0) continue; - char path[PATH_MAX]; - int n = snprintf(path, sizeof(path), "%s/%s", best_run, f->d_name); - if (n < 0 || n >= (int)sizeof(path)) continue; - struct stat st; - if (stat(path, &st) != 0 || !S_ISREG(st.st_mode)) continue; - long long step = 0; - int digits = 0; - for (const char* p = f->d_name; *p >= '0' && *p <= '9'; p++) { - step = step * 10 + (*p - '0'); - digits++; - } - if (digits == 0) step = (long long)st.st_mtime; - if (step > best_step || (step == best_step && st.st_mtime >= best_mt)) { - best_step = step; - best_mt = st.st_mtime; - snprintf(best_path, sizeof(best_path), "%s", path); - } - } - closedir(rd); - } - } - if (best_path[0]) return best_path; - return resources; + return "resources/nethack/nethack_weights.bin"; } // Drain stdin. Returns a bitset: bit0=space, bit1=shift+space/S, bit2=quit. // Hold is detected via OS auto-repeat (and CSI for Shift+Space). -#define DEMO_IN_SPACE 1 -#define DEMO_IN_FAST 2 -#define DEMO_IN_QUIT 4 +#define DEMO_IN_SPACE 1 +#define DEMO_IN_FAST 2 +#define DEMO_IN_QUIT 4 // xterm/kitty encode modifiers as 1 + bitmask (Shift=1, Alt=2, Ctrl=4, ...) static int demo_mod_shift(int mod_param) { @@ -564,16 +575,22 @@ static int demo_poll_input(void) { } else if (c == 0x1b) { // Esc alone (no following bytes in this read) => quit. If CSI, // parse modifyOtherKeys / kitty sequences for Space. - if (i + 1 >= n || buf[i + 1] != '[') { flags |= DEMO_IN_QUIT; continue; } + if (i + 1 >= n || buf[i + 1] != '[') { + flags |= DEMO_IN_QUIT; + continue; + } i++; // at '[' int params[8], np = 0, val = 0, in_num = 0; memset(params, 0, sizeof(params)); for (i++; i < n; i++) { unsigned char d = buf[i]; - if (d >= '0' && d <= '9') { val = val * 10 + (d - '0'); in_num = 1; } - else if (d == ';') { + if (d >= '0' && d <= '9') { + val = val * 10 + (d - '0'); + in_num = 1; + } else if (d == ';') { if (np < 8) params[np++] = in_num ? val : 0; - val = 0; in_num = 0; + val = 0; + in_num = 0; } else if (d >= 0x40 && d <= 0x7e) { if (in_num && np < 8) params[np++] = val; // xterm modifyOtherKeys: CSI 27 ; mod ; keycode ~ @@ -601,18 +618,22 @@ static int demo_poll_input(void) { static void demo_step_once(NethackNet* net, Nethack* env, float* acts_f, float* ep_score, float* ep_len, float* ep_depth, float* ep_xp, float* ep_gt) { - nethack_net_forward(net, env->observations); + 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); - for (int h = 0; h < DEMO_NUM_HEADS; h++) env->actions[h] = acts_f[h]; + for (int h = 0; h < DEMO_NUM_HEADS; h++) env->agents[0].actions[h] = acts_f[h]; puf_step(env); - if (env->terminals[0] > 0.5f) { + 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\n", - env->log.score - *ep_score, env->log.episode_length - *ep_len, d, x, g); + fprintf(stderr, "episode end: 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->log.verb_uses[NETHACK_ACT_EAT], env->log.floor_eats, + env->log.verb_uses[NETHACK_ACT_WEAR], + env->log.verb_uses[NETHACK_ACT_THROW]); *ep_score = env->log.score; *ep_len = env->log.episode_length; *ep_depth = env->log.max_depth; @@ -623,12 +644,201 @@ static void demo_step_once(NethackNet* net, Nethack* env, float* acts_f, } } +// message history ring for the rich view +#define DEMO_MSG_RING 6 +static char demo_msgs[DEMO_MSG_RING][96]; +static int demo_msg_n = 0; + +static void demo_note_message(Nethack* env) { + if (!env->message[0]) return; + char buf[96]; + int j = 0; + for (; j < 95 && env->message[j]; j++) buf[j] = (char)env->message[j]; + buf[j] = 0; + if (demo_msg_n > 0 && strcmp(demo_msgs[(demo_msg_n - 1) % DEMO_MSG_RING], buf) == 0) + return; + strcpy(demo_msgs[demo_msg_n % DEMO_MSG_RING], buf); + demo_msg_n++; +} + +// menucolor-style BUC tint +static const char* demo_inv_clr(Nethack* env, int i) { + if (env->inv_oclasses[i] == 12) return "\x1b[33m"; // COIN_CLASS (objclass.h) + int buc = env->inv_state[i * NLE_INV_STATE_FIELDS]; + return buc == 1 ? "\x1b[31m" : buc == 3 ? "\x1b[32m" + : buc == 2 ? "\x1b[37m" : "\x1b[36m"; +} + +// curses-look cell: walls -> box drawing, floors -> '·', engine colors -> ANSI +#define DEMO_CMAP_OFF 2359 +#define DEMO_PET_LO 381 +#define DEMO_PET_HI 762 +static unsigned char demo_colors[NH_GRID]; // bound via nethack_color_sink +static const char* demo_wall_sym[12] = { // S_stone..S_trwall + " ", "│", "─", "┌", "┐", "└", "┘", + "┼", "┴", "┬", "┤", "├"}; +static const int demo_ansi_clr[16] = // CLR_* 0..15 -> SGR fg + {90, 31, 32, 33, 34, 35, 36, 37, 39, 91, 92, 93, 94, 95, 96, 97}; + +static void demo_map_cell(Nethack* env, int r, int c, int hero) { + // underfoot_glyphs hides the hero glyph; draw the @ explicitly + if (hero) { + printf("\x1b[7;1;97m@\x1b[0m"); + return; + } + int g = env->glyphs[r * NH_COLS + c]; + unsigned char ch = env->chars[r * NH_COLS + c]; + int cmap = g - DEMO_CMAP_OFF; + if (cmap >= 1 && cmap <= 11) { // walls in PUFF_CYAN (connect4.h) + printf("\x1b[38;2;0;187;187m%s\x1b[0m", demo_wall_sym[cmap]); + return; + } + if (cmap == 19) { // S_room + printf("\x1b[37m·\x1b[0m"); + return; + } + if (cmap == 20) { // S_darkroom + printf("\x1b[2m·\x1b[0m"); + return; + } + if (!ch || ch == ' ') { + putchar(' '); + return; + } + int clr = demo_colors[r * NH_COLS + c] & 15; + printf("\x1b[%s%dm%c\x1b[0m", + (g >= DEMO_PET_LO && g < DEMO_PET_HI) ? "4;" : "", + demo_ansi_clr[clr], ch); +} + +#define DEMO_INV_W 42 // inventory pane inner width + +// perm_invent pane: doname text grouped by class in game display order +static unsigned char demo_inv_strs[NLE_INVENTORY_SIZE * NLE_INVENTORY_STR_LENGTH]; +static const int demo_inv_order[] = + {12, 5, 2, 3, 7, 9, 10, 8, 4, 11, 6, 13, 14, 15, 16, 17, 1}; +static const char* demo_class_name[18] = {0, "Illegal objects", "Weapons", + "Armor", "Rings", "Amulets", "Tools", "Comestibles", "Potions", "Scrolls", + "Spellbooks", "Wands", "Coins", "Gems/Stones", "Boulders/Statues", + "Iron balls", "Chains", "Venoms"}; +#define DEMO_PANE_MAX 64 + +static int demo_inv_pane(Nethack* env, char lines[][DEMO_INV_W + 8], + const char* clrs[]) { + int n = 0; + for (int k = 0; k < (int)(sizeof(demo_inv_order)/sizeof(*demo_inv_order)); k++) { + int cls = demo_inv_order[k], first = 1; + for (int i = 0; i < NETHACK_INV_SLOTS && n < DEMO_PANE_MAX - 1; i++) { + if (!env->inv_letters[i] || env->inv_oclasses[i] != cls) continue; + if (first) { + snprintf(lines[n], DEMO_INV_W + 8, "%s", demo_class_name[cls]); + clrs[n++] = "\x1b[1m"; + first = 0; + } + snprintf(lines[n], DEMO_INV_W + 8, "%c) %s", env->inv_letters[i], + demo_inv_strs + i * NLE_INVENTORY_STR_LENGTH); + clrs[n] = demo_inv_clr(env, i); + n++; + } + } + return n; +} + +static void demo_box_edge(const char* l, const char* r, const char* title, int inner) { + printf("%s", l); + int n = 0; + if (title) n = printf("─ %s ", title) - 2; // rule char is 3 bytes, 1 column + for (; n < inner; n++) printf("─"); + printf("%s", r); +} + static void demo_render(Nethack* env, int rate_hz, long steps) { - puf_render(env); - printf("steps %ld | SPACE step/hold 5Hz | Shift+SPACE (or S) 20Hz | q quit", + long* bl = env->blstats; + printf("\x1b[H\x1b[2J"); + // message window + int total = NH_COLS + 3 + DEMO_INV_W; // map box + gap + inv box inner + demo_box_edge("┌", "┐\n", NULL, total); + int shown = demo_msg_n < DEMO_MSG_RING ? demo_msg_n : DEMO_MSG_RING; + for (int k = 0; k < DEMO_MSG_RING; k++) { + printf("│ "); + int len = 0; + if (k >= DEMO_MSG_RING - shown) { + int idx = (demo_msg_n - (DEMO_MSG_RING - k)) % DEMO_MSG_RING; + len = printf("%s%.*s\x1b[0m", k == DEMO_MSG_RING - 1 ? "\x1b[1m" : "\x1b[2m", + total - 2, demo_msgs[idx]) - 8; + } + printf("%*s│\n", total - 1 - len, ""); + } + demo_box_edge("└", "┘\n", NULL, total); + // map box + inventory box + char title[32]; + static char plines[DEMO_PANE_MAX][DEMO_INV_W + 8]; + static const char* pclrs[DEMO_PANE_MAX]; + int pn = demo_inv_pane(env, plines, pclrs); + snprintf(title, sizeof(title), "Dlvl:%ld", bl[NLE_BL_DEPTH]); + demo_box_edge("┌", "┐", title, NH_COLS); + printf(" "); + demo_box_edge("┌", "┐\n", "Inventory", DEMO_INV_W); + for (int r = 0; r < NH_ROWS; r++) { + printf("│"); + for (int c = 0; c < NH_COLS; c++) + demo_map_cell(env, r, c, c == bl[NLE_BL_X] && r == bl[NLE_BL_Y]); + printf("│ │ "); + int len = 0; + if (r < NH_ROWS - 1 || pn <= NH_ROWS) { + if (r < pn) + len = printf("%s%.*s\x1b[0m", pclrs[r], DEMO_INV_W - 2, plines[r]) + - (int)strlen(pclrs[r]) - 4; + } else + len = printf("\x1b[2m... +%d more\x1b[0m", pn - (NH_ROWS - 1)) - 8; + printf("%*s│\n", DEMO_INV_W - 1 - len, ""); + } + demo_box_edge("└", "┘", NULL, NH_COLS); + printf(" "); + demo_box_edge("└", "┘\n", NULL, DEMO_INV_W); + // status: hitpointbar over the agent name, then the stat line + long hp = bl[NLE_BL_HP], hpm = bl[NLE_BL_HPMAX] > 0 ? bl[NLE_BL_HPMAX] : 1; + const char* name = "Agent the Puffer"; + int fill = (int)((hp * 16 + hpm - 1) / hpm); + if (fill > 16) fill = 16; + int hpc = hp * 3 >= hpm * 2 ? 32 : hp * 3 >= hpm ? 33 : 31; + printf("[\x1b[7;%dm%.*s\x1b[0m%s] St:%ld Dx:%ld Co:%ld Score:%ld\n", + hpc, fill, name, name + fill, + bl[NLE_BL_STR25], bl[NLE_BL_DEX], bl[NLE_BL_CON], bl[NLE_BL_SCORE]); + static const char* conds[10] = {"Stone", "Slime", "Strngl", "FoodPois", + "TermIll", "Blind", "Deaf", "Stun", "Conf", "Hallu"}; + static const char* hungers[5] = {"Satiated", "", "Hungry", "Weak", "Fainting"}; + long hu = bl[NLE_BL_HUNGER]; + printf("Dlvl:%ld $:%ld HP:%ld(%ld) Pw:%ld(%ld) AC:%ld Xp:%ld/%ld T:%ld", + bl[NLE_BL_DEPTH], bl[NLE_BL_GOLD], hp, bl[NLE_BL_HPMAX], + bl[NLE_BL_ENE], bl[NLE_BL_ENEMAX], bl[NLE_BL_AC], + bl[NLE_BL_XP], bl[NLE_BL_EXP], bl[NLE_BL_TIME]); + if (hu >= 0 && hu < 5 && hungers[hu][0]) printf(" \x1b[33m%s\x1b[0m", hungers[hu]); + // known spells: name Lv fail%% (env->spell_* is refreshed each pack_obs) + if (env->n_spells > 0) { + static const struct { int id; const char* nm; } spnames[] = { + {344, "sleep"}, {348, "healing"}, {377, "protection"}, + {340, "force bolt"}, {342, "magic missile"}, {361, "cure blindness"}, + }; + printf(" \x1b[36mSp:"); + for (int i = 0; i < env->n_spells && i < 3; i++) { + const char* nm = NULL; + for (unsigned k = 0; k < sizeof(spnames)/sizeof(spnames[0]); k++) { + if (spnames[k].id != env->spell_ids[i]) continue; + nm = spnames[k].nm; + break; + } + printf("%s%s(L%d %d%%)", i ? "," : "", + nm ? nm : "spell", env->spell_levs[i], env->spell_fails[i]); + } + printf("\x1b[0m"); + } + for (int b = 0; b < 10; b++) + if (bl[NLE_BL_CONDITION] & (1L << b)) printf(" \x1b[31;1m%s\x1b[0m", conds[b]); + printf("\n\x1b[2msteps %ld | SPACE step/hold 5Hz | Shift+SPACE (or S) 20Hz | q quit", steps); if (rate_hz > 0) printf(" | running %d Hz", rate_hz); - printf("\n"); + printf("\x1b[0m\n"); fflush(stdout); } @@ -644,9 +854,11 @@ static void run_demo_interactive(long max_steps) { NethackNet* net = make_nethack_net(w); Nethack env; - env_open(&env); + nethack_color_sink = demo_colors; + nethack_invstr_sink = demo_inv_strs; const char* seed_env = getenv("NH_SEED"); srand(seed_env ? (unsigned)strtoul(seed_env, NULL, 10) : (unsigned)time(NULL)); + env_open(&env); demo_raw_term(); @@ -658,8 +870,8 @@ static void run_demo_interactive(long max_steps) { // first SPACE/S -> exactly one step // further events -> continuous 5 Hz (space) or 20 Hz (shift+space / S) // Grace after first press covers the typical OS key-repeat delay (~0.5 s). - int mode = 0; // 0=idle, 1=slow, 2=fast - int confirmed_hold = 0; // saw a second key event (auto-repeat) + int mode = 0; // 0=idle, 1=slow, 2=fast + int confirmed_hold = 0; // saw a second key event (auto-repeat) int edge_pending = 0; double held_until = 0; double next_step_at = 0; @@ -709,6 +921,7 @@ static void run_demo_interactive(long max_steps) { if (do_step) { demo_step_once(net, &env, acts_f, &ep_score, &ep_len, &ep_depth, &ep_xp, &ep_gt); + demo_note_message(&env); steps++; demo_render(&env, confirmed_hold ? rate : 0, steps); } else { @@ -725,7 +938,10 @@ static void run_demo_interactive(long max_steps) { printf("steps=%ld\n", steps); env_close(&env); free_mingru(net->mingru); - free(net->md); free(net->hidden); free(net->e_eff); free(net); + free(net->md); + free(net->hidden); + free(net->e_eff); + free(net); free(w); } @@ -741,9 +957,9 @@ static void run_demo_auto(long max_steps, int frame_ms) { NethackNet* net = make_nethack_net(w); Nethack env; - env_open(&env); const char* seed_env = getenv("NH_SEED"); srand(seed_env ? (unsigned)strtoul(seed_env, NULL, 10) : (unsigned)time(NULL)); + env_open(&env); float ep_score = 0, ep_len = 0, ep_depth = 0, ep_xp = 0, ep_gt = 0; float acts_f[DEMO_NUM_HEADS]; @@ -761,7 +977,10 @@ static void run_demo_auto(long max_steps, int frame_ms) { env.log.max_depth / env.log.n, env.log.max_xp_level / env.log.n); env_close(&env); free_mingru(net->mingru); - free(net->md); free(net->hidden); free(net->e_eff); free(net); + free(net->md); + free(net->hidden); + free(net->e_eff); + free(net); free(w); } @@ -775,9 +994,9 @@ int main(int argc, char** argv) { int frame_ms = 50; if (argc >= 3) { frame_ms = atoi(argv[2]); - interactive = 0; // explicit frame timing => auto mode + interactive = 0; // explicit frame timing => auto mode } else if (!interactive) { - frame_ms = 0; // piped/non-TTY default: headless auto-run + frame_ms = 0; // piped/non-TTY default: headless auto-run } if (interactive) run_demo_interactive(max_steps); else run_demo_auto(max_steps, frame_ms); diff --git a/ocean/nethack/nethack.cu b/ocean/nethack/nethack.cu index f9651c6bda..8ea596cec0 100644 --- a/ocean/nethack/nethack.cu +++ b/ocean/nethack/nethack.cu @@ -12,7 +12,7 @@ __global__ void nh_bias_relu_kernel( data[idx] = from_float(fmaxf(0.0f, to_float(data[idx]) + to_float(bias[idx % dim]))); } -// ---- Nethack constants ---- +// constants // Obs layout (must match ocean/nethack/nethack.h): // [0, 2*NH_MGRID) full 79x21 glyph grid, int16 LE (map memory included) // [2*NH_MGRID, +4*NH_BL_RAW) blstats, int32 LE (x,y first) @@ -21,15 +21,15 @@ __global__ void nh_bias_relu_kernel( static constexpr int NH_MAPW = 79, NH_MAPH = 21; 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_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_CROP = 9, NH_CHALF = 4; // NETHACK_CROP, egocentric +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_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; static constexpr int NH_LOC_HID = 256; // Global branch: per patch, embed+flatten (25 cells x 32 dims) + normalized @@ -41,34 +41,42 @@ static constexpr int NH_LOC_HID = 256; // as concatenating onto the flatten. static constexpr int NH_P1 = 16; 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_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) -static constexpr int NH_BL_RAW = 27; // NLE_BLSTATS_SIZE +static constexpr int NH_HOT_G = 10; // hot-glyph dT smem slots (10x400 int64 = 32KB) +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 -static constexpr int NH_ACTIONS = 22; // NETHACK_NUM_ACTIONS -static constexpr int NH_OCLASSES = 18; // MAXOCLASSES -static constexpr int NH_EX_RAW = 2 + NH_OCLASSES; // NETHACK_EXTRA_INTS +static constexpr int NH_BL_HP = 10, NH_BL_ENE = 14; // hp/hpmax at 10/11, ene/enemax at 14/15 +static constexpr int NH_ACTIONS = 26; // NETHACK_NUM_ACTIONS +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_RAW = 2 + NH_OCLASSES + 2 + 1 + 4 * NH_SPELL_SLOTS + 2; // NETHACK_EXTRA_INTS (+spell block +weight pair) // 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 // branch is nominal, not ordinal. -static constexpr int NH_F_HUNGER = 25; // 7-way one-hot -static constexpr int NH_F_COND = NH_F_HUNGER + 7; // 13 condition bits -static constexpr int NH_F_PREV = NH_F_COND + 13; // prev-action one-hot -static constexpr int NH_F_INV = NH_F_PREV + NH_ACTIONS; // inv class counts -static constexpr int NH_F_FRAC = NH_F_INV + NH_OCLASSES; // hp_frac, ene_frac -static constexpr int NH_F_DNUM = NH_F_FRAC + 2; // 8-way one-hot -static constexpr int NH_F_ENGR = NH_F_DNUM + 8; // engraving bits -static constexpr int NH_BL_FEAT = NH_F_ENGR + 2; +static constexpr int NH_F_HUNGER = 25; // 7-way one-hot +static constexpr int NH_F_COND = NH_F_HUNGER + 7; // 13 condition bits +static constexpr int NH_F_PREV = NH_F_COND + 13; // prev-action one-hot +static constexpr int NH_F_INV = NH_F_PREV + NH_ACTIONS; // inv class counts +static constexpr int NH_F_FRAC = NH_F_INV + NH_OCLASSES; // hp_frac, ene_frac +static constexpr int NH_F_DNUM = NH_F_FRAC + 2; // 8-way one-hot +static constexpr int NH_F_ENGR = NH_F_DNUM + 8; // engraving bits +static constexpr int NH_F_SHOP = NH_F_ENGR + 2; // in-shop, affordability +// spell scalar feature: known count only — per-slot content rides the +// spell-key path (v3 pointer), not the blstats block +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_BL_FEAT = NH_F_WEIGHT + 2; 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 // linear -> relu. The per-slot vectors are the pointer decoder's keys (slot // identity lives there); the trunk only gets the pooled summary below. Fused // per-glyph table T_inv = E @ inv1_w^T (5977xNH_INV_HID) rebuilt per forward. -static constexpr int NH_INV = 55; // NETHACK_INV_SLOTS +static constexpr int NH_INV = 55; // NETHACK_INV_SLOTS // 16-dim per-slot rep: doubles as the pooled-summary bottleneck AND the // decoder pointer key (unified, patch-encoder style). Halves the pool max // MACs, the deterministic max-backward atomics, and the max-kernel smem @@ -84,24 +92,32 @@ static constexpr int NH_INV_POOL = 128; // hashes into NH_MSG_VOCAB buckets, its NH_MSG_HID-dim embed row summed, then // 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_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; -static constexpr int NH_CONCAT = NH_MSG_CONCAT_OFF + NH_MSG_HID; -static constexpr int NH_BL_OFF = 2 * NH_MGRID; // blstats offset, obs elements +// 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 +static constexpr int NH_SPELL_CONCAT_OFF = NH_MSG_CONCAT_OFF + NH_MSG_HID; +static constexpr int NH_CONCAT = NH_SPELL_CONCAT_OFF + NH_SPKEY; +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 // [buc, spe(-128=unknown), quan, ero1, ero2, flags, typeknown, rsvd], // expanded to NH_SFEAT features feeding the slot MLP beside the embed static constexpr int NH_INVST_OFF = NH_INV_OFF + NH_INV * 2; -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 -static constexpr int NH_MSG_OFF = NH_INVST_OFF + NH_INV * NH_ST_RAW; // message block start +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_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_SORT_BLOCKS = 256; // hist grid (smem histograms) -static constexpr int NH_HOT_T = 16; // hot-glyph smem rows (16x32 int64 = 4KB) +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) // Residual factorized glyph embedding: E_eff = E_res + E_kind[kind(g)] + // E_sub[sub(g)]. The (kind, sub) mapping is generated from the engine's own @@ -142,7 +158,7 @@ static __device__ const signed char nh_obj_armcat_dev[NH_NUM_OBJECTS] = { }; static_assert(NH_GM_VOCAB == 5977, "glyph map vocab mismatch"); static constexpr int NH_NKIND = NH_GM_NKIND; -static constexpr int NH_NSUB = NH_GM_NSUB; +static constexpr int NH_NSUB = NH_GM_NSUB; // 2^24 fixed-point gradient accumulators: integer atomics are associative, so // scatter/bias sums are bit-identical run to run (float atomicAdd ordering is @@ -172,7 +188,10 @@ __global__ void nh_fxp_to_precision_rows_kernel( int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= n) return; int row = idx / trow; - if (counts[row] == 0 && hot_map[row] < 0) { dst[idx] = from_float(0.0f); return; } + if (counts[row] == 0 && hot_map[row] < 0) { + dst[idx] = from_float(0.0f); + return; + } dst[idx] = from_float(nh_fxp_to_float(src[idx])); src[idx] = 0; } @@ -180,23 +199,23 @@ __global__ void nh_fxp_to_precision_rows_kernel( // Per-blstat normalization: log1p fields get log1p(max(v,0))*scale, the rest // v*scale. Hunger (21) and condition (25) are expanded, not scaled. __constant__ float NH_BL_SCALE[NH_BL_RAW] = { - 1.f/79, 1.f/21, // x, y + 1.f/79, 1.f/21, // x, y 1.f/25, 1.f/125, 1.f/25, 1.f/25, 1.f/25, 1.f/25, 1.f/25, // str25 str125 dex con int wis cha - 0.1f, // score (log) - 1.f/200, 1.f/200, 1.f/50, // hp, hpmax, depth - 0.1f, // gold (log) - 1.f/100, 1.f/100, 1.f/10, 1.f/10, 1.f/30, // ene, enemax, ac, hd, xp level - 0.1f, 0.1f, // exp points, time (log) - 0.f, // hunger (expanded) - 1.f/4, 0.f, 1.f/50, // cap, dnum (one-hot), dlevel - 0.f, // condition (expanded) - 1.f, // align + 0.1f, // score (log) + 1.f/200, 1.f/200, 1.f/50, // hp, hpmax, depth + 0.1f, // gold (log) + 1.f/100, 1.f/100, 1.f/10, 1.f/10, 1.f/30, // ene, enemax, ac, hd, xp level + 0.1f, 0.1f, // exp points, time (log) + 0.f, // hunger (expanded) + 1.f/4, 0.f, 1.f/50, // cap, dnum (one-hot), dlevel + 0.f, // condition (expanded) + 1.f, // align }; __constant__ int NH_BL_ISLOG[NH_BL_RAW] = { 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, }; -// ---- Nethack kernels ---- +// kernels __device__ __forceinline__ int nh_bl_read_i32(const precision_t* p) { return (int)((unsigned int)(int)to_float(p[0]) @@ -225,8 +244,8 @@ __global__ void nh_crop_kernel( if (t >= B * NH_CGRID) return; int b = t / NH_CGRID, p = t % NH_CGRID; const precision_t* bl = obs + (int64_t)b * NH_OBS_SIZE + NH_BL_OFF; - int r = nh_bl_read_i32(bl + 4) - NH_CHALF + p / NH_CROP; // blstats[1] = y - int c = nh_bl_read_i32(bl) - NH_CHALF + p % NH_CROP; // blstats[0] = x + int r = nh_bl_read_i32(bl + 4) - NH_CHALF + p / NH_CROP; // blstats[1] = y + int c = nh_bl_read_i32(bl) - NH_CHALF + p % NH_CROP; // blstats[0] = x crop[t] = (r < 0 || r >= NH_MAPH || c < 0 || c >= NH_MAPW) ? (float)NH_PAD_GLYPH : idx[b * NH_MGRID + r * NH_MAPW + c]; } @@ -238,7 +257,7 @@ __global__ void nh_local_gather_kernel( int t = blockIdx.x * blockDim.x + threadIdx.x; if (t >= B * NH_LOC_IN) return; int d = t % NH_EMBED_DIM; - int64_t cell = t / NH_EMBED_DIM; // b*NH_CGRID + p + int64_t cell = t / NH_EMBED_DIM; // b*NH_CGRID + p x[t] = E[(int64_t)(int)crop[cell] * NH_EMBED_DIM + d]; } @@ -322,8 +341,8 @@ __global__ void nh_patch_max_kernel( if (b >= B) return; if (threadIdx.x == 0) { const precision_t* bl = obs + (int64_t)b * NH_OBS_SIZE + NH_BL_OFF; - hero[0] = (float)nh_bl_read_i32(bl); // x - hero[1] = (float)nh_bl_read_i32(bl + 4); // y + hero[0] = (float)nh_bl_read_i32(bl); // x + hero[1] = (float)nh_bl_read_i32(bl + 4); // y } for (int i = threadIdx.x; i < NH_GLB_HID * NH_P1; i += blockDim.x) w2s[i] = to_float(w2[i]); @@ -356,7 +375,10 @@ __global__ void nh_patch_max_kernel( float v = 0.0f; for (int k = 0; k < NH_P1; k++) v += w2s[o * NH_P1 + k] * t16s[tk * NH_P1 + k]; - if (v > best) { best = v; bm = tk; } + if (v > best) { + best = v; + bm = tk; + } } glb_out[(int64_t)b * NH_GLB_HID + o] = from_float(fmaxf(best + to_float(b2[o]), 0.0f)); argmax[(int64_t)b * NH_GLB_HID + o] = bm; @@ -444,15 +466,33 @@ __global__ void nh_blstats_kernel( // hp_frac / ene_frac in [0,1]: the "how close to death/empty" ratio int base = (j == NH_F_FRAC) ? NH_BL_HP : NH_BL_ENE; int cur = nh_bl_read_i32(src + 4*base); - int mx = nh_bl_read_i32(src + 4*(base + 1)); + int mx = nh_bl_read_i32(src + 4*(base + 1)); f = fminf(fmaxf((float)cur / (float)(mx > 1 ? mx : 1), 0.0f), 1.0f); } else if (j < NH_F_ENGR) { int v = nh_bl_read_i32(src + 4*NH_BL_DNUM); f = (j - NH_F_DNUM == max(0, min(v, 7))) ? 1.0f : 0.0f; - } else { + } else if (j < NH_F_SHOP) { // underfoot engraving from ex[0]: any engraving, active Elbereth int v = nh_bl_read_i32(ex); f = (j == NH_F_ENGR) ? (v >= 1 ? 1.0f : 0.0f) : (v >= 2 ? 1.0f : 0.0f); + } else if (j < NH_F_SPELL) { + // standing on shop goods, and gold/price capped at 1 + int v = nh_bl_read_i32(ex + 4*(NH_EXTRA_SHOP + (j - NH_F_SHOP))); + f = (j == NH_F_SHOP) ? (float)v : (float)v * 0.01f; + } else if (j < NH_F_WEIGHT) { + // known-spell count/8; per-slot content rides the spell-key path + int v = nh_bl_read_i32(ex + 4*(NH_EXTRA_SHOP + 2)); + f = (float)v * 0.125f; + } else { + // encumbrance: softsign(ratio-1) is 0 at the wall, keeps gradient + // through Overloaded (~3x); capacity /1000 (engine caps there) + int v = nh_bl_read_i32(ex + 4*(NH_EXTRA_SHOP + 2 + 1 + 4 * NH_SPELL_SLOTS + + (j - NH_F_WEIGHT))); + if (j == NH_F_WEIGHT) { + float d = (float)v * 0.01f - 1.0f; + f = d / (1.0f + fabsf(d)); + } else + f = (float)v * 0.001f; } // 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) @@ -461,13 +501,13 @@ __global__ void nh_blstats_kernel( } // concat = [local hid | global hid | bl hid | bl raw feats] -// ---- trigram message branch ---- +// trigram message branch __device__ __forceinline__ int nh_msg_lc(int c) { - return (c >= 'A' && c <= 'Z') ? c + 32 : c; // lowercase; keep spaces/punct + return (c >= 'A' && c <= 'Z') ? c + 32 : c; // lowercase; keep spaces/punct } __device__ __forceinline__ int nh_msg_hash(int c0, int c1, int c2) { unsigned key = ((unsigned)c0 << 16) | ((unsigned)c1 << 8) | (unsigned)c2; - return (int)((key * 2654435761u) >> (32 - NH_MSG_LOG2V)); // top log2V bits + return (int)((key * 2654435761u) >> (32 - NH_MSG_LOG2V)); // top log2V bits } // per-position trigram bucket id (-1 for the padded tail / past the null). Ids // stay contiguous because the topline is null-terminated, so consumers break @@ -477,10 +517,16 @@ __global__ void nh_msg_ids_kernel( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= B * NH_MSG_LEN) return; int b = i / NH_MSG_LEN, t = i % NH_MSG_LEN; - if (t > NH_MSG_LEN - 3) { ids[i] = -1.0f; return; } + if (t > NH_MSG_LEN - 3) { + ids[i] = -1.0f; + return; + } const precision_t* m = obs + (int64_t)b * NH_OBS_SIZE + NH_MSG_OFF; int c0 = (int)to_float(m[t]), c1 = (int)to_float(m[t + 1]), c2 = (int)to_float(m[t + 2]); - if (c0 == 0 || c1 == 0 || c2 == 0) { ids[i] = -1.0f; return; } + if (c0 == 0 || c1 == 0 || c2 == 0) { + ids[i] = -1.0f; + return; + } ids[i] = (float)nh_msg_hash(nh_msg_lc(c0), nh_msg_lc(c1), nh_msg_lc(c2)); } // normalized-sum bag: block per sample, one warp = NH_MSG_HID lanes (lane d @@ -514,7 +560,10 @@ __global__ void nh_msg_bwd_kernel( int d = threadIdx.x; const float* mi = ids + (int64_t)b * NH_MSG_LEN; int count = 0; - for (int t = 0; t < NH_MSG_LEN; t++) { if ((int)mi[t] < 0) break; count++; } + for (int t = 0; t < NH_MSG_LEN; t++) { + if ((int)mi[t] < 0) break; + count++; + } float g = to_float(grad_concat[(int64_t)b * NH_CONCAT + NH_MSG_CONCAT_OFF + d]) * rsqrtf((float)count + 1.0f); if (g == 0.0f) return; @@ -522,6 +571,136 @@ __global__ void nh_msg_bwd_kernel( nh_fxp_atomic_add(&dmsg_acc[(int64_t)(int)mi[t] * NH_MSG_HID + d], g); } +// spell-key forward: per slot, build the 36-dim input [e_eff(book glyph) | +// known, lev/7, fail/100, know/20000], key = spk_w . input. Saves inputs + +// glyph ids for the backward; empty slot -> zero input -> zero key. +__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) { + 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; + 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; + int lev = (int)to_float(src[4]) | ((int)to_float(src[5]) << 8); + int fail = (int)to_float(src[8]) | ((int)to_float(src[9]) << 8); + int know = (int)to_float(src[12]) | ((int)to_float(src[13]) << 8) + | ((int)to_float(src[14]) << 16); + in[NH_EMBED_DIM + 0] = id > 0 ? 1.0f : 0.0f; + 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) + } +} + +// inventory-style pool: project each relu'd slot rep, max over slots, bias, +// relu (mirrors inv2). Saves the argmax slot per (sample, pool dim). +__global__ void nh_sppool_kernel(precision_t* __restrict__ concat, + precision_t* __restrict__ pool, int* __restrict__ amax, + const precision_t* __restrict__ keys, + const precision_t* __restrict__ spk2_w, const precision_t* __restrict__ spk2_b, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_SPKEY) return; + int b = t / NH_SPKEY, d = t % NH_SPKEY; + float best = -1e30f; + int bs = 0; + for (int s = 0; s < NH_SPELL_SLOTS; s++) { + float acc = 0.0f; + for (int k = 0; k < NH_SPKEY; k++) + acc += to_float(spk2_w[d * NH_SPKEY + k]) + * to_float(keys[((int64_t)b * NH_SPELL_SLOTS + s) * NH_SPKEY + k]); + if (acc > best) { + best = acc; + bs = s; + } + } + amax[t] = bs; + float v = fmaxf(best + to_float(spk2_b[d]), 0.0f); + pool[t] = from_float(v); + concat[(int64_t)b * NH_CONCAT + NH_SPELL_CONCAT_OFF + d] = 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; +// finally masked by the rep's own relu. +__global__ void nh_spkey_dk_kernel(precision_t* __restrict__ dkeys, + const precision_t* __restrict__ grad_concat, const precision_t* __restrict__ pool, + const int* __restrict__ amax, const precision_t* __restrict__ spk2_w, + const precision_t* __restrict__ keys, const precision_t* __restrict__ ptr_dkeys, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_SPELL_SLOTS * NH_SPKEY) return; + int b = t / (NH_SPELL_SLOTS * NH_SPKEY); + int s = (t / NH_SPKEY) % NH_SPELL_SLOTS; + int k = t % NH_SPKEY; + float v = ptr_dkeys ? to_float(ptr_dkeys[t]) : 0.0f; + for (int d = 0; d < NH_SPKEY; d++) { + if (amax[(int64_t)b * NH_SPKEY + d] != s) continue; + if (to_float(pool[(int64_t)b * NH_SPKEY + d]) <= 0.0f) + continue; // pool relu gate + v += to_float(grad_concat[(int64_t)b * NH_CONCAT + NH_SPELL_CONCAT_OFF + d]) + * to_float(spk2_w[d * NH_SPKEY + k]); + } + if (to_float(keys[t]) <= 0.0f) v = 0.0f; // rep relu gate + dkeys[t] = from_float(v); +} + +// spk2 grads: (16x16 + 16) threads loop the batch; small enough to be cheap +__global__ void nh_spk2_grad_kernel(precision_t* __restrict__ spk2_wgrad, + precision_t* __restrict__ spk2_bgrad, const precision_t* __restrict__ grad_concat, + const precision_t* __restrict__ pool, const int* __restrict__ amax, + const precision_t* __restrict__ keys, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= NH_SPKEY * (NH_SPKEY + 1)) return; + int d = t / (NH_SPKEY + 1), k = t % (NH_SPKEY + 1); + float acc = 0.0f; + for (int b = 0; b < B; b++) { + if (to_float(pool[(int64_t)b * NH_SPKEY + d]) <= 0.0f) + continue; + float g = to_float(grad_concat[(int64_t)b * NH_CONCAT + NH_SPELL_CONCAT_OFF + d]); + if (g == 0.0f) continue; + if (k == NH_SPKEY) { // bias + acc += g; + continue; + } + int s = amax[(int64_t)b * NH_SPKEY + d]; + acc += g * to_float(keys[((int64_t)b * NH_SPELL_SLOTS + s) * NH_SPKEY + k]); + } + if (k == NH_SPKEY) spk2_bgrad[d] = from_float(acc); + else spk2_wgrad[d * NH_SPKEY + k] = from_float(acc); +} + +// stage 2: scatter d_embed = spk_w[:, :32]^T . dkey into the fxp dE staging +__global__ void nh_spkey_dE_kernel(long long* __restrict__ dE_i, + const precision_t* __restrict__ dkeys, const precision_t* __restrict__ spk_w, + const float* __restrict__ sp_idx, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_SPELL_SLOTS * NH_EMBED_DIM) return; + int bs = t / NH_EMBED_DIM; + int d = t % NH_EMBED_DIM; + int g = (int)sp_idx[bs]; + if (g < 0) return; + float acc = 0.0f; + for (int r = 0; r < NH_SPKEY; r++) + acc += to_float(spk_w[r * NH_SPIN + d]) * to_float(dkeys[(int64_t)bs * NH_SPKEY + r]); + 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, @@ -541,8 +720,10 @@ __global__ void nh_concat_kernel( val = bl_out[(int64_t)b * NH_BL_HID + (c - NH_LOC_HID - NH_GLB_HID - NH_INV_POOL)]; 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)]; - else + else if (c < NH_SPELL_CONCAT_OFF) val = msg[(int64_t)b * NH_MSG_HID + (c - NH_MSG_CONCAT_OFF)]; + else + return; // spell slice: nh_spell_gather_kernel fills it afterwards out[idx] = val; } @@ -555,14 +736,14 @@ __global__ void nh_slice_kernel( dst[idx] = src[(idx / n) * stride + offset + idx % n]; } -// ---- inventory entity branch ---- +// inventory entity branch __global__ void nh_inv_decode_kernel( - float* __restrict__ idx, const precision_t* __restrict__ obs, int B) { + float* __restrict__ idx, const precision_t* __restrict__ obs, int B, int off) { int t = blockIdx.x * blockDim.x + threadIdx.x; if (t >= B * NH_INV) return; int b = t / NH_INV, s = t % NH_INV; - const precision_t* src = obs + (int64_t)b * NH_OBS_SIZE + NH_INV_OFF + 2 * s; + const precision_t* src = obs + (int64_t)b * NH_OBS_SIZE + off + 2 * s; int g = (int)to_float(src[0]) | ((int)to_float(src[1]) << 8); idx[t] = (float)max(0, min(g, NH_GLYPH_VOCAB - 1)); } @@ -590,14 +771,14 @@ __global__ void nh_inv_sfeat_kernel(precision_t* __restrict__ out, int st[NH_ST_RAW]; for (int j = 0; j < NH_ST_RAW; j++) { int v = (int)to_float(src[j]); - st[j] = v >= 128 ? v - 256 : v; // bytes -> int8 + st[j] = v >= 128 ? v - 256 : v; // bytes -> int8 } precision_t* f = out + (int64_t)t * NH_SFEAT; 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 + f[6] = from_float(log1pf(fmaxf((float)st[2], 0.0f)) * 0.5f); // guard vs unclamped export drift 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++) @@ -616,13 +797,18 @@ __global__ void nh_inv_sfeat_kernel(precision_t* __restrict__ out, __global__ void nh_inv_gather_kernel(precision_t* __restrict__ out, const precision_t* __restrict__ T, const precision_t* __restrict__ b1, const precision_t* __restrict__ ws, const precision_t* __restrict__ sfeat, - const float* __restrict__ idx, int B) { + const float* __restrict__ idx, const precision_t* __restrict__ Tt, + const float* __restrict__ idxt, int B) { int t = blockIdx.x * blockDim.x + threadIdx.x; if (t >= B * NH_INV_FLAT) return; int b = t / NH_INV_FLAT, r = t % NH_INV_FLAT; int s = r / NH_INV_HID, k = r % NH_INV_HID; int g = (int)idx[b * NH_INV + s]; float v = to_float(T[g * NH_INV_HID + k]) + to_float(b1[k]); + // discovered-type channel: pad = unknown identity, hard zero contribution + int gt = (int)idxt[b * NH_INV + s]; + if (gt != NH_GLYPH_VOCAB - 1) + v += to_float(Tt[gt * NH_INV_HID + k]); const precision_t* f = sfeat + ((int64_t)b * NH_INV + s) * NH_SFEAT; for (int j = 0; j < NH_SFEAT; j++) v += to_float(ws[k * NH_SFEAT + j]) * to_float(f[j]); @@ -653,7 +839,10 @@ __global__ void nh_inv_max_kernel(precision_t* __restrict__ pool_out, float v = 0.0f; for (int k = 0; k < NH_INV_HID; k++) v += w2s[o * NH_INV_HID + k] * ss[s * NH_INV_HID + k]; - if (v > best) { best = v; bm = s; } + if (v > best) { + best = v; + bm = s; + } } pool_out[(int64_t)b * NH_INV_POOL + o] = from_float(fmaxf(best + to_float(b2[o]), 0.0f)); argmax[(int64_t)b * NH_INV_POOL + o] = bm; @@ -700,12 +889,15 @@ __global__ void nh_inv_max_bwd_kernel(precision_t* __restrict__ inv_grad, // dT_inv scatter, plain global fxp atomics: 55x16 per sample is too small // for the hot-row machinery to pay __global__ void nh_dTinv_scatter_kernel(long long* __restrict__ dT, - const precision_t* __restrict__ dflat, const float* __restrict__ idx, int64_t n) { + const precision_t* __restrict__ dflat, const float* __restrict__ idx, int64_t n, + int skip_pad) { int64_t t = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; if (t >= n) return; float v = to_float(dflat[t]); if (v == 0.0f) return; int g = (int)idx[t / NH_INV_HID]; + // true-glyph channel: pad slots contributed hard zero in the forward + if (skip_pad && g == NH_GLYPH_VOCAB - 1) return; nh_fxp_atomic_add(&dT[(int64_t)g * NH_INV_HID + t % NH_INV_HID], v); } @@ -766,7 +958,7 @@ __global__ void nh_bias_flush_kernel( else if ((i -= n5) < n6) d6[i] = from_float(v); } -// ---- embedding backward ---- +// embedding backward // Both views are linear in the embeddings, so dE is a scatter-add of per-cell // 32-dim grad vectors into glyph rows. The dominant glyphs (unexplored stone, // floor, walls cover ~80% of cells) contend on the same rows; per-block smem @@ -798,7 +990,10 @@ __global__ void nh_hot_select_kernel( for (int k = 0; k < K; k++) { int bv = 0, bg = -1; for (int g = tid; g < NH_GLYPH_VOCAB; g += blockDim.x) - if (counts[g] > bv) { bv = counts[g]; bg = g; } + if (counts[g] > bv) { + bv = counts[g]; + bg = g; + } best_v[tid] = bv; best_g[tid] = bg; __syncthreads(); for (int off = blockDim.x / 2; off > 0; off >>= 1) { @@ -824,7 +1019,7 @@ __global__ void nh_dE_scatter_kernel( const float* __restrict__ gidx, const int* __restrict__ hot_map, const int* __restrict__ hot_list, const int* __restrict__ hot_n, int64_t ncell) { - extern __shared__ long long acc_s[]; // NH_HOT_T x NH_EMBED_DIM + extern __shared__ long long acc_s[]; // NH_HOT_T x NH_EMBED_DIM for (int i = threadIdx.x; i < NH_HOT_T * NH_EMBED_DIM; i += blockDim.x) acc_s[i] = 0; __syncthreads(); @@ -859,7 +1054,7 @@ __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) { - extern __shared__ long long acc_s[]; // NH_HOT_G x NH_TROW + 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; __syncthreads(); @@ -934,51 +1129,62 @@ __global__ void nh_count_pad_kernel(int* __restrict__ counts, int B) { counts[NH_PAD_GLYPH] += NH_PAD_PER_SAMPLE * B; } -// ---- Nethack encoder structs ---- +// 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, inv2_w, inv2_b; + Prec inv1_w, inv1_b, inv1s_w, invt_w, inv2_w, inv2_b; Prec bl_w, bl_b, proj_w, proj_b; - Prec msg_w; // trigram embedding table (NH_MSG_VOCAB, NH_MSG_HID) + 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) int obs_size, hidden; }; 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) - Float inv_idx; // inventory slot glyph ids - Prec inv_sfeat; // per-slot state features (B, 55*NH_SFEAT) - Prec inv_T, inv_out; // fused inv table + relu'd flat slots - Prec inv_pool; // pooled 128-dim trunk summary - Int inv_amax; // winning slot per (sample, pool dim) + 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) + 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) + 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) + 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 dTinv, dE_tmp; // inv-table grad + its dE staging - Long dT_i, dTinv_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 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 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 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, inv2_wgrad, inv2_bgrad; + Prec inv1_wgrad, inv1_bgrad, inv1s_wgrad, invt_wgrad, inv2_wgrad, inv2_bgrad; Prec bl_wgrad, bl_bgrad, proj_wgrad, proj_bgrad; - Prec msg_wgrad; + Prec msg_wgrad, spk_wgrad, spk2_wgrad, spk2_bgrad; }; static NethackEncoderWeights* nethack_encoder_create(int obs_size, int hidden) { @@ -993,9 +1199,9 @@ static NethackEncoderWeights* nethack_encoder_create(int obs_size, int hidden) { return ew; } -// ---- Nethack encoder interface ---- +// encoder interface -// ---- encoder <-> pointer-decoder wiring ---- +// encoder <-> pointer-decoder wiring // The decoder's slot head is a pointer over the inventory branch's per-slot // vectors: it reads the encoder's inv_out (keys), and the encoder backward // adds the decoder's key gradients into the inv slice (inv_out has two grad @@ -1006,7 +1212,8 @@ struct NethackDecoderActivations; // each arch_reg_* call; the decoder captures this at its own reg time, so // every rollout buffer's decoder reads its own buffer's inv_out. static NethackEncoderActivations* nh_enc_last = NULL; -static Prec* nh_ptr_keygrad = NULL; // train decoder's (B_TT, NH_INV_FLAT) +static Prec* nh_ptr_keygrad = NULL; // train decoder's (B_TT, NH_INV_FLAT) +static Prec* nh_ptr_spkeygrad = NULL; // train decoder's spell-key grads (B_TT, 8*NH_SPKEY) static Prec nethack_encoder_forward(void* w, void* activations, Prec input, cudaStream_t stream) { NethackEncoderWeights* ew = (NethackEncoderWeights*)w; @@ -1035,14 +1242,18 @@ static Prec nethack_encoder_forward(void* w, void* activations, Prec input, cuda ew->glb2_b.data, a->glyph_idx.data, input.data, B); nh_inv_decode_kernel<<>>( - a->inv_idx.data, input.data, B); + a->inv_idx.data, input.data, B, NH_INV_OFF); + nh_inv_decode_kernel<<>>( + a->invt_idx.data, input.data, B, NH_INVTRUE_OFF); nh_inv_table_kernel<<>>( a->inv_T.data, a->e_eff.data, ew->inv1_w.data); + nh_inv_table_kernel<<>>( + a->invt_T.data, a->e_eff.data, ew->invt_w.data); nh_inv_sfeat_kernel<<>>( a->inv_sfeat.data, input.data, B); 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, B); + 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); @@ -1061,6 +1272,12 @@ static Prec nethack_encoder_forward(void* w, void* activations, Prec input, cuda nh_concat_kernel<<>>( a->concat.data, a->loc_out.data, a->glb_out.data, a->inv_pool.data, a->bl_out.data, a->bl_feats.data, a->msg_out.data, 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); 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); @@ -1109,7 +1326,7 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda 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}}; + Prec dxyv = {.data = a->dxy.data, .shape = {B * NH_TOK, 2}}; puf_mm_tn(&dt16v, &dxyv, &a->glb1_xygrad, stream); // Inventory branch: slice the pooled-summary grad, relu-mask it (inv2 @@ -1142,10 +1359,17 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda 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); nh_dTinv_scatter_kernel<<>>( - (long long*)a->dTinv_i.data, a->inv_grad.data, a->inv_idx.data, (int64_t)B * NH_INV_FLAT); + (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); 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); + 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); + 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<<>>( @@ -1182,8 +1406,8 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda 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 + 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); @@ -1208,6 +1432,29 @@ static void nethack_encoder_backward(void* w, void* activations, Prec grad, cuda puf_mm_nn(&a->dTinv, &ew->inv1_w, &a->dE_tmp, stream); nh_add_inplace_kernel<<>>( a->embed_wgrad.data, a->dE_tmp.data, dE_n); + // ...and the discovered-type channel's: dE += dT_true @ invt_w. + puf_mm_nn(&a->dTtrue, &ew->invt_w, &a->dE_tmp, stream); + nh_add_inplace_kernel<<>>( + a->embed_wgrad.data, a->dE_tmp.data, dE_n); + + // 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); + 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}}; + Prec inf = {.data = a->spk_in.data, .shape = {B * NH_SPELL_SLOTS, NH_SPIN}}; + puf_mm_tn(&dkf, &inf, &a->spk_wgrad, stream); } + 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 @@ -1241,38 +1488,47 @@ static void nethack_encoder_init_weights(void* w, uint64_t* seed, cudaStream_t s 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 + puf_normal_init(&ew->msg_w, 1.0f, (*seed)++, stream); // trigram embedding + 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); } // Param and grad registration orders must match pairwise (muon walks both flat). static void nethack_encoder_reg_params(void* w, Allocator* alloc) { NethackEncoderWeights* ew = (NethackEncoderWeights*)w; 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->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->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}}; // 32x17=544, mult of 8 - ew->inv2_w = {.shape = {NH_INV_POOL, NH_INV_HID}}; - ew->inv2_b = {.shape = {NH_INV_POOL}}; - 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->glb1_b = {.shape = {NH_P1}}; + ew->glb2_w = {.shape = {NH_GLB_HID, NH_P1}}; + ew->glb2_b = {.shape = {NH_GLB_HID}}; + 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}}; + 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 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); @@ -1280,76 +1536,95 @@ static void nethack_encoder_reg_params(void* w, Allocator* alloc) { alloc_register(alloc,&ew->glb1_b); alloc_register(alloc,&ew->glb2_w); alloc_register(alloc,&ew->glb2_b); alloc_register(alloc,&ew->inv1_w); alloc_register(alloc,&ew->inv1_b); - alloc_register(alloc,&ew->inv1s_w); + 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); } static void nethack_encoder_reg_train(void* w, void* activations, Allocator* acts, Allocator* grads, int B_TT) { NethackEncoderWeights* ew = (NethackEncoderWeights*)w; NethackEncoderActivations* a = (NethackEncoderActivations*)activations; *a = {}; - a->glyph_idx = {.shape = {B_TT, NH_MGRID}}; + a->glyph_idx = {.shape = {B_TT, NH_MGRID}}; 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->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->inv_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->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->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}}; - a->msg_ids = {.shape = {B_TT, NH_MSG_LEN}}; - a->msg_out = {.shape = {B_TT, NH_MSG_HID}}; - a->concat = {.shape = {B_TT, NH_CONCAT}}; - a->out = {.shape = {B_TT, ew->hidden}}; + 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_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->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}}; + a->msg_ids = {.shape = {B_TT, NH_MSG_LEN}}; + a->msg_out = {.shape = {B_TT, NH_MSG_HID}}; + 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}}; + 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->inv_idx); alloc_register(acts,&a->inv_sfeat); - alloc_register(acts,&a->inv_T); alloc_register(acts,&a->inv_out); + 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); + 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->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->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->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->dmsg_acc = {.shape = {NH_MSG_VOCAB * NH_MSG_HID}}; + 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->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->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->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}}; 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); 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); @@ -1358,24 +1633,28 @@ static void nethack_encoder_reg_train(void* w, void* activations, Allocator* act 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->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->inv1_wgrad = {.shape = {NH_INV_HID, NH_EMBED_DIM}}; - a->inv1_bgrad = {.shape = {NH_INV_HID}}; + a->glb1_bgrad = {.shape = {NH_P1}}; + a->glb2_wgrad = {.shape = {NH_GLB_HID, NH_P1}}; + a->glb2_bgrad = {.shape = {NH_GLB_HID}}; + 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->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->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}}; 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); @@ -1383,52 +1662,66 @@ static void nethack_encoder_reg_train(void* w, void* activations, Allocator* act alloc_register(grads,&a->glb1_bgrad); alloc_register(grads,&a->glb2_wgrad); alloc_register(grads,&a->glb2_bgrad); alloc_register(grads,&a->inv1_wgrad); alloc_register(grads,&a->inv1_bgrad); - alloc_register(grads,&a->inv1s_wgrad); + 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); nh_enc_last = a; } static void nethack_encoder_reg_rollout(void* w, void* activations, Allocator* alloc, int B) { NethackEncoderWeights* ew = (NethackEncoderWeights*)w; NethackEncoderActivations* a = (NethackEncoderActivations*)activations; - a->glyph_idx = {.shape = {B, NH_MGRID}}; + a->glyph_idx = {.shape = {B, NH_MGRID}}; 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->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->inv_idx = {.shape = {B, NH_INV}}; - a->inv_sfeat = {.shape = {B, NH_INV * NH_SFEAT}}; - a->inv_T = {.shape = {NH_GLYPH_VOCAB, 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->glb_out = {.shape = {B, NH_GLB_HID}}; - a->bl_feats = {.shape = {B, NH_BL_FEAT}}; - a->bl_out = {.shape = {B, NH_BL_HID}}; - a->msg_ids = {.shape = {B, NH_MSG_LEN}}; - a->msg_out = {.shape = {B, NH_MSG_HID}}; - a->concat = {.shape = {B, NH_CONCAT}}; - a->out = {.shape = {B, ew->hidden}}; + 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_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->glb_out = {.shape = {B, NH_GLB_HID}}; + a->bl_feats = {.shape = {B, NH_BL_FEAT}}; + a->bl_out = {.shape = {B, NH_BL_HID}}; + a->msg_ids = {.shape = {B, NH_MSG_LEN}}; + 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}}; + 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->inv_idx); alloc_register(alloc,&a->inv_sfeat); - alloc_register(alloc,&a->inv_T); alloc_register(alloc,&a->inv_out); + alloc_register(alloc,&a->inv_idx); alloc_register(alloc,&a->invt_idx); + alloc_register(alloc,&a->spell_idx); + 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->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->concat); alloc_register(alloc,&a->out); nh_enc_last = a; } @@ -1452,7 +1745,7 @@ static void create_nethack_encoder(Encoder* enc) { }; } -// ---- Nethack decoder: per-verb pointer slot heads ---- +// decoder: per-verb pointer slot heads // Output layout matches DefaultDecoder: [14 verb | 5x55 slots | 8 dir | value]. // verb/dir/value are one small linear. Each item verb (wear/eat/quaff/throw/ // zap) owns a query q_h = W_qh . hidden; slot logit i = exp(ltau_h) * cos(q_h, k_i) @@ -1462,13 +1755,18 @@ static void create_nethack_encoder(Encoder* enc) { // shared across heads, so the item->action mapping is position-invariant and // every item use trains the same projections. -static constexpr int NH_DIRS = 8; -static constexpr int NH_HEADS = 12; // wear|eat|quaff|throw|zap|takeoff|puton|remove|wield|apply|read|drop -static constexpr int NH_SLOT_OD = NH_HEADS * NH_INV; // 660 slot logits -static constexpr int NH_DEC_OD = NH_ACTIONS + NH_SLOT_OD + NH_DIRS; // 690 logits -static constexpr int NH_DEC_LIN = NH_ACTIONS + NH_DIRS + 1; // 31 verb|dir|value rows -static constexpr int NH_DEC_PAD = 32; // lin rows padded to mult of 8 (cublasLt alignment) -static constexpr int NH_QDIM = NH_HEADS * NH_INV_HID; // stacked queries +static constexpr int NH_DIRS = 8; +static constexpr int NH_DIRHEADS = 6; // move|run|kick|throw|zap|apply +static constexpr int NH_HEADS = 12; // wear|eat|quaff|throw|zap|takeoff|puton|remove|wield|apply|read|drop +static constexpr int NH_SLOT_OD = NH_HEADS * NH_INV; // 660 slot logits +static constexpr int NH_DEC_OD = NH_ACTIONS + NH_SLOT_OD + NH_DIRHEADS * NH_DIRS + + NH_SPELL_SLOTS; // logits: verbs|slots|dirs|spell +static constexpr int NH_DEC_LIN = NH_ACTIONS + NH_DIRHEADS * NH_DIRS + 1; // verbs|dirs|value +static constexpr int NH_DEC_PAD = (NH_DEC_LIN + 7) / 8 * 8; // lin rows padded to mult of 8 (cublasLt alignment) +// queries: 12 cosine inv heads + 1 dot-product spell head (row NH_HEADS) +static constexpr int NH_QHEADS = NH_HEADS + 1; +static constexpr int NH_QDIM = NH_QHEADS * NH_INV_HID; // stacked queries +static constexpr int NH_SPELL_BASE = NH_ACTIONS + NH_SLOT_OD + NH_DIRHEADS * NH_DIRS; // tau is padded to 8 entries (first NH_HEADS live): checkpoints are saved // compactly and the puffernet loader assumes every tensor is a multiple of // 8 floats (16-byte bf16 alignment). Pad slots are dead but NOT frozen @@ -1484,30 +1782,32 @@ struct NethackDecoderWeights { int hidden_dim, output_dim; bool continuous; // pointer-head weights (v3: per-verb queries, shared cosine keys) - Prec lin_w; // (NH_DEC_PAD rows, hidden); first NH_DEC_LIN used - Prec q_w; // (NH_QDIM, hidden) stacked per-head query projections - Prec k_w; // (NH_INV_HID, NH_INV_HID) key projection over inv features - Prec tau; // (NH_TAU_PAD,) learnable LOG temperatures, first NH_HEADS live + Prec lin_w; // (NH_DEC_PAD rows, hidden); first NH_DEC_LIN used + Prec q_w; // (NH_QDIM, hidden) stacked per-head query projections + Prec k_w; // (NH_INV_HID, NH_INV_HID) key projection over inv features + Prec tau; // (NH_TAU_PAD,) learnable LOG temperatures, first NH_HEADS live }; struct NethackDecoderActivations { - NethackEncoderActivations* enc; // partner encoder acts (keys source) - Prec out; // (B, NH_DEC_OD+1) - Prec tmp, q; // (B, NH_DEC_PAD), (B, NH_QDIM) + NethackEncoderActivations* enc; // partner encoder acts (keys source) + Prec out; // (B, NH_DEC_OD+1) + Prec tmp, q; // (B, NH_DEC_PAD), (B, NH_QDIM) Prec saved_input, grad_input, grad_input2; - Prec grad_out; // assembled logits+value grad + Prec grad_out; // assembled logits+value grad 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_HEADS) - Prec slot_logits; // (B, NH_SLOT_OD) tau_h * cos - Prec dkmat; // backward scratch - Long tau_acc; // fixed-point dtau staging (NH_TAU_PAD,) + 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 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) + Long tau_acc; // fixed-point dtau staging (NH_TAU_PAD,) Prec lin_wgrad, q_wgrad, k_wgrad, tau_grad; }; __global__ void nh_dec_assemble_kernel(precision_t* __restrict__ out, - const precision_t* __restrict__ tmp, const precision_t* __restrict__ slot_logits, int B) { + const precision_t* __restrict__ tmp, const precision_t* __restrict__ slot_logits, + const precision_t* __restrict__ q, const precision_t* __restrict__ spkeys, int B) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int od1 = NH_DEC_OD + 1; if (idx >= B * od1) return; @@ -1516,7 +1816,18 @@ __global__ void nh_dec_assemble_kernel(precision_t* __restrict__ out, if (c < NH_ACTIONS) v = to_float(tmp[(int64_t)b * NH_DEC_PAD + c]); else if (c < NH_ACTIONS + NH_SLOT_OD) v = to_float(slot_logits[(int64_t)b * NH_SLOT_OD + (c - NH_ACTIONS)]); - else v = to_float(tmp[(int64_t)b * NH_DEC_PAD + NH_ACTIONS + (c - NH_ACTIONS - NH_SLOT_OD)]); + else if (c < NH_SPELL_BASE) // per-verb dir rows from lin + v = to_float(tmp[(int64_t)b * NH_DEC_PAD + NH_ACTIONS + (c - NH_ACTIONS - NH_SLOT_OD)]); + else if (c < NH_SPELL_BASE + NH_SPELL_SLOTS) { + // spell head: dot(q_spell, key_s) / sqrt(keydim) + int s = c - NH_SPELL_BASE; + const precision_t* qs = q + ((int64_t)b * NH_QHEADS + NH_HEADS) * NH_INV_HID; + const precision_t* ks = spkeys + ((int64_t)b * NH_SPELL_SLOTS + s) * NH_SPKEY; + float dot = 0.0f; + for (int k = 0; k < NH_SPKEY; k++) dot += to_float(qs[k]) * to_float(ks[k]); + v = dot * 0.25f; + } else // value + v = to_float(tmp[(int64_t)b * NH_DEC_PAD + NH_DEC_LIN - 1]); out[idx] = from_float(v); } @@ -1543,13 +1854,13 @@ __global__ void nh_ptr3_cos_kernel(precision_t* __restrict__ slot_logits, if (idx >= B * NH_SLOT_OD) return; int b = idx / NH_SLOT_OD, hi = idx % NH_SLOT_OD; int h = hi / NH_INV, i = hi % NH_INV; - const precision_t* qb = q + ((int64_t)b * NH_HEADS + h) * NH_INV_HID; + const precision_t* qb = q + ((int64_t)b * NH_QHEADS + h) * NH_INV_HID; const precision_t* ki = kmat + ((int64_t)b * NH_INV + i) * NH_INV_HID; float dot = 0.0f; 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_HEADS + h]) * to_float(kn[(int64_t)b * NH_INV + i]))); + (to_float(qn[(int64_t)b * NH_QHEADS + h]) * to_float(kn[(int64_t)b * NH_INV + i]))); } __global__ void nh_dec_dtmp_kernel(precision_t* __restrict__ dtmp, @@ -1557,8 +1868,14 @@ __global__ void nh_dec_dtmp_kernel(precision_t* __restrict__ dtmp, int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= B * NH_DEC_PAD) return; int b = idx / NH_DEC_PAD, c = idx % NH_DEC_PAD; - if (c >= NH_DEC_LIN) { dtmp[idx] = from_float(0.0f); return; } // pad rows - int src = c < NH_ACTIONS ? c : NH_ACTIONS + NH_SLOT_OD + (c - NH_ACTIONS); + if (c >= NH_DEC_LIN) { // pad rows + dtmp[idx] = from_float(0.0f); + return; + } + // lin rows: verbs, dirs, then value (spell logits are pointer-derived) + int src = c < NH_ACTIONS ? c + : c < NH_DEC_LIN - 1 ? NH_ACTIONS + NH_SLOT_OD + (c - NH_ACTIONS) + : NH_DEC_OD; dtmp[idx] = g[(int64_t)b * (NH_DEC_OD + 1) + src]; } @@ -1573,11 +1890,12 @@ __global__ void nh_ptr3_dq_kernel(precision_t* __restrict__ dq, int bh = blockIdx.x * blockDim.x + threadIdx.x; if (bh >= B * NH_HEADS) return; int b = bh / NH_HEADS, h = bh % NH_HEADS; + int64_t qrow = (int64_t)b * NH_QHEADS + h; float tauv = expf(to_float(tau[h])); - float qnv = to_float(qn[bh]); + float qnv = to_float(qn[qrow]); float vhat[NH_INV_HID], dv[NH_INV_HID]; for (int k = 0; k < NH_INV_HID; k++) { - vhat[k] = to_float(q[(int64_t)bh * NH_INV_HID + k]) / qnv; + vhat[k] = to_float(q[qrow * NH_INV_HID + k]) / qnv; dv[k] = 0.0f; } float dtau = 0.0f; @@ -1595,10 +1913,30 @@ __global__ void nh_ptr3_dq_kernel(precision_t* __restrict__ dq, float vdv = 0.0f; for (int k = 0; k < NH_INV_HID; k++) vdv += vhat[k] * dv[k]; for (int k = 0; k < NH_INV_HID; k++) - dq[(int64_t)bh * NH_INV_HID + k] = from_float((dv[k] - vhat[k] * vdv) / qnv); + dq[qrow * NH_INV_HID + k] = from_float((dv[k] - vhat[k] * vdv) / qnv); if (dtau != 0.0f) nh_fxp_atomic_add(&tau_acc[h], dtau * tauv); } +// spell head backward (dot-product pointer): dq_spell = 0.25 sum_s g_s k_s, +// spdk_s = 0.25 g_s q_spell (+ the pool grad, added encoder-side). Always +// writes its dq row and all spdk entries (zeros when CAST saw no gradient). +__global__ void nh_spq_bwd_kernel(precision_t* __restrict__ dq, + precision_t* __restrict__ spdk, const precision_t* __restrict__ g, + const precision_t* __restrict__ q, const precision_t* __restrict__ spkeys, int B) { + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= B * NH_SPKEY) return; + int b = t / NH_SPKEY, k = t % NH_SPKEY; + int64_t qrow = ((int64_t)b * NH_QHEADS + NH_HEADS) * NH_INV_HID; + float qk = to_float(q[qrow + k]); + float dqk = 0.0f; + for (int s = 0; s < NH_SPELL_SLOTS; s++) { + float gs = to_float(g[(int64_t)b * (NH_DEC_OD + 1) + NH_SPELL_BASE + s]); + dqk += 0.25f * gs * to_float(spkeys[((int64_t)b * NH_SPELL_SLOTS + s) * NH_SPKEY + k]); + spdk[((int64_t)b * NH_SPELL_SLOTS + s) * NH_SPKEY + k] = from_float(0.25f * gs * qk); + } + dq[qrow + k] = from_float(dqk); +} + // dk_i = sum_h tau_h * g_hi * (v_h - u_i * cos_hi) / ||k_i|| __global__ void nh_ptr3_dkmat_kernel(precision_t* __restrict__ dkmat, const precision_t* __restrict__ g, const precision_t* __restrict__ out, @@ -1620,8 +1958,8 @@ __global__ void nh_ptr3_dkmat_kernel(precision_t* __restrict__ dkmat, 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_HEADS + h) * NH_INV_HID + k]) - / to_float(qn[(int64_t)b * NH_HEADS + h]); + 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); } dkmat[idx] = from_float(acc / knv); @@ -1638,14 +1976,14 @@ static Prec nethack_decoder_forward(void* w, void* activations, Prec input, cuda Prec sflat = {.data = ea->inv_out.data, .shape = {B * NH_INV, NH_INV_HID}}; Prec kflat = {.data = a->kmat.data, .shape = {B * NH_INV, NH_INV_HID}}; puf_mm(&sflat, &dw->k_w, &kflat, stream); - nh_ptr_rownorm_kernel<<>>( - a->qn.data, a->q.data, B * NH_HEADS); + 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); nh_dec_assemble_kernel<<>>( - a->out.data, a->tmp.data, a->slot_logits.data, B); + a->out.data, a->tmp.data, a->slot_logits.data, a->q.data, ea->spk_keys.data, B); return a->out; } @@ -1666,6 +2004,8 @@ static Prec nethack_decoder_backward(void* w, void* activations, a->q.data, a->qn.data, a->kmat.data, a->kn.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); @@ -1700,9 +2040,9 @@ static void nethack_decoder_init_weights(void* w, uint64_t* seed, cudaStream_t s static void nethack_decoder_reg_params(void* w, Allocator* alloc) { NethackDecoderWeights* dw = (NethackDecoderWeights*)w; dw->lin_w = {.shape = {NH_DEC_PAD, dw->hidden_dim}}; - dw->q_w = {.shape = {NH_QDIM, dw->hidden_dim}}; - dw->k_w = {.shape = {NH_INV_HID, NH_INV_HID}}; - dw->tau = {.shape = {NH_TAU_PAD}}; + dw->q_w = {.shape = {NH_QDIM, dw->hidden_dim}}; + dw->k_w = {.shape = {NH_INV_HID, NH_INV_HID}}; + dw->tau = {.shape = {NH_TAU_PAD}}; alloc_register(alloc,&dw->lin_w); alloc_register(alloc,&dw->q_w); alloc_register(alloc,&dw->k_w); @@ -1713,26 +2053,27 @@ static void nethack_decoder_reg_train(void* w, void* activations, Allocator* act NethackDecoderWeights* dw = (NethackDecoderWeights*)w; NethackDecoderActivations* a = (NethackDecoderActivations*)activations; *a = {}; - a->out = {.shape = {B_TT, NH_DEC_OD + 1}}; - a->tmp = {.shape = {B_TT, NH_DEC_PAD}}; - a->q = {.shape = {B_TT, NH_QDIM}}; + a->out = {.shape = {B_TT, NH_DEC_OD + 1}}; + a->tmp = {.shape = {B_TT, NH_DEC_PAD}}; + a->q = {.shape = {B_TT, NH_QDIM}}; a->saved_input = {.shape = {B_TT, dw->hidden_dim}}; - a->grad_input = {.shape = {B_TT, dw->hidden_dim}}; + a->grad_input = {.shape = {B_TT, dw->hidden_dim}}; a->grad_input2 = {.shape = {B_TT, dw->hidden_dim}}; - a->grad_out = {.shape = {B_TT, NH_DEC_OD + 1}}; - a->dtmp = {.shape = {B_TT, NH_DEC_PAD}}; - 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_HEADS}}; + a->grad_out = {.shape = {B_TT, NH_DEC_OD + 1}}; + a->dtmp = {.shape = {B_TT, NH_DEC_PAD}}; + 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}}; - a->tau_acc = {.shape = {NH_TAU_PAD}}; - a->lin_wgrad = {.shape = {NH_DEC_PAD, dw->hidden_dim}}; - a->q_wgrad = {.shape = {NH_QDIM, dw->hidden_dim}}; - a->k_wgrad = {.shape = {NH_INV_HID, NH_INV_HID}}; - a->tau_grad = {.shape = {NH_TAU_PAD}}; + a->dkmat = {.shape = {B_TT, NH_INV_FLAT}}; + a->spdk = {.shape = {B_TT, NH_SPELL_SLOTS * NH_SPKEY}}; + a->tau_acc = {.shape = {NH_TAU_PAD}}; + a->lin_wgrad = {.shape = {NH_DEC_PAD, dw->hidden_dim}}; + a->q_wgrad = {.shape = {NH_QDIM, dw->hidden_dim}}; + a->k_wgrad = {.shape = {NH_INV_HID, NH_INV_HID}}; + a->tau_grad = {.shape = {NH_TAU_PAD}}; alloc_register(acts,&a->out); alloc_register(acts,&a->tmp); alloc_register(acts,&a->q); alloc_register(acts,&a->saved_input); alloc_register(acts,&a->grad_input); alloc_register(acts,&a->grad_input2); @@ -1741,22 +2082,24 @@ static void nethack_decoder_reg_train(void* w, void* activations, Allocator* act alloc_register(acts,&a->kmat); alloc_register(acts,&a->kn); 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); alloc_register(grads,&a->lin_wgrad); alloc_register(grads,&a->q_wgrad); alloc_register(grads,&a->k_wgrad); alloc_register(grads,&a->tau_grad); a->enc = nh_enc_last; nh_ptr_keygrad = &a->keygrad; + nh_ptr_spkeygrad = &a->spdk; } static void nethack_decoder_reg_rollout(void* w, void* activations, Allocator* alloc, int B) { - (void)w; // rollout shapes are all compile-time constants + (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_HEADS}}; + 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); diff --git a/ocean/nethack/nethack.h b/ocean/nethack/nethack.h index ead5871b2c..220bfba60a 100644 --- a/ocean/nethack/nethack.h +++ b/ocean/nethack/nethack.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -26,7 +27,19 @@ extern "C" { extern nle_ctx_t* nle_start(nle_obs*, FILE*, nle_settings*); extern nle_ctx_t* nle_step(nle_ctx_t*, nle_obs*); extern nle_ctx_t* nle_obs_refresh(nle_ctx_t*, nle_obs*); -extern void nle_end(nle_ctx_t*); +extern int nle_path_drain(nle_ctx_t*, short*, int); +extern long nle_shop_price(nle_ctx_t*); +extern int nle_terrain_underfoot(nle_ctx_t*); +extern int nle_inside_shop(nle_ctx_t*); +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 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 void nle_end(nle_ctx_t*); #ifdef __cplusplus } #endif @@ -34,12 +47,14 @@ extern void nle_end(nle_ctx_t*); #include "netlib.h" #define OBS_SIZE NETHACK_OBS_SIZE -#define NUM_ATNS 14 +#define NUM_ATNS 20 #define ACT_SIZES {NETHACK_NUM_ACTIONS, \ NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, \ NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, \ NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, NETHACK_INV_SLOTS, \ - NETHACK_NUM_DIRS} + NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, \ + NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, NETHACK_NUM_DIRS, \ + NETHACK_SPELL_SLOTS} typedef unsigned char obs_t; typedef Env Nethack; @@ -48,7 +63,7 @@ struct Env { Agent agents[1]; unsigned char* action_mask; int num_agents; - int pending_reset; // NLE's coroutine must reset on a stepping thread + int pending_reset; // NLE's coroutine must reset on a stepping thread int tag; int boundary_reached; @@ -59,16 +74,23 @@ struct Env { char vardir[1024]; // NLE-written buffers - short glyphs[NH_GRID]; - long blstats[NLE_BLSTATS_SIZE]; - unsigned char chars[NH_GRID]; - unsigned char message[NLE_MESSAGE_SIZE]; - int misc[NLE_MISC_SIZE]; - int internal[NLE_INTERNAL_SIZE]; - short inv_glyphs[NLE_INVENTORY_SIZE]; - unsigned char inv_letters[NLE_INVENTORY_SIZE]; - unsigned char inv_oclasses[NLE_INVENTORY_SIZE]; - signed char inv_state[NLE_INVENTORY_SIZE * NLE_INV_STATE_FIELDS]; + short glyphs[NH_GRID]; + long blstats[NLE_BLSTATS_SIZE]; + unsigned char chars[NH_GRID]; + unsigned char message[NLE_MESSAGE_SIZE]; + int misc[NLE_MISC_SIZE]; + int internal[NLE_INTERNAL_SIZE]; + short inv_glyphs[NLE_INVENTORY_SIZE]; + short inv_true[NLE_INVENTORY_SIZE]; + unsigned char inv_letters[NLE_INVENTORY_SIZE]; + unsigned char inv_oclasses[NLE_INVENTORY_SIZE]; + signed char inv_state[NLE_INVENTORY_SIZE * NLE_INV_STATE_FIELDS]; + short spell_ids[8]; + signed char spell_levs[8], spell_fails[8]; + int spell_knows[8]; // retention turns, 0 = forgotten (slot-faithful) + long stall_prev_turn; // stall watchdog: last seen game turn + int stall_ctr; // consecutive same-turn steps + int n_spells; Stats stats; @@ -79,12 +101,13 @@ struct Env { long prev_exp; long prev_gold; long start_gold; - long prev_hp; - long prev_hunger; + 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; - long prev_ac; - int prev_bad_cond; + 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 // reward coefs float gold_coef; @@ -92,36 +115,41 @@ struct Env { float descent_coef; float xp_coef; float scout_coef; - float hp_coef; - float hunger_coef; - float illegal_penalty; - float death_penalty; float ac_coef; - float heal_coef; - float status_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 - unsigned int rng; // required by vecenv.h - unsigned long seed; // advanced each reset + unsigned int rng; // required by vecenv.h + unsigned long seed; // advanced each reset }; -#include "macros.h" // keystroke utils +#include "macros.h" // init +// demo-only obs planes; NULL in training (fills skipped) +static unsigned char* nethack_color_sink; +static unsigned char* nethack_invstr_sink; + static void nethack_bind_obs(Nethack* env) { nle_obs* o = &env->obs; memset(o, 0, sizeof(*o)); - o->glyphs = env->glyphs; - o->blstats = env->blstats; - o->chars = env->chars; - o->message = env->message; - o->misc = env->misc; + o->colors = nethack_color_sink; + o->inv_strs = nethack_invstr_sink; + o->glyphs = env->glyphs; + o->blstats = env->blstats; + o->chars = env->chars; + o->message = env->message; + o->misc = env->misc; o->internal = env->internal; - o->inv_glyphs = env->inv_glyphs; - o->inv_letters = env->inv_letters; + o->inv_glyphs = env->inv_glyphs; + o->inv_true_glyphs = env->inv_true; + o->inv_letters = env->inv_letters; o->inv_oclasses = env->inv_oclasses; - o->inv_state = env->inv_state; - // partial fills + o->inv_state = env->inv_state; o->partial = 1; } @@ -137,14 +165,16 @@ static void nethack_init_settings(Nethack* env) { strncpy(env->settings.hackdir, env->vardir, sizeof(env->settings.hackdir) - 1); } env->settings.spawn_monsters = 1; - env->settings.underfoot_glyphs = 1; // underfoot shows objects + env->settings.underfoot_glyphs = 1; // underfoot shows objects snprintf(env->settings.options, sizeof(env->settings.options), "@%s", nethack_rc_path(NETHACK_DEFAULT_OPTIONS)); - env->settings.fix_moon_phase = true; // moon phase from seed + env->settings.fix_moon_phase = true; // moon phase from seed } void init(Nethack* env) { - env->seed = 0xCAFEBEEFUL + (unsigned long)env->rng; // rng = env index + env->seed = 0xCAFEBEEFUL + (unsigned long)env->rng; // rng = env index + // opt into consumed-head gating (env_head_consume_map below); =0 still disables + setenv("PUFFER_HEAD_GATING", "1", 0); // nle_start deferred to first puf_reset nethack_init_settings(env); } @@ -153,16 +183,82 @@ void init(Nethack* env) { static int nethack_slot_usable(const Nethack* env, const Verb* verb, int i) { if (!(verb->item_classes & (1u << env->inv_oclasses[i]))) 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))) { + if (env->blstats[NLE_BL_CONDITION] & 0x20L) return 0; + if (env->inv_oclasses[i] == 10 + && env->inv_true[i] != NETHACK_PAD_GLYPH) { + // only the discoveries channel names the spell (glyphs are shuffled + // appearances), exactly when the agent could know it + int otyp = env->inv_true[i] - NH_GLYPH_OBJ_OFF; + for (int j = 0; j < NETHACK_SPELL_SLOTS; j++) + if (env->spell_ids[j] == otyp && env->spell_knows[j] > 2000) + return 0; + } + } + // WIELD while welded (known-cursed wield) refuses for free and the weld + // never expires; the revealing first attempt on unknown BUC stays legal + if (verb == &NETHACK_VERBS[NETHACK_ACT_WIELD]) { + for (int j = 0; j < NETHACK_INV_SLOTS && env->inv_letters[j]; j++) + if ((env->inv_state[j * NLE_INV_STATE_FIELDS + 5] & 2) + && env->inv_state[j * NLE_INV_STATE_FIELDS + 0] == 1) + return 0; + } + // TAKEOFF/REMOVE of known-cursed gear refuses for free; first attempt on + // unknown BUC consumes a move and reveals, so it stays legal + if ((verb == &NETHACK_VERBS[NETHACK_ACT_TAKEOFF] + || verb == &NETHACK_VERBS[NETHACK_ACT_REMOVE]) + && env->inv_state[i * NLE_INV_STATE_FIELDS + 0] == 1) + return 0; int worn = env->inv_state[i * NLE_INV_STATE_FIELDS + 5] & 1; if (verb->wornreq == WORN_ONLY) return worn; - if (verb->wornreq == UNWORN_ONLY) return !worn; + if (verb->wornreq == UNWORN_ONLY) { + if (worn) return 0; + // WEAR: mirror the engine's layering rules and known-cursed same-slot + // swaps (56% of measured prompt aborts) + if (verb->item_classes == (1u << 3)) { + int gn = env->inv_glyphs[i] - NH_GLYPH_OBJ_OFF; + int cat = (gn >= 0 && gn < NH_NUM_OBJECTS) ? nh_obj_armcat[gn] : -1; + if (cat >= 0) { + for (int j = 0; j < NETHACK_INV_SLOTS && env->inv_letters[j]; j++) { + if (!(env->inv_state[j * NLE_INV_STATE_FIELDS + 5] & 1)) continue; + int gj = env->inv_glyphs[j] - NH_GLYPH_OBJ_OFF; + int cj = (gj >= 0 && gj < NH_NUM_OBJECTS) ? nh_obj_armcat[gj] : -1; + if ((cat == 0 || cat == 6) && (cj == 5 || (cat == 6 && cj == 0))) return 0; + if (cj == cat + && env->inv_state[j * NLE_INV_STATE_FIELDS + 0] == 1) return 0; + } + } + } + return 1; + } return 1; } +// visible HOSTILE target (monster/detected/warning glyph) on the ray within +// 8 tiles. A peaceful in the line blocks the shot: ranged attacks have no +// "really attack?" prompt, so aiming at (or through) a watchman starts a war +static int nethack_ray_target(Nethack* env, int dx, int dy) { + long hx = env->blstats[NLE_BL_X], hy = env->blstats[NLE_BL_Y]; + for (int k = 1; k <= 8; k++) { + long x = hx + dx * k, y = hy + dy * k; + if (x < 0 || x >= NH_COLS || y < 0 || y >= NH_ROWS) return 0; + int gl = env->glyphs[y * NH_COLS + x]; + if ((gl >= 0 && gl < NETHACK_NUMMONS) + || (gl >= 762 && gl < 1144) + || (gl >= 5589 && gl < 5595)) + return !nle_peaceful_at(env->ctx, (int)x + 1, (int)y); + } + return 0; +} + static void nethack_compute_mask(Nethack* env) { unsigned char* mask = env->action_mask; memset(mask, 1, NETHACK_NUM_ACTIONS); - if (env->blstats[NLE_BL_HUNGER] == 0) mask[NETHACK_ACT_EAT] = 0; // choke gate + if (env->mask_search20 != 0.0f) mask[NETHACK_ACT_SEARCH20] = 0; + if (env->mask_run != 0.0f) mask[NETHACK_ACT_RUN] = 0; + if (env->blstats[NLE_BL_HUNGER] == 0) mask[NETHACK_ACT_EAT] = 0; // choke gate // underfoot long hero_x = env->blstats[NLE_BL_X], hero_y = env->blstats[NLE_BL_Y]; @@ -172,43 +268,154 @@ static void nethack_compute_mask(Nethack* env) { int on_corpse = (underfoot >= NETHACK_GLYPH_BODY_OFF && underfoot < NETHACK_GLYPH_BODY_OFF + NETHACK_NUMMONS); - if (underfoot != NETHACK_GLYPH_DNSTAIR && underfoot != NETHACK_GLYPH_DNLADDER) + // terrain, not the map glyph: an object on the tile occludes the stairs + // (underfoot_glyphs shows the top item), which used to mask DOWN off and + // strand the agent on a littered staircase + int terrain = nle_terrain_underfoot(env->ctx); + if (terrain != NETHACK_GLYPH_DNSTAIR && terrain != NETHACK_GLYPH_DNLADDER) mask[NETHACK_ACT_DOWN] = 0; - if (underfoot != NETHACK_GLYPH_UPSTAIR && underfoot != NETHACK_GLYPH_UPLADDER) + if (terrain != NETHACK_GLYPH_UPSTAIR && terrain != NETHACK_GLYPH_UPLADDER) mask[NETHACK_ACT_UP] = 0; - if (env->blstats[NLE_BL_DEPTH] <= 1) mask[NETHACK_ACT_UP] = 0; // declined exit + if (env->blstats[NLE_BL_DEPTH] <= 1) mask[NETHACK_ACT_UP] = 0; // declined exit if (!on_object && !on_corpse) mask[NETHACK_ACT_PICKUP] = 0; + if (terrain != NETHACK_GLYPH_ALTAR) mask[NETHACK_ACT_ALTAR_ID] = 0; + // presence is public (the container renders); locked/empty is learnable + if (!nle_container_at(env->ctx)) mask[NETHACK_ACT_TIP] = 0; + + // engrave-test needs an unidentified wand not yet tested this episode: + // one test prints all its information, re-testing just drains charges + int unid_wand = 0; + for (int i = 0; i < NETHACK_INV_SLOTS && env->inv_letters[i]; i++) { + if (env->inv_oclasses[i] != 11 + || env->inv_state[i * NLE_INV_STATE_FIELDS + 6] != 0) continue; + int lb = nethack_letter_bit(env->inv_letters[i]); + if (lb >= 0 && (env->engid_tested & (1ULL << lb))) continue; + unid_wand = 1; + break; + } + if (!unid_wand) mask[NETHACK_ACT_ENGRAVE_ID] = 0; + + // ELBERETH: unengravable terrain (fountain/water/lava/air/cloud) and + // levitating/engulfed both refuse for free + int tg = terrain - 2359; + if (tg == 31 || tg == 32 || tg == 34 || tg == 39 || tg == 40 || tg == 41) + mask[NETHACK_ACT_ELBERETH] = 0; + if ((env->blstats[NLE_BL_CONDITION] & 0x400L) || (env->internal[6] & 4)) + mask[NETHACK_ACT_ELBERETH] = 0; + + // spell-slot head: castable iff still known (retention > 0) and Pw covers + // 5 * level; fail% is deliberately NOT masked (it's in the obs) + unsigned char* sp = mask + NETHACK_NUM_ACTIONS + 12 * NETHACK_INV_SLOTS + + NETHACK_DIR_HEADS * NETHACK_NUM_DIRS; + memset(sp, 0, NETHACK_SPELL_SLOTS); + int castable = 0; + for (int s = 0; s < env->n_spells && s < NETHACK_SPELL_SLOTS; s++) { + if (env->spell_ids[s] > 0 && env->spell_knows[s] > 0 + && env->blstats[NLE_BL_ENE] >= 5L * (long)env->spell_levs[s]) { + sp[s] = 1; + castable = 1; + } + } + if (!castable) sp[0] = 1; // unconsumed head still needs a legal entry + // CAST zero-turn refusal mirror (engine predicate: stun, chant, freehand, + // too-weak, hunger): a free refusal never advances the clock, so the + // blocking condition can never expire -- self-sealing wedge + if (!castable || env->internal[7] <= 10 || nle_cast_blocked(env->ctx)) + mask[NETHACK_ACT_CAST] = 0; + // shop goods we can't pay for: picking them up incurs a bill the agent + // has no way to settle, so gate on affordability (price is quoted to the + // player on arrival, so this is public information) + long shop_price = nle_shop_price(env->ctx); + if (shop_price > env->blstats[NLE_BL_GOLD]) mask[NETHACK_ACT_PICKUP] = 0; // item slot heads for (int a = 0; a < NETHACK_NUM_ACTIONS; a++) { const Verb* verb = &NETHACK_VERBS[a]; - if (verb->head < 0) continue; // direct verb, no item argument + if (verb->head < 0) continue; // direct verb, no item argument unsigned char* slots = mask + NETHACK_NUM_ACTIONS + verb->head * NETHACK_INV_SLOTS; memset(slots, 0, NETHACK_INV_SLOTS); int has_usable = 0; for (int i = 0; i < NETHACK_INV_SLOTS && env->inv_letters[i]; i++) { - if (env->inv_oclasses[i] >= NETHACK_NUM_OCLASSES) break; // padded tail - if (nethack_slot_usable(env, verb, i)) { slots[i] = 1; has_usable = 1; } + if (env->inv_oclasses[i] >= NETHACK_NUM_OCLASSES) break; // padded tail + if (nethack_slot_usable(env, verb, i)) { + slots[i] = 1; + has_usable = 1; + } } if (has_usable) continue; - // no usable item + // no usable item: verb off unless EAT has actual floor food underfoot + // (any-object gating kept EAT legal on inedible piles -- refusal loop) slots[0] = 1; - int floor_food = (a == NETHACK_ACT_EAT) && (on_object || on_corpse); + int floor_food = (a == NETHACK_ACT_EAT) && (on_object || on_corpse) + && nle_food_underfoot(env->ctx); if (!floor_food) mask[a] = 0; } - // directions + // per-verb dir rows: wall legality; THROW/ZAP use target rays (all-1 fallback) unsigned char* dirs = mask + NETHACK_NUM_ACTIONS + 12 * NETHACK_INV_SLOTS; memset(dirs, 1, NETHACK_NUM_DIRS); int legal_dirs = 0; for (int d = 0; d < NETHACK_NUM_DIRS; d++) { long col = hero_x + NETHACK_DIR_DX[d], row = hero_y + NETHACK_DIR_DY[d]; - if (row < 0 || row >= NH_ROWS || col < 0 || col >= NH_COLS) { dirs[d] = 0; continue; } + if (row < 0 || row >= NH_ROWS || col < 0 || col >= NH_COLS) { + dirs[d] = 0; + continue; + } int g = env->glyphs[row * NH_COLS + col]; if (g >= NETHACK_WALL_GLYPH_LO && g <= NETHACK_WALL_GLYPH_HI) dirs[d] = 0; else legal_dirs++; } if (!legal_dirs) memset(dirs, 1, NETHACK_NUM_DIRS); + for (int h = 1; h < NETHACK_DIR_HEADS; h++) + memcpy(dirs + h * NETHACK_NUM_DIRS, dirs, NETHACK_NUM_DIRS); + + // MOVE-head refinements: peaceful-adjacent and diagonal-door moves are void + unsigned char keep[NETHACK_NUM_DIRS]; + memcpy(keep, dirs, NETHACK_NUM_DIRS); + int on_door = (tg >= 12 && tg <= 14); + for (int d = 0; d < NETHACK_NUM_DIRS; d++) { + if (!dirs[d]) continue; + long col = hero_x + NETHACK_DIR_DX[d], row = hero_y + NETHACK_DIR_DY[d]; + if (row < 0 || row >= NH_ROWS || col < 0 || col >= NH_COLS) continue; + int g = env->glyphs[row * NH_COLS + col]; + if (g >= 0 && g < 381 + && nle_peaceful_at(env->ctx, (int)col + 1, (int)row)) { + dirs[d] = 0; + continue; + } + if (d >= 4 && (on_door || (g >= 2371 && g <= 2373))) dirs[d] = 0; + } + int open_moves = 0; + for (int d = 0; d < NETHACK_NUM_DIRS; d++) open_moves |= dirs[d]; + if (!open_moves) memcpy(dirs, keep, NETHACK_NUM_DIRS); // cornered: restore + + // RUN head: a run aimed at an adjacent hostile is void with p=1.0; + // directional only -- escape dirs stay + unsigned char* rdirs = dirs + 1 * NETHACK_NUM_DIRS; + unsigned char rkeep[NETHACK_NUM_DIRS]; + memcpy(rkeep, rdirs, NETHACK_NUM_DIRS); + for (int d = 0; d < NETHACK_NUM_DIRS; d++) { + if (!rdirs[d]) continue; + long col = hero_x + NETHACK_DIR_DX[d], row = hero_y + NETHACK_DIR_DY[d]; + if (row < 0 || row >= NH_ROWS || col < 0 || col >= NH_COLS) continue; + int g = env->glyphs[row * NH_COLS + col]; + if (g >= 0 && g < 381 && !nle_peaceful_at(env->ctx, (int)col + 1, (int)row)) + rdirs[d] = 0; + } + int open_runs = 0; + for (int d = 0; d < NETHACK_NUM_DIRS; d++) open_runs |= rdirs[d]; + if (!open_runs) memcpy(rdirs, rkeep, NETHACK_NUM_DIRS); + + static const int ray_verbs[2] = {NETHACK_ACT_THROW, NETHACK_ACT_ZAP}; + for (int v = 0; v < 2; v++) { + unsigned char* vdirs = dirs + nethack_dir_head(ray_verbs[v]) * NETHACK_NUM_DIRS; + int any = 0; + for (int d = 0; d < NETHACK_NUM_DIRS; d++) { + vdirs[d] = nethack_ray_target(env, NETHACK_DIR_DX[d], NETHACK_DIR_DY[d]); + any |= vdirs[d]; + } + if (!any) memset(vdirs, 1, NETHACK_NUM_DIRS); + } } // observations @@ -224,14 +431,42 @@ static void nethack_pack_obs(Nethack* env) { bl[4*i + 3] = (unsigned char)((v >> 24) & 0xffu); } int32_t extra[NETHACK_EXTRA_INTS] = {0}; - // engraving state 0/1/2 - extra[0] = env->internal[6]; + // engraving state 0/1/2 (bit 2 = engulfed, mask-only -- keep it out of obs) + extra[0] = env->internal[6] & 3; extra[1] = env->prev_action; for (int i = 0; i < NLE_INVENTORY_SIZE; i++) { int oc = env->inv_oclasses[i]; - if (oc >= NETHACK_NUM_OCLASSES) break; // padded tail + if (oc >= NETHACK_NUM_OCLASSES) break; // padded tail extra[2 + oc]++; } + + // 8-slot spell channel with retention (know 0 = forgotten, the re-read cue) + env->n_spells = nle_spells(env->ctx, env->spell_ids, env->spell_levs, + env->spell_fails, env->spell_knows, + NETHACK_SPELL_SLOTS); + extra[NETHACK_EXTRA_SPELL] = env->n_spells; + for (int s = 0; s < NETHACK_SPELL_SLOTS; s++) { + int* q = extra + NETHACK_EXTRA_SPELL + 1 + 4 * s; + int known = s < env->n_spells && env->spell_ids[s] > 0; + q[0] = known ? env->spell_ids[s] : 0; + q[1] = known ? env->spell_levs[s] : 0; + q[2] = known ? env->spell_fails[s] : 0; + q[3] = known ? env->spell_knows[s] : 0; + } + + int wt, wcap; + nle_weight(env->ctx, &wt, &wcap); + if (wcap < 1) wcap = 1; + extra[NETHACK_EXTRA_WEIGHT + 0] = (int)(100L * wt / wcap); + extra[NETHACK_EXTRA_WEIGHT + 1] = wcap; + + // gold/price as a percent, capped at 100 (0 = no purchase available here) + long price = nle_shop_price(env->ctx); + long gold = env->blstats[NLE_BL_GOLD]; + extra[NETHACK_EXTRA_SHOP] = nle_inside_shop(env->ctx); + extra[NETHACK_EXTRA_SHOP + 1] = (price > 0) + ? (int32_t)(gold >= price ? 100 : (gold * 100) / price) : 0; + 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]; @@ -240,16 +475,48 @@ static void nethack_pack_obs(Nethack* env) { ex[4*i + 2] = (unsigned char)((v >> 16) & 0xffu); ex[4*i + 3] = (unsigned char)((v >> 24) & 0xffu); } - // slot glyphs + // slot glyphs; NH_DISC_SWAP=1: discovered identity REPLACES the appearance + // glyph (stable rep once known; the add-channel below goes silent) + static int dsw = -1; + if (dsw < 0) { + const char* e = getenv("NH_DISC_SWAP"); + dsw = e && e[0] && e[0] != '0'; + } 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 - ? (uint16_t)env->inv_glyphs[i] : (uint16_t)NETHACK_PAD_GLYPH; + ? (dsw && env->inv_true[i] != NETHACK_PAD_GLYPH + ? (uint16_t)env->inv_true[i] : (uint16_t)env->inv_glyphs[i]) + : (uint16_t)NETHACK_PAD_GLYPH; iv[2*i + 0] = (unsigned char)(g & 0xffu); iv[2*i + 1] = (unsigned char)((g >> 8) & 0xffu); } // item 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_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; + it[2*i + 0] = (unsigned char)(g & 0xffu); + it[2*i + 1] = (unsigned char)((g >> 8) & 0xffu); + } + // stall watchdog: a key-eating modal (getpos-class) can freeze the game + // turn while swallowing every action; at 96 frozen steps fire an ESC burst + // (ESC is a harmless no-op at the main prompt) + long turn = env->blstats[NLE_BL_TIME]; + if (turn == env->stall_prev_turn && !env->obs.done) { + if (++env->stall_ctr == 96) { + for (int k = 0; k < 8 && !env->obs.done; k++) { + env->obs.action = 27; + env->ctx = nle_step(env->ctx, &env->obs); + } + env->stall_ctr = 0; // re-arm; recovery shows as turn advance + } + } else { + env->stall_ctr = 0; + env->stall_prev_turn = turn; + } // topline unsigned char* mv = ((obs_t*)env->agents[0].observations) + NETHACK_OFF_MSG; size_t mlen = strnlen((const char*)env->message, NETHACK_MSG_LEN); @@ -261,49 +528,46 @@ static void nethack_pack_obs(Nethack* env) { // logging -static void nethack_add_log(Nethack* env, int how) { // how: nle how_done, -1 = truncated +static void nethack_add_log(Nethack* env, int how) { // how: nle how_done, -1 = truncated for (int v = 0; v < NETHACK_NUM_ACTIONS; v++) env->log.verb_uses[v] += (float)env->stats.verb_uses[v]; - env->log.perf += (float)env->prev_score; - env->log.score += (float)env->prev_score; - env->log.valid_moves += (float)env->stats.valid_moves; + env->log.perf += (float)env->prev_score; + env->log.score += (float)env->prev_score; + env->log.valid_moves += (float)env->stats.valid_moves; env->log.illegal_actions += (float)env->stats.illegal_actions; - env->log.new_tiles += (float)env->stats.new_tiles; - env->log.max_depth += (float)env->stats.max_depth; - env->log.enhances += (float)env->stats.enhances; - env->log.prayers_low_hp += (float)env->stats.prayers_low_hp; - env->log.prayers_starving += (float)env->stats.prayers_starving; - env->log.floor_eats += (float)env->stats.floor_eats; - env->log.damage_taken += (float)env->stats.damage; - env->log.ac += env->stats.length > 0 - ? (float)env->stats.ac_sum / (float)env->stats.length : 0.0f; + env->log.new_tiles += (float)env->stats.new_tiles; + env->log.max_depth += (float)env->stats.max_depth; + env->log.floors += (float)env->stats.floors; + env->log.depth_5 += env->stats.max_depth >= 5 ? 1.0f : 0.0f; + env->log.depth_10 += env->stats.max_depth >= 10 ? 1.0f : 0.0f; + 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; + env->log.reads_book += (float)env->stats.reads_book; + env->log.sells += (float)env->stats.sells; + env->log.buys += (float)env->stats.buys; + env->log.discoveries += (float)(nle_discoveries(env->ctx) - env->disc0); env->log.min_ac += (float)env->stats.min_ac; - env->log.armor_swaps += (float)env->stats.armor_swaps; - env->log.heal_hp += (float)env->stats.heal_hp; - env->log.cures += (float)env->stats.cures; env->log.burdened_frac += env->stats.length > 0 ? (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.episode_return += env->stats.ret; - env->log.episode_length += env->stats.length; - if (how == -1) env->log.truncated += 1.0f; - else if (how == 0) env->log.death_combat += 1.0f; - else if (how == 3) env->log.death_starved += 1.0f; - else if (how == NLE_HOW_WRATH) env->log.death_smited += 1.0f; - else env->log.death_other += 1.0f; - // combat anatomy - if (how == 0) { - env->log.death_mon_level += (float)env->internal[NETHACK_INTERNAL_KILLER_MLEV]; - env->log.death_adj_monsters += (float)env->stats.last_adj; - env->log.death_maxhp += (float)env->stats.last_maxhp; - } - env->log.reach_mines += (env->stats.areas & NETHACK_AREA_MINES) ? 1.0f : 0.0f; - env->log.reach_minetown += (env->stats.areas & NETHACK_AREA_MINETOWN) ? 1.0f : 0.0f; - env->log.reach_deep_mines += (env->stats.areas & NETHACK_AREA_DEEP_MINES) ? 1.0f : 0.0f; - env->log.reach_main_d5 += (env->stats.areas & NETHACK_AREA_MAIN_D5) ? 1.0f : 0.0f; - env->log.reach_sokoban += (env->stats.areas & NETHACK_AREA_SOKOBAN) ? 1.0f : 0.0f; - env->log.n += 1.0f; + env->log.game_time += (float)env->prev_time; + env->log.max_xp_level += (float)env->stats.max_xp; + env->log.episode_return += env->stats.ret; + env->log.episode_length += env->stats.length; + if (how == -1) env->log.truncated += 1.0f; + else if (how == 0) env->log.death_combat += 1.0f; + else if (how == 3) env->log.death_starved += 1.0f; + else env->log.death_other += 1.0f; + if (how >= 0 && env->stats.last_hunger >= NETHACK_HUNGER_WEAK) + env->log.death_weak += 1.0f; + if (how == 0) + 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 @@ -315,7 +579,10 @@ static void nethack_do_reset(Nethack* env) { } nethack_bind_obs(env); - env->obs.how_done = -2; // only really_done() sets it + env->obs.how_done = -2; // only really_done() sets it + + // discovered types render as true glyphs on the map (engine opt-in) + setenv("NLE_TRUE_GLYPHS", "1", 1); // seed advance env->seed = env->seed * 6364136223846793005UL + 1442695040888963407UL; @@ -327,26 +594,28 @@ static void nethack_do_reset(Nethack* env) { env->ctx = nle_start(&env->obs, NULL, &env->settings); nethack_drain_prompts(env); - nle_obs_refresh(env->ctx, &env->obs); // full fill: prev_* seeds read blstats + nle_obs_refresh(env->ctx, &env->obs); // full fill: prev_* seeds read blstats env->prev_score = 0; env->prev_exp = env->blstats[NLE_BL_EXP]; env->start_gold = env->blstats[NLE_BL_GOLD]; - env->prev_gold = 0; // clamped net gold - env->prev_hp = env->blstats[NLE_BL_HP]; - env->prev_hunger = env->blstats[NLE_BL_HUNGER]; - if (env->prev_hunger < 1) env->prev_hunger = 1; - else if (env->prev_hunger > 6) env->prev_hunger = 6; + env->prev_gold = 0; // clamped net gold env->prev_depth = (int)env->blstats[NLE_BL_DEPTH]; - env->prev_ac = env->blstats[NLE_BL_AC]; - env->prev_bad_cond = __builtin_popcount((unsigned)env->blstats[NLE_BL_CONDITION] & NETHACK_COND_BAD); env->prev_time = env->blstats[NLE_BL_TIME]; env->prev_action = -1; + short scratch[2 * NETHACK_PATH_MAX]; // discard boot-walk residue + nle_path_drain(env->ctx, scratch, NETHACK_PATH_MAX); + env->prev_floor = (unsigned)(env->blstats[NLE_BL_DNUM] << 8 | env->blstats[NLE_BL_DLEVEL]); + env->disc0 = nle_discoveries(env->ctx); + env->engid_tested = 0; env->enh_ready = 0; memset(&env->stats, 0, sizeof(env->stats)); 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); } @@ -357,61 +626,66 @@ void puf_reset(Nethack* env) { // reward static void nethack_update_stats(Nethack* env) { - int depth = (int)env->blstats[NLE_BL_DEPTH]; - long dnum = env->blstats[NLE_BL_DNUM]; - if (dnum == 2) { - env->stats.areas |= NETHACK_AREA_MINES; - long mlvl = env->blstats[NLE_BL_DLEVEL]; - if (mlvl >= 3) env->stats.areas |= NETHACK_AREA_MINETOWN; - if (mlvl >= 5) env->stats.areas |= NETHACK_AREA_DEEP_MINES; - } - else if (dnum == 0 && depth >= 5) env->stats.areas |= NETHACK_AREA_MAIN_D5; - else if (dnum == 4) env->stats.areas |= NETHACK_AREA_SOKOBAN; - - long hp = env->blstats[NLE_BL_HP]; - if (hp < env->prev_hp) env->stats.damage += env->prev_hp - hp; if (env->blstats[NLE_BL_CAP] > 0) env->stats.burdened_steps++; - - // death anatomy, read back at death - env->stats.last_maxhp = env->blstats[NLE_BL_HPMAX]; - long hx = env->blstats[NLE_BL_X], hy = env->blstats[NLE_BL_Y]; - int adj = 0; - for (int dy = -1; dy <= 1; dy++) - for (int dx = -1; dx <= 1; dx++) { - if (!dx && !dy) continue; - long r = hy + dy, c = hx + dx; - if (r < 0 || r >= NH_ROWS || c < 0 || c >= NH_COLS) continue; - int g = env->glyphs[r * NH_COLS + c]; - if (g >= 0 && g < NETHACK_NUMMONS) adj++; - } - env->stats.last_adj = adj; - + env->stats.last_hunger = (int)env->blstats[NLE_BL_HUNGER]; env->prev_score = env->blstats[NLE_BL_SCORE]; env->prev_time = env->blstats[NLE_BL_TIME]; - env->prev_depth = depth; + env->prev_depth = (int)env->blstats[NLE_BL_DEPTH]; } -static int nethack_first_visit(Nethack* env, int depth, long px, long py) { - if (px < 0 || px >= NH_COLS || py < 0 || py >= NH_ROWS) return 0; - int d = depth < 1 ? 0 : (depth > NETHACK_MAX_DEPTH ? NETHACK_MAX_DEPTH - 1 : depth - 1); - int bit = (int)py * NH_COLS + (int)px; - unsigned char mask = (unsigned char)(1 << (bit & 7)); - if (env->stats.visited[d][bit >> 3] & mask) return 0; - env->stats.visited[d][bit >> 3] |= mask; - return 1; +// 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. +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; + unsigned short key = (unsigned short)(dn << 8 | dl); + int d = -1; + for (int i = 0; i < env->stats.n_visited_floors; i++) { + if (env->stats.visited_key[i] != key) continue; + d = i; + break; + } + if (d < 0) { + if (env->stats.n_visited_floors >= NETHACK_MAX_DEPTH) return 0.0f; + d = env->stats.n_visited_floors++; + 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; } -static float nethack_reward(Nethack* env, int illegal) { +static float nethack_reward(Nethack* env) { // death payout - if (env->obs.done) - return env->death_penalty - env->hp_coef * (float)env->prev_hp; + if (env->obs.done) return env->death_penalty; nethack_update_stats(env); int depth = (int)env->blstats[NLE_BL_DEPTH]; + // term order is the record run's summation order: exp gold descent xp + // scout ac. Float adds are order-sensitive; reordering breaks bit repro. + float r = 0.0f; + // exp, gains only long exp = env->blstats[NLE_BL_EXP]; - float r = exp > env->prev_exp ? env->exp_coef * (float)(exp - env->prev_exp) : 0.0f; + if (exp > env->prev_exp) + r += env->exp_coef * (float)(exp - env->prev_exp); env->prev_exp = exp; // gold, net of start @@ -420,44 +694,22 @@ static float nethack_reward(Nethack* env, int illegal) { r += env->gold_coef * (float)(g - env->prev_gold); env->prev_gold = g; - // descent, max-depth only + // unique-floor stat (no reward) + long dn = env->blstats[NLE_BL_DNUM], dl = env->blstats[NLE_BL_DLEVEL]; + if (dn >= 0 && dn < 16 && dl >= 1 && dl <= 64) { + unsigned long long fb = 1ULL << (dl - 1); + if (!(env->stats.floors_bits[dn] & fb)) { + env->stats.floors_bits[dn] |= fb; + env->stats.floors++; + } + } + + // descent pays per max-depth delta if (depth > env->stats.max_depth) { r += env->descent_coef * (float)(depth - env->stats.max_depth); env->stats.max_depth = depth; } - // hp potential - long hp = env->blstats[NLE_BL_HP]; - long hp_delta = hp - env->prev_hp; - r += env->hp_coef * (float)hp_delta; - // gain-only heal credit - if (hp_delta > 0 && (env->prev_action == NETHACK_ACT_QUAFF - || env->prev_action == NETHACK_ACT_PRAY)) { - r += env->heal_coef * (float)hp_delta; - env->stats.heal_hp += hp_delta; - } - env->prev_hp = hp; - - // ac potential - long ac = env->blstats[NLE_BL_AC]; - r += env->ac_coef * (float)(env->prev_ac - ac); - env->prev_ac = ac; - env->stats.ac_sum += ac; - if ((int)ac < env->stats.min_ac) env->stats.min_ac = (int)ac; - - // status potential - int bad_cond = __builtin_popcount((unsigned)env->blstats[NLE_BL_CONDITION] & NETHACK_COND_BAD); - r += env->status_coef * (float)(env->prev_bad_cond - bad_cond); - if (bad_cond < env->prev_bad_cond) env->stats.cures += env->prev_bad_cond - bad_cond; - env->prev_bad_cond = bad_cond; - - // hunger potential - long hunger = env->blstats[NLE_BL_HUNGER]; - if (hunger < 1) hunger = 1; - else if (hunger > 6) hunger = 6; - r += env->hunger_coef * (float)(env->prev_hunger - hunger); - env->prev_hunger = hunger; - // xp level, max only int xp = (int)env->blstats[NLE_BL_XP]; if (xp > env->stats.max_xp) { @@ -465,13 +717,56 @@ static float nethack_reward(Nethack* env, int illegal) { env->stats.max_xp = xp; } - // scout - if (nethack_first_visit(env, depth, env->blstats[NLE_BL_X], env->blstats[NLE_BL_Y])) { - r += env->scout_coef; - env->stats.new_tiles++; + // 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) + // leaves path coords from the old floor -- skip those. + unsigned floor = (unsigned)(dn << 8 | dl); + short path[2 * NETHACK_PATH_MAX]; + int n = nle_path_drain(env->ctx, path, NETHACK_PATH_MAX); + if (floor != env->prev_floor) n = 0; + env->prev_floor = floor; + + float fresh = 0.0f; + int touched = 0; + for (int i = 0; i < n; i++) { + float c = nethack_tile_claim(env, dn, dl, path[2 * i], path[2 * i + 1]); + fresh += c; + touched += (c > 0.0f); } + // no usable path: non-move verbs or never left the tile + if (!n) { + float c = nethack_tile_claim(env, dn, dl, + env->blstats[NLE_BL_X], env->blstats[NLE_BL_Y]); + fresh += c; + 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 (illegal) r += env->illegal_penalty; return r; } @@ -484,7 +779,8 @@ static void nethack_execute(Nethack* env, int verb, int slot, int dirkey, int* b nethack_send_key(env, dirkey); break; case NETHACK_ACT_RUN: - nethack_send_key(env, dirkey - 32); // uppercase = run + st->verb_uses[verb]++; + nethack_send_key(env, dirkey - 32); // uppercase = run break; case NETHACK_ACT_DOWN: nethack_send_key(env, '>'); @@ -493,7 +789,7 @@ static void nethack_execute(Nethack* env, int verb, int slot, int dirkey, int* b nethack_send_key(env, '<'); break; case NETHACK_ACT_KICK: - nethack_send_key(env, 4); // ^D + nethack_send_key(env, 4); // ^D nethack_answer_direction(env, dirkey); break; case NETHACK_ACT_SEARCH: @@ -506,19 +802,23 @@ static void nethack_execute(Nethack* env, int verb, int slot, int dirkey, int* b break; case NETHACK_ACT_SEARCH20: st->verb_uses[verb]++; + nethack_send_key(env, '2'); + nethack_send_key(env, '0'); nethack_send_key(env, 's'); - //nethack_send_key(env, '2'); - //if (!env->obs.done) nethack_send_key(env, '0'); - //if (!env->obs.done) nethack_send_key(env, 's'); break; - case NETHACK_ACT_PICKUP: + case NETHACK_ACT_PICKUP: { + int purchase = nle_shop_price(env->ctx) > 0; // 0 = own/no-charge pile, nothing to pay st->verb_uses[verb]++; nethack_send_key(env, ','); nethack_answer_menu(env); + // shop pickup bills you; settle it now (the mask guarantees we can) + if (purchase && !env->obs.done) { + nethack_send_key(env, 'p'); + st->buys++; + } break; + } case NETHACK_ACT_PRAY: - if (4 * env->blstats[NLE_BL_HP] <= env->blstats[NLE_BL_HPMAX]) st->prayers_low_hp++; - if (env->blstats[NLE_BL_HUNGER] >= NETHACK_HUNGER_WEAK) st->prayers_starving++; st->verb_uses[verb]++; nethack_send_key(env, 0x80 | 'p'); break; @@ -560,17 +860,69 @@ static void nethack_execute(Nethack* env, int verb, int slot, int dirkey, int* b (otyp == 234 /* PICK_AXE */ || otyp == 50 /* MATTOCK */) ? '>' : dirkey); } break; - case NETHACK_ACT_READ: - if (nethack_item_use(env, 'r', "read", NULL, slot, &st->verb_uses[verb], bad_pick)) + case NETHACK_ACT_READ: { + int oc = env->inv_oclasses[slot]; + if (nethack_item_use(env, 'r', "read", NULL, slot, &st->verb_uses[verb], bad_pick)) { + if (oc == 9) st->reads_scroll++; + else if (oc == 10) st->reads_book++; nethack_answer_menu(env); + } break; + } + case NETHACK_ACT_ALTAR_ID: { + // each drop onto an altar flashes the item's curse state (sets bknown), + // so dump everything still unknown and take it straight back + st->verb_uses[verb]++; + for (int i = 0; i < NETHACK_INV_SLOTS && env->inv_letters[i] && !env->obs.done; i++) { + if (env->inv_oclasses[i] >= NETHACK_NUM_OCLASSES) break; + if (env->inv_oclasses[i] == 12) continue; // gold: no flash + const signed char* st8 = &env->inv_state[i * NLE_INV_STATE_FIELDS]; + if (st8[0] != 0 || (st8[5] & 1)) continue; // BUC known, or worn + nethack_item_use(env, 'd', "drop", NULL, i, NULL, NULL); + } + if (!env->obs.done) { + nethack_send_key(env, ','); + nethack_answer_menu(env); + } + break; + } + case NETHACK_ACT_TIP: + // M('T') = #tip; the "tip it? [ynq]" prompt auto-commits, spillage + // lands underfoot for PICKUP. Locked -> "It's locked." (kick first). + st->verb_uses[verb]++; + nethack_send_key(env, 0x80 | 'T'); + break; + case NETHACK_ACT_ENGRAVE_ID: + st->verb_uses[verb]++; + nethack_do_engrave_id(env); + break; + case NETHACK_ACT_CAST: { + // Z, then 'a' = first known spell; directional spells then prompt and + // take the (ZAP-masked) direction, self-spells resolve immediately + st->verb_uses[verb]++; + nethack_send_key(env, 'Z'); + // proceed only if the spell chooser menu actually opened (xwait): + // a zero-turn refusal (unmirrored rejectcasting case) leaves the + // main prompt, where 'a' would open an apply prompt instead + if (!env->obs.done && env->misc[NETHACK_MISC_XWAIT]) { + nethack_send_key(env, 'a' + (slot >= 0 && slot < NETHACK_SPELL_SLOTS ? slot : 0)); + nethack_answer_direction(env, dirkey); + } + break; + } case NETHACK_ACT_DROP: nethack_item_use(env, 'd', "drop", NULL, slot, &st->verb_uses[verb], bad_pick); break; } } +// compute_mask writes through this flat alias of agents[0].action_mask +static void nethack_sync_buffers(Nethack* env) { + env->action_mask = env->agents[0].action_mask; +} + void puf_step(Nethack* env) { + nethack_sync_buffers(env); // agent pointers are re-dealt between epochs if (env->pending_reset) { env->pending_reset = 0; nethack_do_reset(env); @@ -579,27 +931,29 @@ void puf_step(Nethack* env) { int verb = (int)env->agents[0].actions[0]; int head = NETHACK_VERBS[verb].head; int slot = (head >= 0) ? (int)env->agents[0].actions[1 + head] : 0; - int dirkey = NETHACK_DIR_KEYS[(int)env->agents[0].actions[13]]; + if (verb == NETHACK_ACT_CAST) slot = (int)env->agents[0].actions[19]; + int dh = nethack_dir_head(verb); + int dirkey = NETHACK_DIR_KEYS[dh >= 0 ? (int)env->agents[0].actions[13 + dh] : 0]; long time_before = env->blstats[NLE_BL_TIME]; int bad_pick = 0; nethack_execute(env, verb, slot, dirkey, &bad_pick); env->prev_action = verb; - int illegal = nethack_handle_prompts(env); + nethack_handle_prompts(env); if (!env->obs.done) nle_obs_refresh(env->ctx, &env->obs); nethack_auto_enhance(env); - if (bad_pick) { illegal = 1; env->stats.illegal_actions++; } + if (bad_pick) env->stats.illegal_actions++; if (env->blstats[NLE_BL_TIME] > time_before) env->stats.valid_moves++; env->stats.length++; - float reward = nethack_reward(env, illegal); + float reward = nethack_reward(env); env->agents[0].rewards[0] = reward; env->stats.ret += reward; int done = env->obs.done || env->stats.length >= NETHACK_MAX_EPISODE_STEPS; - env->agents[0].terminals[0] = done ? 1.0f : 0.0f; // truncation reported as terminal too + env->agents[0].terminals[0] = done ? 1.0f : 0.0f; // truncation reported as terminal too if (done) { nethack_add_log(env, env->obs.done ? env->obs.how_done : -1); // eager same-thread reset: the terminal step returns the fresh obs @@ -631,15 +985,13 @@ void puf_render(Nethack* env) { env->blstats[NLE_BL_HP], env->blstats[NLE_BL_HPMAX], env->blstats[NLE_BL_AC], env->blstats[NLE_BL_DEPTH], env->blstats[NLE_BL_SCORE], env->blstats[NLE_BL_TIME]); + if (env->n_spells > 0) + printf("Sp: id%d L%d fail%d%% x%d\n", env->spell_ids[0], + env->spell_levs[0], env->spell_fails[0], env->n_spells); printf("Msg: %.*s\n", NLE_MESSAGE_SIZE, env->message); fflush(stdout); } -static void nethack_sync_buffers(Nethack* env) { - env->agents[0].observations = (unsigned char*)env->agents[0].observations; - env->action_mask = env->agents[0].action_mask; -} - void puf_init(Env* env, Dict* kwargs) { env->num_agents = 1; env->agents[0].policy = 0; @@ -648,14 +1000,13 @@ void puf_init(Env* env, Dict* kwargs) { env->exp_coef = dict_get(kwargs, "exp_coef"); env->descent_coef = dict_get(kwargs, "descent_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->hp_coef = dict_get(kwargs, "hp_coef"); - env->hunger_coef = dict_get(kwargs, "hunger_coef"); - env->illegal_penalty = dict_get(kwargs, "illegal_penalty"); env->death_penalty = dict_get(kwargs, "death_penalty"); - env->ac_coef = dict_get(kwargs, "ac_coef"); - env->heal_coef = dict_get(kwargs, "heal_coef"); - env->status_coef = dict_get(kwargs, "status_coef"); + env->mask_search20 = dict_get(kwargs, "mask_search20"); + env->mask_run = dict_get(kwargs, "mask_run"); } // Export order: outcomes first (score/depth/reaches/deaths), then action @@ -668,18 +1019,20 @@ void puf_log(Log* log, Dict* out) { dict_set(out, "episode_length", log->episode_length); dict_set(out, "max_depth", log->max_depth); dict_set(out, "min_ac", log->min_ac); - dict_set(out, "reach_mines", log->reach_mines); - dict_set(out, "reach_minetown", log->reach_minetown); - dict_set(out, "reach_deep_mines", log->reach_deep_mines); - dict_set(out, "reach_main_d5", log->reach_main_d5); - dict_set(out, "reach_sokoban", log->reach_sokoban); + dict_set(out, "depth_5", log->depth_5); + dict_set(out, "depth_10", log->depth_10); + dict_set(out, "depth_15", log->depth_15); + dict_set(out, "mines_depth", log->mines_depth); + dict_set(out, "sokoban_depth", log->sokoban_depth); + dict_set(out, "sells", log->sells); + dict_set(out, "buys", log->buys); + dict_set(out, "discoveries", log->discoveries); dict_set(out, "death_combat", log->death_combat); + dict_set(out, "death_weak", log->death_weak); dict_set(out, "death_starved", log->death_starved); - dict_set(out, "death_smited", log->death_smited); dict_set(out, "death_other", log->death_other); dict_set(out, "death_mon_level", log->death_mon_level); - dict_set(out, "death_adj_monsters", log->death_adj_monsters); - dict_set(out, "death_maxhp", log->death_maxhp); + dict_set(out, "death_ac", log->death_ac); for (int v = 0; v < NETHACK_NUM_ACTIONS; v++) { if (NETHACK_VERB_STAT[v]) dict_set(out, NETHACK_VERB_STAT[v], log->verb_uses[v]); @@ -689,22 +1042,20 @@ void puf_log(Log* log, Dict* out) { dict_set(out, "new_tiles", log->new_tiles); dict_set(out, "enhances", log->enhances); dict_set(out, "floor_eats", log->floor_eats); - dict_set(out, "prayers_low_hp", log->prayers_low_hp); - dict_set(out, "prayers_starving", log->prayers_starving); + dict_set(out, "reads_scroll", log->reads_scroll); + dict_set(out, "reads_book", log->reads_book); dict_set(out, "burdened_frac", log->burdened_frac); - dict_set(out, "damage_taken", log->damage_taken); - dict_set(out, "ac", log->ac); - dict_set(out, "armor_swaps", log->armor_swaps); - dict_set(out, "heal_hp", log->heal_hp); - dict_set(out, "cures", log->cures); dict_set(out, "game_time", log->game_time); dict_set(out, "max_xp_level", log->max_xp_level); + dict_set(out, "floors", log->floors); + dict_set(out, "scout_held", log->scout_held); dict_set(out, "truncated", log->truncated); } // Per-(verb,head) consumption map for PPO consumed-head gating (weak symbol -// read by src/algo.cu / src/pufferl.cu). heads: [0]=verb, [1..12]=slot heads 0..11, -// [13]=direction. A head is "consumed" iff the sampled verb actually uses it. +// read by src/algo.cu). heads: [0]=verb, [1..12]=slot heads 0..11, +// [13..18]=per-verb dir heads, [19]=spell-slot head (CAST). A head is +// "consumed" iff the sampled verb actually uses it. #define PUFFER_PROVIDES_HEAD_CONSUME_MAP 1 const signed char* env_head_consume_map(int* n_verbs, int* n_atns) { static signed char map[NETHACK_NUM_ACTIONS * NUM_ATNS]; @@ -713,13 +1064,12 @@ const signed char* env_head_consume_map(int* n_verbs, int* n_atns) { memset(map, 0, sizeof(map)); for (int v = 0; v < NETHACK_NUM_ACTIONS; v++) { signed char* row = map + v * NUM_ATNS; - row[0] = 1; // verb head: always - int sh = NETHACK_VERBS[v].head; // slot head 0..11 or -1 + row[0] = 1; // verb head: always + int sh = NETHACK_VERBS[v].head; // slot head 0..11 or -1 if (sh >= 0) row[1 + sh] = 1; - if (v == NETHACK_ACT_MOVE || v == NETHACK_ACT_RUN - || v == NETHACK_ACT_KICK || v == NETHACK_ACT_THROW - || v == NETHACK_ACT_ZAP || v == NETHACK_ACT_APPLY) - row[NUM_ATNS - 1] = 1; // direction head + int dh = nethack_dir_head(v); + if (dh >= 0) row[13 + dh] = 1; + if (v == NETHACK_ACT_CAST) row[19] = 1; // spell-slot head } built = 1; } diff --git a/ocean/nethack/netlib.h b/ocean/nethack/netlib.h index c558f7374e..a4029e8a5b 100644 --- a/ocean/nethack/netlib.h +++ b/ocean/nethack/netlib.h @@ -1,8 +1,9 @@ -// Static data for the NetHack env: layout and glyph constants, the action +// Static data for the NetHack env: glyph and layout constants, the action // space and verb table, engine options, and the telemetry structs. -// Included by nethack.h after nletypes.h. #pragma once +// object tables + // object-type -> armor slot (ARM_SUIT=0..ARM_SHIRT=6, -1 = not armor), indexed // by otyp = glyph - NH_GLYPH_OBJ_OFF; generated from the engine's objects[] // (gen_obj_armcat, NetHack 3.6.6). Device copy inlined in ocean/nethack/nethack.cu. @@ -34,192 +35,263 @@ static const signed char nh_obj_armcat[NH_NUM_OBJECTS] = { -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, }; +// glyphs + +#define NETHACK_PAD_GLYPH 5976 + +// corpse glyphs: [GLYPH_BODY_OFF, +NUMMONS), display.h +#define NETHACK_GLYPH_BODY_OFF 1144 +#define NETHACK_NUMMONS 381 + +// object glyphs [GLYPH_OBJ_OFF, GLYPH_CMAP_OFF); underfoot objects win over terrain +#define NETHACK_GLYPH_OBJ_LO 1906 +#define NETHACK_GLYPH_OBJ_HI 2359 + +// cmap wall glyphs S_vwall..S_trwall; S_stone excluded (= "unexplored") +#define NETHACK_WALL_GLYPH_LO 2360 +#define NETHACK_WALL_GLYPH_HI 2370 + +// stair/ladder/altar cmap glyphs: GLYPH_CMAP_OFF(2359) + S_upstair(23).. +#define NETHACK_GLYPH_UPSTAIR 2382 +#define NETHACK_GLYPH_DNSTAIR 2383 +#define NETHACK_GLYPH_UPLADDER 2384 +#define NETHACK_GLYPH_DNLADDER 2385 +#define NETHACK_GLYPH_ALTAR 2386 + +// obs layout: glyphs | blstats | extra | inventory | item state | discovered | message + #define NH_ROWS 21 #define NH_COLS 79 #define NH_GRID (NH_ROWS * NH_COLS) // encoder views (GPU-side): egocentric crop + 5x5 patches over the grid -#define NETHACK_CROP 9 -#define NETHACK_CROP_GRID (NETHACK_CROP * NETHACK_CROP) -#define NETHACK_PAD_GLYPH 5976 +#define NETHACK_CROP 9 +#define NETHACK_CROP_GRID (NETHACK_CROP * NETHACK_CROP) -// corpse glyphs: [GLYPH_BODY_OFF, +NUMMONS), display.h -#define NETHACK_GLYPH_BODY_OFF 1144 -#define NETHACK_NUMMONS 381 - -// obs layout: glyphs | blstats | extra | inventory | item state | message -#define NETHACK_NUM_OCLASSES 18 // MAXOCLASSES; inv_oclasses pads with 18 -#define NETHACK_OFF_GLYPHS 0 +#define NETHACK_NUM_OCLASSES 18 // MAXOCLASSES; inv_oclasses pads with 18 +#define NETHACK_OFF_GLYPHS 0 #define NETHACK_OFF_BLSTATS (NH_GRID * 2) -#define NETHACK_OFF_EXTRA (NETHACK_OFF_BLSTATS + NLE_BLSTATS_SIZE * 4) -#define NETHACK_EXTRA_INTS (2 + NETHACK_NUM_OCLASSES) + +// extra ints: [0] engraving, [1] prev_action, [2..] per-class inv counts, +// in-shop bit + affordability percent, spell block (known count + 8 quads of +// id/level/fail%/retention turns; know 0 = forgotten, the re-read cue), +// 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) +#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) + // inventory: 55 slot glyphs (slot heads index these), then 8 gated int8 -// state fields per slot [buc, spe, quan, ero1, ero2, flags, typeknown, rsvd] -#define NETHACK_INV_SLOTS NLE_INVENTORY_SIZE -#define NETHACK_OFF_INV (NETHACK_OFF_EXTRA + NETHACK_EXTRA_INTS * 4) -#define NETHACK_OFF_INVST (NETHACK_OFF_INV + NETHACK_INV_SLOTS * 2) +// state fields per slot [buc, spe, quan, ero1, ero2, flags, typeknown, rsvd], +// then discovered-type glyphs (true otyp once dknown && oc_name_known, else pad) +#define NETHACK_INV_SLOTS NLE_INVENTORY_SIZE +#define NETHACK_OFF_INV (NETHACK_OFF_EXTRA + NETHACK_EXTRA_INTS * 4) +#define NETHACK_OFF_INVST (NETHACK_OFF_INV + NETHACK_INV_SLOTS * 2) +#define NETHACK_OFF_INVTRUE (NETHACK_OFF_INVST + NETHACK_INV_SLOTS * NLE_INV_STATE_FIELDS) + // raw topline chars, null-padded; must match NH_MSG_LEN in ocean/nethack/nethack.cu -#define NETHACK_OFF_MSG (NETHACK_OFF_INVST + NETHACK_INV_SLOTS * NLE_INV_STATE_FIELDS) -#define NETHACK_MSG_LEN 128 -#define NETHACK_OBS_SIZE (NETHACK_OFF_MSG + NETHACK_MSG_LEN) -#define NETHACK_INTERNAL_KILLER_MNUM 9 // killer monster index + 1 (0 = not a monster), death only +#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) + +// engine state + +#define NETHACK_PATH_MAX 128 // engine records up to 128 hero tiles/step #define NETHACK_INTERNAL_KILLER_MLEV 10 // killer monster level, death only #define NETHACK_MAX_EPISODE_STEPS 10000 -#define NETHACK_AUTODISMISS_MAX 64 // cap on prompt-dismiss keystrokes per step -#define NETHACK_MAX_DEPTH 64 // scout bitmaps tracked per episode +#define NETHACK_AUTODISMISS_MAX 64 // cap on prompt-dismiss keystrokes per step +#define NETHACK_MAX_DEPTH 64 // scout bitmaps: max distinct (dnum, dlevel) floors per episode -// Stats.areas bits; logged as reach_* proportions -#define NETHACK_AREA_MINES 1u // Gnomish Mines (dnum 2) -#define NETHACK_AREA_MINETOWN 2u // Mines level 3+ (Minetown band) -#define NETHACK_AREA_DEEP_MINES 4u // Mines level 5+ (past Minetown) -#define NETHACK_AREA_MAIN_D5 8u // Dungeons of Doom depth 5+ (Oracle route) -#define NETHACK_AREA_SOKOBAN 16u // Sokoban (dnum 4) +// hunger states (hack.h): SATIATED 0, NOT_HUNGRY 1, HUNGRY 2, WEAK 3, FAINTING 4 +#define NETHACK_HUNGER_WEAK 3 // nle_obs.misc[] prompt-state flags enum { NETHACK_MISC_YN = 0, NETHACK_MISC_GETLIN = 1, NETHACK_MISC_XWAIT = 2 }; -// action space: verb head (22) + 12 item-slot heads (55) + direction head (8) -#define NETHACK_NUM_ACTIONS 22 -#define NETHACK_NUM_DIRS 8 +// actions + +// action space: verb head (26) + 12 item-slot heads (55) + 6 per-verb +// direction heads (8 each: MOVE RUN KICK THROW ZAP APPLY) + spell-slot head (8) +#define NETHACK_NUM_ACTIONS 26 +#define NETHACK_NUM_DIRS 8 +#define NETHACK_DIR_HEADS 6 + static const int NETHACK_DIR_KEYS[NETHACK_NUM_DIRS] = - {'k','j','h','l','y','u','b','n'}; // N S W E NW NE SW SE + {'k','j','h','l','y','u','b','n'}; // N S W E NW NE SW SE static const int NETHACK_DIR_DX[NETHACK_NUM_DIRS] = { 0, 0,-1, 1,-1, 1,-1, 1}; static const int NETHACK_DIR_DY[NETHACK_NUM_DIRS] = {-1, 1, 0, 0,-1,-1, 1, 1}; -// cmap wall glyphs S_vwall..S_trwall; S_stone excluded (= "unexplored") -#define NETHACK_WALL_GLYPH_LO 2360 -#define NETHACK_WALL_GLYPH_HI 2370 -// hunger states (hack.h): SATIATED 0, NOT_HUNGRY 1, HUNGRY 2, WEAK 3, FAINTING 4 -#define NETHACK_HUNGER_WEAK 3 -// major-trouble condition bits (botl.h BL_MASK_): STONE|SLIME|STRNGL|FOODPOIS|TERMILL -#define NETHACK_COND_MAJOR 0x1Fu -#define NETHACK_COND_BAD 0x3FFu // all afflictions STONE..HALLU; excludes LEV/FLY/RIDE -// stair/ladder cmap glyphs: GLYPH_CMAP_OFF(2359) + S_upstair(23)..S_dnladder(26) -#define NETHACK_GLYPH_UPSTAIR 2382 -#define NETHACK_GLYPH_DNSTAIR 2383 -#define NETHACK_GLYPH_UPLADDER 2384 -#define NETHACK_GLYPH_DNLADDER 2385 -// object glyphs [GLYPH_OBJ_OFF, GLYPH_CMAP_OFF); underfoot objects win over terrain -#define NETHACK_GLYPH_OBJ_LO 1906 -#define NETHACK_GLYPH_OBJ_HI 2359 enum { - NETHACK_ACT_MOVE = 0, - NETHACK_ACT_RUN = 1, - NETHACK_ACT_DOWN = 2, - NETHACK_ACT_UP = 3, - NETHACK_ACT_KICK = 4, - NETHACK_ACT_SEARCH = 5, + NETHACK_ACT_MOVE = 0, + NETHACK_ACT_RUN = 1, + NETHACK_ACT_DOWN = 2, + NETHACK_ACT_UP = 3, + NETHACK_ACT_KICK = 4, + NETHACK_ACT_SEARCH = 5, NETHACK_ACT_ELBERETH = 6, - NETHACK_ACT_WEAR = 7, - NETHACK_ACT_EAT = 8, - NETHACK_ACT_QUAFF = 9, - NETHACK_ACT_PRAY = 10, - NETHACK_ACT_THROW = 11, - NETHACK_ACT_ZAP = 12, - NETHACK_ACT_SEARCH20 = 13, // count-prefixed search: ~20 turns of rest - NETHACK_ACT_PICKUP = 14, // no slot: grab the pile underfoot (beyond narrow autopickup) - NETHACK_ACT_TAKEOFF = 15, - NETHACK_ACT_PUTON = 16, - NETHACK_ACT_REMOVE = 17, - NETHACK_ACT_WIELD = 18, - NETHACK_ACT_APPLY = 19, - NETHACK_ACT_READ = 20, - NETHACK_ACT_DROP = 21, + NETHACK_ACT_WEAR = 7, + NETHACK_ACT_EAT = 8, + NETHACK_ACT_QUAFF = 9, + NETHACK_ACT_PRAY = 10, + NETHACK_ACT_THROW = 11, + NETHACK_ACT_ZAP = 12, + NETHACK_ACT_SEARCH20 = 13, // count-prefixed search: ~20 turns of rest + NETHACK_ACT_PICKUP = 14, // no slot: grab the pile underfoot (beyond autopickup) + NETHACK_ACT_TAKEOFF = 15, + NETHACK_ACT_PUTON = 16, + NETHACK_ACT_REMOVE = 17, + NETHACK_ACT_WIELD = 18, + NETHACK_ACT_APPLY = 19, + NETHACK_ACT_READ = 20, + NETHACK_ACT_DROP = 21, + NETHACK_ACT_ALTAR_ID = 22, // bulk BUC-identify on an altar + NETHACK_ACT_TIP = 23, // empty a floor container (chest/box) underfoot + NETHACK_ACT_ENGRAVE_ID = 24, // engrave-test the first unidentified wand + NETHACK_ACT_CAST = 25, // cast a known spell (fail% is in the obs) }; +// dir-head index (0..NETHACK_DIR_HEADS-1) for verbs that take a direction +static inline int nethack_dir_head(int verb) { + if (verb == NETHACK_ACT_CAST) return 4; // ZAP's dir head + its hostile-ray mask + switch (verb) { + case NETHACK_ACT_MOVE: return 0; + case NETHACK_ACT_RUN: return 1; + case NETHACK_ACT_KICK: return 2; + case NETHACK_ACT_THROW: return 3; + case NETHACK_ACT_ZAP: return 4; + case NETHACK_ACT_APPLY: return 5; + } + return -1; +} + +// inventory letter -> bit index (a-z 0-25, A-Z 26-51), -1 for non-letters +static inline int nethack_letter_bit(int c) { + if (c >= 'a' && c <= 'z') return c - 'a'; + if (c >= 'A' && c <= 'Z') return 26 + c - 'A'; + return -1; +} + +// engine options + // !status_updates skips the status renderer + recalc_mapseen (~25% of engine) #define NETHACK_DEFAULT_OPTIONS \ - "name:Agent-val-dwa-law-fem," \ + "name:Agent-mon-hum-neu-mal," \ "autopickup,color,disclose:+i +a +v +g +c +o," \ "mention_walls,nobones,nocmdassist,nolegacy,nosparkle," \ - "pickup_burden:unencumbered,pickup_types:$[%!)/," \ + "pickup_burden:unencumbered," \ "runmode:teleport,showexp,showscore,time," \ "!status_updates" +// verb table + // slot-head legality per verb (masking + decode; execution is nethack_execute) enum { WORN_ANY = 0, WORN_ONLY = 1, UNWORN_ONLY = 2 }; typedef struct { - signed char head; // -1 direct verb, 0..11 = item slot head + signed char head; // -1 direct verb, 0..11 = item slot head unsigned int item_classes; unsigned char wornreq; } Verb; +// DROP excludes COIN: drop-gold was a parser-vetoed no-op exploit static const Verb NETHACK_VERBS[NETHACK_NUM_ACTIONS] = { - {-1}, - {-1}, - {-1}, - {-1}, - {-1}, - {-1}, - {-1}, - {0, 1u<<3, UNWORN_ONLY}, - {1, 1u<<7, WORN_ANY}, - {2, 1u<<8, WORN_ANY}, - {-1}, - {3, 1u<<2, WORN_ANY}, - {4, 1u<<11, WORN_ANY}, - {-1}, - {-1}, - {5, 1u<<3, WORN_ONLY}, - {6, (1u<<4)|(1u<<5), UNWORN_ONLY}, - {7, (1u<<4)|(1u<<5), WORN_ONLY}, - {8, 1u<<2, WORN_ANY}, - {9, 1u<<6, WORN_ANY}, - {10, (1u<<9)|(1u<<10), WORN_ANY}, - {11, 0x3FFFFu, UNWORN_ONLY}, + {-1} /* MOVE */, + {-1} /* RUN */, + {-1} /* DOWN */, + {-1} /* UP */, + {-1} /* KICK */, + {-1} /* SEARCH */, + {-1} /* ELBERETH */, + {0, 1u<<3, UNWORN_ONLY} /* WEAR */, + {1, 1u<<7, WORN_ANY} /* EAT */, + {2, 1u<<8, WORN_ANY} /* QUAFF */, + {-1} /* PRAY */, + {3, 1u<<2, WORN_ANY} /* THROW */, + {4, 1u<<11, WORN_ANY} /* ZAP */, + {-1} /* SEARCH20 */, + {-1} /* PICKUP */, + {5, 1u<<3, WORN_ONLY} /* TAKEOFF */, + {6, (1u<<4)|(1u<<5), UNWORN_ONLY} /* PUTON */, + {7, (1u<<4)|(1u<<5), WORN_ONLY} /* REMOVE */, + {8, 1u<<2, WORN_ANY} /* WIELD */, + {9, 1u<<6, WORN_ANY} /* APPLY */, + {10, (1u<<9)|(1u<<10), WORN_ANY} /* READ */, + {11, 0x3EFFFu, UNWORN_ONLY} /* DROP */, + {-1} /* ALTAR_ID */, + {-1} /* TIP */, + {-1} /* ENGRAVE_ID */, + {-1} /* CAST */ }; // wandb key per verb success counter (NULL = not logged) static const char* NETHACK_VERB_STAT[NETHACK_NUM_ACTIONS] = { - NULL, NULL, NULL, NULL, NULL, - "searches", "engraves", "wears", "eats", "quaffs", - "prayers", "throws", "zaps", "search20", "pickups", - "takeoffs", "putons", "removes", "wields", "applies", - "reads", "drops", + NULL /* MOVE */, + "runs" /* RUN */, + NULL /* DOWN */, + NULL /* UP */, + NULL /* KICK */, + "searches" /* SEARCH */, + "engraves" /* ELBERETH */, + "wears" /* WEAR */, + "eats" /* EAT */, + "quaffs" /* QUAFF */, + "prayers" /* PRAY */, + "throws" /* THROW */, + "zaps" /* ZAP */, + "search20" /* SEARCH20 */, + "pickups" /* PICKUP */, + "takeoffs" /* TAKEOFF */, + "putons" /* PUTON */, + "removes" /* REMOVE */, + "wields" /* WIELD */, + "applies" /* APPLY */, + NULL /* READ: exported as reads_scroll + reads_book */, + "drops" /* DROP */, + "altar_ids" /* ALTAR_ID */, + "tips" /* TIP */, + "engrave_ids" /* ENGRAVE_ID */, + "casts" /* CAST */ }; +// telemetry + typedef struct Log { float perf; - float verb_uses[NETHACK_NUM_ACTIONS]; // success counters, keys = NETHACK_VERB_STAT + float verb_uses[NETHACK_NUM_ACTIONS]; // success counters, keys = NETHACK_VERB_STAT float score; float episode_return; float episode_length; - float valid_moves; // steps that advanced NetHack's turn counter - float illegal_actions; // steps that hit a sub-prompt we ESC'd + float valid_moves; // steps that advanced NetHack's turn counter + float illegal_actions; // steps that hit a sub-prompt we ESC'd float new_tiles; - float max_depth; // deepest level reached (depth under-reports at death) - float enhances; // #enhance presses (skill advancement claims) - float floor_eats; // eats that accepted a floor "eat it?" offer - float prayers_low_hp; // prayers at <=25% max HP (looser than real trouble) - float prayers_starving; // prayers at hunger >= Weak: TROUBLE_STARVING, prayer feeds you - // rest/retreat/burden diagnostics (log-only) - float burdened_frac; // steps with encumbrance > Unencumbered - float damage_taken; - float ac; // mean armor class over the episode (lower = better) - float min_ac; // best (lowest) AC reached this episode - float armor_swaps; // atomic WEAR swaps (auto-takeoff + wear in one step) - float heal_hp; // HP restored by heal actions (quaff/pray) this episode - float cures; // bad conditions cleared this episode - float game_time; // NetHack turns survived + float max_depth; // deepest level reached (depth under-reports at death) + float floors; // unique (dnum, dlevel) floors visited + 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; + float reads_book; + float discoveries; // object types discovered this episode (oc_name_known delta) + float sells; // shop sale offers accepted (deliberate drop in a shop) + float buys; // shop pickups paid for + float burdened_frac; // steps with encumbrance > Unencumbered + float min_ac; // best (lowest) AC reached this episode + float game_time; // NetHack turns survived float max_xp_level; - // episode end reason, one-hot (game_end_types in hack.h) + // episode end reason (game_end_types in hack.h); other = traps/wrath/poison/... float death_combat; float death_starved; - float death_smited; // god's wrath (NLE_HOW_WRATH), not a monster kill float death_other; - // combat-death anatomy (0 for non-combat episodes; ~95% are combat) - float death_mon_level; // killer's monster level (vs max_xp_level = the mismatch) - float death_adj_monsters; // hostile monsters adjacent on the last obs before death - float death_maxhp; // max HP at death (progression measure) - float truncated; // hit NETHACK_MAX_EPISODE_STEPS - // 0/1 per episode; the logged mean is the proportion - float reach_mines; - float reach_minetown; - float reach_deep_mines; - float reach_main_d5; - float reach_sokoban; + float death_weak; // any death at Weak+ hunger (fainting included) + float death_mon_level; // killer's monster level (vs max_xp_level = the mismatch) + float death_ac; // AC on the last obs before death + float truncated; // hit NETHACK_MAX_EPISODE_STEPS float n; } Log; @@ -229,24 +301,27 @@ typedef struct Stats { long valid_moves; long illegal_actions; long new_tiles; + long scout_held; long enhances; - long armor_swaps; // atomic WEAR that auto-took-off an occupant long burdened_steps; - long heal_hp; // HP restored by heal actions (quaff/pray) - long cures; // bad conditions cleared - int min_ac; // best (lowest) AC reached this episode - long last_maxhp; - int last_adj; // hostile monsters adjacent, last obs - long prayers_low_hp; - long prayers_starving; + int min_ac; + int last_ac; // AC on the last living obs (death-step blstats are torn down) long floor_eats; - long damage; - long ac_sum; // sum of AC over living steps; mean = ac_sum/length + long reads_scroll; + long reads_book; + long sells; + long buys; + int last_hunger; int max_depth; + int floors; + unsigned long long floors_bits[16]; // dnum 0..15, bit dlevel-1 int max_xp; - unsigned areas; // NETHACK_AREA_* bits float ret; int length; - // per-level first-visit bitmaps; branch levels sharing a depth share one - unsigned char visited[NETHACK_MAX_DEPTH][(NH_GRID + 7) / 8]; + // per-tile scout claim ledger: the xp level at which the tile was last + // claimed (0 = unclaimed), so tiles re-claim as the hero levels. + // Branch levels sharing a depth share one slot. + unsigned char visited[NETHACK_MAX_DEPTH][NH_GRID]; + unsigned short visited_key[NETHACK_MAX_DEPTH]; // dnum << 8 | dlevel per slot + int n_visited_floors; } Stats; diff --git a/resources/nethack/nethack_weights.bin b/resources/nethack/nethack_weights.bin index bc8174259c..c163242e39 100644 Binary files a/resources/nethack/nethack_weights.bin and b/resources/nethack/nethack_weights.bin differ diff --git a/src/algo.cu b/src/algo.cu index 1934931291..4deaffca70 100644 --- a/src/algo.cu +++ b/src/algo.cu @@ -1382,6 +1382,77 @@ static const signed char* get_head_consume_dev(int* stride) { return g_hc_dev; } +// Verb-eps exploration floor: head 0 is sampled and trained as the mixture +// (1-eps)*softmax + eps*uniform(legal), annealed per the train.verb_eps_* +// keys. Must init before CUDA graph capture (cudaMalloc is illegal +// mid-capture; the pointer bakes into captured kernels); the annealed value +// is memcpy'd each rollout. +static float* g_veps_dev = NULL; +void init_verb_eps(float base) { + if (base > 0.0f && g_veps_dev == NULL) { + cudaMalloc(&g_veps_dev, sizeof(float)); + cudaMemcpy(g_veps_dev, &base, sizeof(float), cudaMemcpyHostToDevice); + } +} + +// constant to verb_eps_anneal_start, then linear to 0 at verb_eps_anneal_end +void verb_eps_update(float base, float start, float end, long step, long total) { + if (base <= 0.0f) return; + float a = start; + if (a > 0.99f) a = 0.99f; + if (a < 0.0f) a = 0.0f; + float ae = end; + if (ae > 1.0f || ae <= 0.0f) ae = 1.0f; + if (ae < a + 0.01f) ae = a + 0.01f; + double frac = total > 0 ? (double)step / (double)total : 0.0; + float eps = frac < a ? base + : frac >= ae ? 0.0f + : base * (float)((ae - frac) / (ae - a)); + if (g_veps_dev != NULL) + cudaMemcpy(g_veps_dev, &eps, sizeof(float), cudaMemcpyHostToDevice); +} + +// sampling-side mixture (see sample_logits); the training side is below +__device__ inline float verb_eps_load(const float* verb_eps, + const precision_t* mask_row, int A, float* inv_K) { + *inv_K = 0.0f; + if (verb_eps == NULL) return 0.0f; + float eps = *verb_eps; + if (eps <= 0.0f) return 0.0f; + int K = 0; + for (int a = 0; a < A; a++) + if (to_float(mask_row[a]) != 0.0f) K++; + if (K == 0) K = A; + *inv_K = 1.0f / (float)K; + return eps; +} + +__device__ inline float verb_eps_mix(float prob, precision_t mask, float eps, float inv_K) { + float legal = to_float(mask) != 0.0f ? 1.0f : 0.0f; + return (1.0f - eps) * prob + eps * legal * inv_K; +} + +__device__ inline float verb_eps_logp(float log_prob, float eps, float inv_K) { + return logf((1.0f - eps) * expf(log_prob) + eps * inv_K); +} + +// training-side mixture (see ppo_loss_compute): converts the chosen verb's +// softmax logp to the mixture logp; *scale gets the head's gradient factor +__device__ inline float verb_eps_train_logp(const float* verb_eps, + const precision_t* mask_row, int A, float lp, float* scale) { + if (verb_eps == NULL) return lp; + float eps = *verb_eps; + if (eps <= 0.0f) return lp; + int K = 0; + for (int j = 0; j < A; ++j) + if (to_float(mask_row[j]) != 0.0f) K++; + if (K == 0) K = A; + float p_act = __expf(lp); + float p_mix = (1.0f - eps) * p_act + eps / (float)K; + *scale = (1.0f - eps) * p_act / p_mix; + return __logf(p_mix); +} + constexpr int PPO_THREADS = 256; // Per-env from ENV_HEADER (ocean//.h). @@ -1431,6 +1502,7 @@ struct PPOKernelArgs { const precision_t* action_mask; // (N, T, A_total); always present const signed char* head_consume; // (nverbs, num_atns) or NULL int hc_stride; + const float* verb_eps; // device scalar, NULL = floor off int num_atns; float clip_coef, vf_clip_coef, vf_coef; const float* ent_coef; // device ptr — host by-value bakes into CUDA graphs @@ -1559,6 +1631,7 @@ __global__ void ppo_loss_compute( float total_log_prob = 0.0f; float total_entropy = 0.0f; + float verb_mix_scale = 1.0f; // Stash across policy fwd → d_new_logp → bwd (need total logp before head grads). float head_logsumexp[NUM_ATNS]; float head_entropy[NUM_ATNS]; @@ -1599,7 +1672,13 @@ __global__ void ppo_loss_compute( head_logsumexp[h] = lse; head_entropy[h] = ent; if (head_used[h]) { - total_log_prob += cache[act] - lse; + // verb-eps floor; entropy stays on the bare softmax + if (h == 0 && a.verb_eps != NULL) + total_log_prob += verb_eps_train_logp(a.verb_eps, + a.action_mask + at_base + logits_offset, A, + cache[act] - lse, &verb_mix_scale); + else + total_log_prob += cache[act] - lse; total_entropy += ent; } logits_offset += A; @@ -1646,6 +1725,9 @@ __global__ void ppo_loss_compute( float ent = head_entropy[h]; int act = head_act[h]; float* cache = logit_cache[h]; + // verb-eps floor: head 0's gradient scales by (1-eps)p/p_mix + float d_logp_save = d_new_logp; + if (h == 0) d_new_logp *= verb_mix_scale; for (int j = 0; j < A; ++j) { float logp = cache[j] - lse; float p = __expf(logp); @@ -1653,6 +1735,7 @@ __global__ void ppo_loss_compute( ((j == act ? 1.0f : 0.0f) - p) * d_new_logp + d_entropy_term * p * (-ent - logp); } + d_new_logp = d_logp_save; logits_offset += A; } } @@ -1799,6 +1882,7 @@ void ppo_loss_fwd_bwd( .action_mask = graph.mb_action_mask.data, .head_consume = hc_dev_l, .hc_stride = hc_stride_l, + .verb_eps = g_veps_dev, .num_atns = NUM_ATNS, .clip_coef = clip_coef, .vf_clip_coef = vf_clip_coef, .vf_coef = vf_coef, .ent_coef = ent_coef, diff --git a/src/puffercpu.h b/src/puffercpu.h index c0f87eb4c7..e9e65ec819 100644 --- a/src/puffercpu.h +++ b/src/puffercpu.h @@ -228,14 +228,20 @@ void _multidiscrete(float* input, float* output, int batch_size, int logit_sizes } } } else { + float max_logit = input[in_adr]; + for (int i = 1; i < n; i++) { + if (input[in_adr + i] > max_logit) { + max_logit = input[in_adr + i]; + } + } float logit_exp_sum = 0.0f; for (int i = 0; i < n; i++) { - logit_exp_sum += expf(input[in_adr + i]); + logit_exp_sum += expf(input[in_adr + i] - max_logit); } float prob = rand() / (float)RAND_MAX; float logit_prob = 0.0f; for (int i = 0; i < n; i++) { - logit_prob += expf(input[in_adr + i]) / logit_exp_sum; + logit_prob += expf(input[in_adr + i] - max_logit) / logit_exp_sum; if (prob < logit_prob) { output[out_adr] = (float)i; break; diff --git a/src/pufferl.cu b/src/pufferl.cu index 51628446cc..5403319fb6 100644 --- a/src/pufferl.cu +++ b/src/pufferl.cu @@ -302,6 +302,9 @@ typedef struct { int num_layers; float lr; float min_lr_ratio; + float verb_eps; + float verb_eps_anneal_start; + float verb_eps_anneal_end; bool anneal_lr; float momentum; int minibatch_size; @@ -609,7 +612,8 @@ __global__ void sample_logits( precision_t* action_mask, // (B, A_total); always allocated int mask_stride, // 0 when unused const signed char* head_consume, // (nverbs, num_atns) or NULL - int hc_stride) { + int hc_stride, + const float* verb_eps) { // floor on head 0, NULL = off int B = dec_out.shape[0]; int fused_cols = dec_out.shape[1]; int num_atns = NUM_ATNS; @@ -652,16 +656,35 @@ __global__ void sample_logits( float logsumexp = ppo_discrete_logsumexp( logits, logits_base, logits_offset, A, action_mask, mask_base, cache); + float inv_K = 0.0f; + float eps = h == 0 ? verb_eps_load(verb_eps, + action_mask + mask_base + logits_offset, A, &inv_K) : 0.0f; + float rand_val = curand_uniform(&state); float cumsum = 0.0f; int sampled = A - 1; for (int a = 0; a < A; a++) { - cumsum += expf(cache[a] - logsumexp); + if (eps > 0.0f) + cumsum += verb_eps_mix(expf(cache[a] - logsumexp), + action_mask[mask_base + logits_offset + a], eps, inv_K); + else + cumsum += expf(cache[a] - logsumexp); if (rand_val < cumsum) { sampled = a; break; } } + // CDF fall-through (float rounding) lands on A - 1, which may be + // masked; snap to the last legal action. A legitimate A - 1 pick + // is always legal, so the snap is an exact no-op for it. + if (sampled == A - 1) { + for (int a = A - 1; a >= 0; a--) { + if (to_float(action_mask[mask_base + logits_offset + a]) != 0.0f) { + sampled = a; + break; + } + } + } // Float32 preserves large categorical IDs that BF16 cannot represent. int aidx = idx * num_atns + h; float action = (float)sampled; @@ -672,7 +695,10 @@ __global__ void sample_logits( int used = (head_consume == NULL || h == 0) ? 1 : (int)head_consume[verb * hc_stride + h]; if (used) { - total_log_prob += cache[sampled] - logsumexp; + if (eps > 0.0f) + total_log_prob += verb_eps_logp(cache[sampled] - logsumexp, eps, inv_K); + else + total_log_prob += cache[sampled] - logsumexp; } logits_offset += A; } @@ -871,7 +897,8 @@ void pufferl_forward(PuffeRL* pufferl, int buf, int t, cudaStream_t stream) { act_b.data, env->actions.data + (long)sub * act_cols, lp_b.data, val_b.data, pufferl->rng_states[buf] + off, - mask_b.data, mask_stride, hc_dev_s, hc_stride_s); + mask_b.data, mask_stride, hc_dev_s, hc_stride_s, + g_veps_dev); } if (hypers->cudagraphs) { @@ -1837,6 +1864,9 @@ PuffeRL* create_pufferl(Ini* ini, TrainContext* ctx) { .num_layers = puf_ini_get(ini, "policy", "num_layers"), .lr = puf_ini_get(ini, "train", "learning_rate"), .min_lr_ratio = puf_ini_get(ini, "train", "min_lr_ratio"), + .verb_eps = puf_ini_get(ini, "train", "verb_eps"), + .verb_eps_anneal_start = puf_ini_get(ini, "train", "verb_eps_anneal_start"), + .verb_eps_anneal_end = puf_ini_get(ini, "train", "verb_eps_anneal_end"), .anneal_lr = puf_ini_get(ini, "train", "anneal_lr") != 0, .momentum = puf_ini_get(ini, "train", "momentum"), .minibatch_size = puf_ini_get(ini, "train", "minibatch_size"), @@ -2113,6 +2143,7 @@ PuffeRL* create_pufferl(Ini* ini, TrainContext* ctx) { // Upload head-consume map before any stream capture (see init_head_consume_map). init_head_consume_map(); + init_verb_eps(hypers.verb_eps); // CUDA graphs: allocate graph array only; capture on first real use. if (hypers.cudagraphs) { @@ -2515,6 +2546,8 @@ static void log_history_bin_mean(PufLogHistory* h, const char* key, } double rollout_start(PuffeRL* p, int slot) { + verb_eps_update(p->hypers.verb_eps, p->hypers.verb_eps_anneal_start, + p->hypers.verb_eps_anneal_end, p->global_step, p->hypers.total_timesteps); p->write_slot = slot; if (p->hypers.async) { Prec* param = &p->policies[0].param; @@ -2955,6 +2988,7 @@ static PuffeRL* eval_make(Ini* ini, TrainContext* ctx, int mode) { char buf[64]; snprintf(buf, sizeof(buf), "%ld", eval_agents); puf_ini_put(ini, "vec.total_agents", buf); + puf_ini_put(ini, "train.verb_eps", "0"); } if (match) { int h = puf_ini_get(ini, "policy", "hidden_size"); diff --git a/tests/test_nethack_cuda.cu b/tests/test_nethack_cuda.cu new file mode 100644 index 0000000000..6a444e060c --- /dev/null +++ b/tests/test_nethack_cuda.cu @@ -0,0 +1,199 @@ +// Test harness for the Nethack encoder — thin wrapper around nethack.cu's real +// implementation. Built as a float (PRECISION_FLOAT) shared lib so finite-diff +// gradient checking is numerically meaningful. +// +// Build (from the 5c repo root; see tests/build_test.sh): +// nvcc -shared -o tests/nethack_test.so tests/test_nethack_cuda.cu ... +#define PRECISION_FLOAT +#include +#include +#include "../src/pufferl.cu" + +extern "C" { + +static Encoder g_enc; +static NethackEncoderWeights* g_w = nullptr; +static NethackEncoderActivations* g_a = nullptr; +static Allocator g_pa = {}, g_aa = {}, g_ga = {}; +static Decoder g_dec; +static NethackDecoderWeights* g_dw = nullptr; +static NethackDecoderActivations* g_da = nullptr; +static Allocator g_dpa = {}, g_daa = {}, g_dga = {}; +static int g_hidden = 32; + +void nh_init(int B, int hidden) { + cublas_init_handle(); + g_hidden = hidden; + g_enc = {}; + g_enc.in_dim = NH_OBS_SIZE; + g_enc.out_dim = hidden; + create_nethack_encoder(&g_enc); + g_w = (NethackEncoderWeights*)g_enc.create_weights(&g_enc); + g_pa = {}; + g_enc.reg_params(g_w, &g_pa); + alloc_create(&g_pa); + g_a = (NethackEncoderActivations*)calloc(1, sizeof(NethackEncoderActivations)); + g_aa = {}; g_ga = {}; + g_enc.reg_train(g_w, g_a, &g_aa, &g_ga, B); + alloc_create(&g_aa); + alloc_create(&g_ga); + uint64_t seed = 1234; + g_enc.init_weights(g_w, &seed, 0); + // pointer decoder, fed by g_a's inv_out (nh_enc_last set by reg_train + // above). Its keygrad buffer feeds encoder backward — zero it so the + // encoder-only checks stay exact until nh_dec_backward runs. + g_dec = {}; + g_dec.hidden_dim = hidden; + g_dec.output_dim = NH_DEC_OD; + create_nethack_decoder(&g_dec); + g_dw = (NethackDecoderWeights*)g_dec.create_weights(&g_dec); + g_dpa = {}; + g_dec.reg_params(g_dw, &g_dpa); + alloc_create(&g_dpa); + g_da = (NethackDecoderActivations*)calloc(1, g_dec.activation_size); + g_daa = {}; g_dga = {}; + g_dec.reg_train(g_dw, g_da, &g_daa, &g_dga, B); + alloc_create(&g_daa); + alloc_create(&g_dga); + g_dec.init_weights(g_dw, &seed, 0); + cudaMemset(g_da->keygrad.data, 0, (size_t)B * NH_INV_FLAT * sizeof(float)); + cudaDeviceSynchronize(); +} + +int nh_obs_size() { return NH_OBS_SIZE; } +int nh_bl_feat() { return NH_BL_FEAT; } +int nh_glyph_vocab() { return NH_GLYPH_VOCAB; } +int nh_embed_dim() { return NH_EMBED_DIM; } +int nh_concat() { return NH_CONCAT; } +int nh_grid() { return NH_MGRID; } +int nh_dec_od() { return NH_DEC_OD; } +int nh_dec_pad() { return NH_DEC_PAD; } +int nh_num_actions() { return NH_ACTIONS; } +int nh_heads() { return NH_HEADS; } + +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); + cudaMemcpy(out, r.data, (size_t)B * g_hidden * sizeof(float), cudaMemcpyDeviceToDevice); + cudaDeviceSynchronize(); +} + +void nh_backward(void* grad, int B) { + Prec g = {.data = (precision_t*)grad, .shape = {B, g_hidden}}; + g_enc.backward(g_w, g_a, g, 0); + cudaDeviceSynchronize(); +} + +// value get/set + grad get for each learnable tensor (all device float ptrs) +#define TENSOR_ACC(name, field) \ + void nh_get_##name(void* dst) { cudaMemcpy(dst, g_w->field.data, numel(g_w->field.shape) * sizeof(float), cudaMemcpyDeviceToDevice); } \ + void nh_set_##name(void* src) { cudaMemcpy(g_w->field.data, src, numel(g_w->field.shape) * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } \ + int nh_numel_##name() { return (int)numel(g_w->field.shape); } +TENSOR_ACC(embed_w, embed_w) +TENSOR_ACC(ekind_w, ekind_w) +TENSOR_ACC(esub_w, esub_w) +TENSOR_ACC(bl_w, bl_w) +TENSOR_ACC(bl_b, bl_b) +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(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) + +#define GRAD_ACC(name, field) \ + void nh_grad_##name(void* dst) { cudaMemcpy(dst, g_a->field.data, numel(g_a->field.shape) * sizeof(float), cudaMemcpyDeviceToDevice); } +GRAD_ACC(embed_w, embed_wgrad) +GRAD_ACC(ekind_w, ekind_wgrad) +GRAD_ACC(esub_w, esub_wgrad) +GRAD_ACC(bl_w, bl_wgrad) +GRAD_ACC(bl_b, bl_bgrad) +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(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) +GRAD_ACC(msg_w, msg_wgrad) +GRAD_ACC(spk_w, spk_wgrad) +GRAD_ACC(spk2_w, spk2_wgrad) +GRAD_ACC(spk2_b, spk2_bgrad) + +// ---- pointer decoder (fed by the encoder's inv_out keys) ---- +// forward: encoder -> decoder directly (no mingru in the harness); the +// decoder kernels see the same activations either way. + +void nh_dec_forward(void* out, void* obs, int B) { + Prec in = {.data = (precision_t*)obs, .shape = {B, NH_OBS_SIZE}}; + Prec h = g_enc.forward(g_w, g_a, in, 0); + Prec r = g_dec.forward(g_dw, g_da, h, 0); + cudaMemcpy(out, r.data, (size_t)B * (NH_DEC_OD + 1) * sizeof(float), cudaMemcpyDeviceToDevice); + cudaDeviceSynchronize(); +} + +// glogits (B, NH_DEC_OD) + gvalue (B,) device floats; grad wrt the decoder's +// hidden-state input lands in dinput (B, hidden) +void nh_dec_backward(void* glogits, void* gvalue, void* dinput, int B) { + Float gl = {.data = (float*)glogits, .shape = {B, NH_DEC_OD}}; + Float gs = {}; + Float gv = {.data = (float*)gvalue, .shape = {B, 1}}; + Prec gi = g_dec.backward(g_dw, g_da, gl, gs, gv, 0); + cudaMemcpy(dinput, gi.data, (size_t)B * g_hidden * sizeof(float), cudaMemcpyDeviceToDevice); + cudaDeviceSynchronize(); +} + +// decoder fed with an EXTERNAL hidden state (e.g. the demo's mingru output); +// runs the encoder first so the decoder's inv keys (ea->inv_out) are fresh +// for the same obs. +void nh_dec_forward_hidden(void* out, void* obs, void* hidden_in, int B) { + Prec in = {.data = (precision_t*)obs, .shape = {B, NH_OBS_SIZE}}; + g_enc.forward(g_w, g_a, in, 0); + Prec h = {.data = (precision_t*)hidden_in, .shape = {B, g_hidden}}; + Prec r = g_dec.forward(g_dw, g_da, h, 0); + cudaMemcpy(out, r.data, (size_t)B * (NH_DEC_OD + 1) * sizeof(float), cudaMemcpyDeviceToDevice); + cudaDeviceSynchronize(); +} + +void nh_dec_keygrad(void* dst, int B) { + cudaMemcpy(dst, g_da->keygrad.data, (size_t)B * NH_INV_FLAT * sizeof(float), cudaMemcpyDeviceToDevice); +} + +#define DEC_ACC(name, field) \ + void nh_get_##name(void* dst) { cudaMemcpy(dst, g_dw->field.data, numel(g_dw->field.shape) * sizeof(float), cudaMemcpyDeviceToDevice); } \ + void nh_set_##name(void* src) { cudaMemcpy(g_dw->field.data, src, numel(g_dw->field.shape) * sizeof(float), cudaMemcpyDeviceToDevice); cudaDeviceSynchronize(); } \ + int nh_numel_##name() { return (int)numel(g_dw->field.shape); } +DEC_ACC(dec_lin_w, lin_w) +DEC_ACC(dec_q_w, q_w) +DEC_ACC(dec_k_w, k_w) +DEC_ACC(dec_tau, tau) + +#define DEC_GRAD(name, field) \ + void nh_grad_##name(void* dst) { cudaMemcpy(dst, g_da->field.data, numel(g_da->field.shape) * sizeof(float), cudaMemcpyDeviceToDevice); } +DEC_GRAD(dec_lin_w, lin_wgrad) +DEC_GRAD(dec_q_w, q_wgrad) +DEC_GRAD(dec_k_w, k_wgrad) +DEC_GRAD(dec_tau, tau_grad) + +} // extern "C" diff --git a/tests/test_nethack_encoder.py b/tests/test_nethack_encoder.py new file mode 100644 index 0000000000..db5920c7aa --- /dev/null +++ b/tests/test_nethack_encoder.py @@ -0,0 +1,621 @@ +"""Numeric gradient check for the Nethack CUDA encoder + pointer decoder +(src/nethack.cu). + +Builds tests/test_nethack_cuda.cu as a float shared lib and verifies the +analytic gradients from encoder backward() against central finite differences +of a scalar loss L = sum(out * g_out), plus exact float64 torch references for +the whole encoder and the 5-head pointer decoder (values, weight grads, +keygrad, grad_input). + +Run: python tests/test_nethack_encoder.py +""" +import ctypes +import os +import glob +import shutil +import subprocess +import sys + +import numpy as np + +HERE = os.path.dirname(os.path.abspath(__file__)) +SRC = os.path.join(os.path.dirname(HERE), "src") +LIB = os.path.join(HERE, "nethack_test.so") + +VP = ctypes.c_void_p + +WEIGHT_NAMES = [ + "embed_w", "ekind_w", "esub_w", "bl_w", "bl_b", "proj_w", "proj_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", "msg_w", "spk_w", "spk2_w", "spk2_b", + "dec_lin_w", "dec_q_w", "dec_k_w", "dec_tau", +] + + +def build(): + root = os.path.dirname(HERE) + cuda = os.path.dirname(os.path.dirname(shutil.which("nvcc"))) + raylib = glob.glob(os.path.join(root, "raylib-*"))[0] + cmd = [ + "nvcc", "-shared", "-o", LIB, os.path.join(HERE, "test_nethack_cuda.cu"), + "-std=c++17", "-arch=native", + "-I" + root, "-I" + os.path.join(root, "src"), + "-I" + os.path.join(root, "ocean", "nethack"), + "-I" + os.path.join(root, "vendor"), + "-I" + os.path.join(root, "vendor", "fast-nle", "include"), + "-I" + os.path.join(root, "vendor", "fast-nle", "build", + "_deps", "deboost_context-src", "include"), + "-I" + os.path.join(cuda, "include"), + "-I" + os.path.join(cuda, "include", "cccl"), + "-I" + os.path.join(raylib, "include"), + '-DENV_HEADER="ocean/nethack/nethack.h"', + "-DPUFFER_NETHACK", "-DENV_NAME=nethack", '-DPUFFER_ENV_NAME="nethack"', + "-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", + ] + try: + import nvidia.nccl + nccl = nvidia.nccl.__path__[0] + cmd += ["-I" + os.path.join(nccl, "include"), + "-L" + os.path.join(nccl, "lib"), "-lnccl"] + except ImportError: + cmd += ["-lnccl"] + print("building:", " ".join(cmd)) + subprocess.run(cmd, check=True) + + +def load(): + lib = ctypes.CDLL(LIB) + for name in WEIGHT_NAMES: + for fn in [f"nh_get_{name}", f"nh_set_{name}"]: + getattr(lib, fn).argtypes = [VP] + getattr(lib, fn).restype = None + if name.startswith("dec_"): + getattr(lib, f"nh_grad_{name}").argtypes = [VP] + getattr(lib, f"nh_numel_{name}").restype = ctypes.c_int + for name in WEIGHT_NAMES: + if not name.startswith("dec_"): + getattr(lib, f"nh_grad_{name}").argtypes = [VP] + getattr(lib, f"nh_grad_{name}").restype = None + lib.nh_init.argtypes = [ctypes.c_int, ctypes.c_int] + for fn in ["nh_obs_size", "nh_bl_feat", "nh_glyph_vocab", "nh_embed_dim", + "nh_concat", "nh_grid", "nh_dec_od", "nh_heads"]: + getattr(lib, fn).restype = ctypes.c_int + lib.nh_forward.argtypes = [VP, VP, ctypes.c_int] + lib.nh_backward.argtypes = [VP, ctypes.c_int] + 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] + return lib + + +def make_obs(B, obs_size, grid, max_glyph_used): + """Build a valid packed obs: glyphs int16 LE, then blstats int32 LE, as + byte-valued float32 (matching cast_dispatch's per-byte float cast). + Returns (obs, glyphs, bl_vals, ex_vals) so the torch reference can rebuild + the exact inputs.""" + obs = np.zeros((B, obs_size), dtype=np.float32) + rng = np.random.default_rng(0) + # glyphs @0: grid cells, 2 bytes each (restrict to a small glyph set so many + # embedding rows receive gradient and are individually checkable). + glyphs = rng.integers(0, max_glyph_used, size=(B, grid)).astype(np.int32) + lo = (glyphs & 0xFF).astype(np.float32) + hi = ((glyphs >> 8) & 0xFF).astype(np.float32) + obs[:, 0:2 * grid:2] = lo + obs[:, 1:2 * grid:2] = hi + # blstats @ 2*grid: 27 int32, mixed magnitudes incl. negatives (AC/align). + bl_off = 2 * grid + vals = rng.integers(-5, 500, size=(B, 27)).astype(np.int64) + vals[:, 0] = rng.integers(0, 79, size=B) # hero x: crop center + vals[:, 1] = rng.integers(0, 21, size=B) # hero y + vals[:, 25] = rng.integers(0, 8192, size=B) # CONDITION bitmask + u = vals.astype(np.uint32) + for k in range(4): + obs[:, bl_off + k::4][:, :27] = ((u >> (8 * k)) & 0xFF).astype(np.float32) + # extra stats @ +27*4: engraving state, prev action (-1..21 valid; sampled + # -1..13 to preserve the original FD test batch — the higher onehot columns + # are linear and covered by the analytic torch check), 18 class counts + spell_cols = [] + for _ in range(8): + spell_cols += [ + rng.integers(0, 500, size=(B, 1)), # slot id (otyp; 0 = empty) + rng.integers(0, 8, size=(B, 1)), # slot level + rng.integers(0, 101, size=(B, 1)), # slot fail% + rng.integers(0, 20001, size=(B, 1)), # slot retention turns + ] + ex = np.concatenate([ + rng.integers(0, 3, size=(B, 1)), # engraving state 0/1/2 + rng.integers(-1, 14, size=(B, 1)), + rng.integers(0, 6, size=(B, 18)), + rng.integers(0, 2, size=(B, 1)), # in-shop bit + rng.integers(0, 101, size=(B, 1)), # affordability percent + rng.integers(0, 9, size=(B, 1)), # known-spell count + ] + spell_cols + [ + rng.integers(0, 320, size=(B, 1)), # encumbrance percent (unclipped) + rng.integers(50, 1001, size=(B, 1)), # carry capacity + ], axis=1).astype(np.int64).astype(np.uint32) + for k in range(4): + obs[:, bl_off + k::4][:, 27:84] = ((ex >> (8 * k)) & 0xFF).astype(np.float32) + # inventory entities: 55 slot glyphs int16 LE, tail padded (5976) + inv_off = bl_off + 84 * 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 + 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 + st_off = inv_off + 55 * 2 + st = np.zeros((B, 55, 8), dtype=np.int64) + st[:, :, 0] = rng.integers(0, 4, size=(B, 55)) + spe_known = rng.integers(0, 2, size=(B, 55)).astype(bool) + st[:, :, 1] = np.where(spe_known, rng.integers(-3, 6, size=(B, 55)), -128) + st[:, :, 2] = rng.integers(1, 20, size=(B, 55)) + st[:, :, 3] = rng.integers(0, 4, size=(B, 55)) + st[:, :, 4] = rng.integers(0, 4, size=(B, 55)) + st[:, :, 5] = rng.integers(0, 128, size=(B, 55)) + st[:, :, 6] = rng.integers(0, 2, size=(B, 55)) + obs[:, st_off:st_off + 55 * 8] = (st.reshape(B, -1) & 0xFF).astype(np.float32) + # discovered-type glyphs @ +55*8: true otyp glyph on a random identified + # subset, pad (5976) elsewhere and past the item tail + itr_off = st_off + 55 * 8 + itr = np.full((B, 55), 5976, dtype=np.int32) + known = rng.integers(0, 2, size=(B, 55)).astype(bool) + itr[known] = rng.integers(1906, 2359, size=int(known.sum())).astype(np.int32) + for b in range(B): + itr[b, n_items[b]:] = 5976 + obs[:, itr_off + 0::2][:, :55] = (itr & 0xFF).astype(np.float32) + obs[:, itr_off + 1::2][:, :55] = ((itr >> 8) & 0xFF).astype(np.float32) + # 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 = np.zeros((B, msg_len), dtype=np.int64) + alpha = np.frombuffer(b"abcdefghijklmnopqrstuvwxyz ", dtype=np.uint8).astype(np.int64) + for b in range(B): + ln = int(rng.integers(6, min(40, msg_len))) + msg[b, :ln] = alpha[rng.integers(0, len(alpha), size=ln)] + obs[:, msg_off:msg_off + msg_len] = msg.astype(np.float32) + return obs, glyphs, vals, ex.astype(np.int64).astype(np.int32), inv, st, itr, msg + + +def dev(nbytes): + import ctypes + p = VP() + _cudart.cudaMalloc(ctypes.byref(p), ctypes.c_size_t(nbytes)) + return p + + +def h2d(arr): + arr = np.ascontiguousarray(arr, dtype=np.float32) + p = dev(arr.nbytes) + _cudart.cudaMemcpy(p, arr.ctypes.data_as(VP), ctypes.c_size_t(arr.nbytes), 1) # H2D + return p, arr.nbytes + + +def d2h(p, n): + out = np.empty(n, dtype=np.float32) + _cudart.cudaMemcpy(out.ctypes.data_as(VP), p, ctypes.c_size_t(n * 4), 2) # D2H + return out + + +_cudart = ctypes.CDLL("libcudart.so") +_cudart.cudaMalloc.argtypes = [VP, ctypes.c_size_t] +_cudart.cudaMemcpy.argtypes = [VP, VP, ctypes.c_size_t, ctypes.c_int] + + +def glyph_map(): + """Parse the generated (kind, sub) mapping straight from the header.""" + import re + txt = open(os.path.join(HERE, "..", "ocean", "nethack", "glyph_map.h")).read() + def arr(name): + m = re.search(name + r"\[\d+\] = \{([0-9,\-]+)\};", txt) + return np.array([int(x) for x in m.group(1).split(",") if x], dtype=np.int64) + return arr("nh_glyph_kind"), arr("nh_glyph_sub") + + +def getw(lib, name, shape): + import torch + fn = getattr(lib, f"nh_get_{name}") + n = getattr(lib, f"nh_numel_{name}")() + a = np.empty(n, dtype=np.float32) + fn(a.ctypes.data_as(VP)) + return torch.tensor(a.astype(np.float64).reshape(shape), requires_grad=True) + + +def torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg, H): + """float64 torch replica of the encoder forward. Returns (out, invh, w) + where invh is the (B,55,16) post-relu slot features (the decoder's keys) + and w maps weight names to the torch leaf tensors.""" + import torch + B = glyphs.shape[0] + ROWS, COLS, CROP, PAD = 21, 79, 9, 5976 + PW, PH, PX, PY = 5, 5, 16, 5 + BL_SCALE = np.array([ + 1/79, 1/21, 1/25, 1/125, 1/25, 1/25, 1/25, 1/25, 1/25, 0.1, + 1/200, 1/200, 1/50, 0.1, 1/100, 1/100, 1/10, 1/10, 1/30, + 0.1, 0.1, 0.0, 1/4, 0.0, 1/50, 0.0, 1.0], dtype=np.float64) + 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)) + kind_map, sub_map = glyph_map() + 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,)) + 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,)) + + # local: crop glyph ids with pad off-map + hx, hy = bl_vals[:, 0], bl_vals[:, 1] + crop_idx = np.full((B, CROP * CROP), PAD, 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: + crop_idx[b, p] = glyphs[b, r * COLS + c] + 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) + # blstats features + f = np.zeros((B, lib.nh_bl_feat()), dtype=np.float64) + j = 0 + for i in range(27): + if i in (21, 25): + continue + v = bl_vals[:, i].astype(np.float64) + f[:, j] = np.log1p(np.maximum(v, 0)) * BL_SCALE[i] if BL_ISLOG[i] else v * BL_SCALE[i] + j += 1 + hunger = np.clip(bl_vals[:, 21], 0, 6) + for h in range(7): + f[:, j] = (hunger == h); j += 1 + for k in range(13): + f[:, j] = (bl_vals[:, 25].astype(np.uint32) >> k) & 1; j += 1 + for h in range(lib.nh_num_actions()): + f[:, j] = (ex_vals[:, 1] == h); j += 1 + for k in range(18): + f[:, j] = ex_vals[:, 2 + k] * 0.125; j += 1 + # hp_frac (hp/hpmax), ene_frac (ene/enemax), clamped to [0,1] + hp = bl_vals[:, 10].astype(np.float64); hpmax = bl_vals[:, 11].astype(np.float64) + ene = bl_vals[:, 14].astype(np.float64); enemax = bl_vals[:, 15].astype(np.float64) + f[:, j] = np.clip(hp / np.maximum(hpmax, 1), 0, 1); j += 1 + f[:, j] = np.clip(ene / np.maximum(enemax, 1), 0, 1); j += 1 + # dnum one-hot (nominal dungeon branch; scalar scale zeroed) + dnum = np.clip(bl_vals[:, 23], 0, 7) + for h in range(8): + f[:, j] = (dnum == h); j += 1 + # underfoot engraving bits: any-engraving, active-Elbereth + f[:, j] = (ex_vals[:, 0] >= 1); j += 1 + f[:, j] = (ex_vals[:, 0] >= 2); j += 1 + # shop: standing on goods, and gold/price capped at 1 + f[:, j] = ex_vals[:, 20]; j += 1 + f[:, j] = ex_vals[:, 21] * 0.01; j += 1 + # spell scalar, mirrors NH_F_SPELL: known count/8 only (per-slot content + # rides the spell-key path) + f[:, j] = ex_vals[:, 22] * 0.125; j += 1 + # encumbrance pair, mirrors NH_F_WEIGHT: softsign(ratio-1), cap/1000 + d = ex_vals[:, 55] * 0.01 - 1.0 + f[:, j] = d / (1.0 + np.abs(d)); j += 1 + f[:, j] = ex_vals[:, 56] * 0.001; 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) + # inventory entities: per-slot embed + gated state features -> 32, relu + # (the decoder's keys), then pooled 32 -> 128 with max over slots + sf = np.zeros(st_vals.shape[:2] + (24,), dtype=np.float64) + for c in range(4): + 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 + sf[:, :, 7] = st_vals[:, :, 3] / 3.0 + sf[:, :, 8] = st_vals[:, :, 4] / 3.0 + for c in range(7): + sf[:, :, 9 + c] = (st_vals[:, :, 5] >> c) & 1 + sf[:, :, 16] = st_vals[:, :, 6] + # armor slot category one-hot from the slot glyph (baked otyp->ARM_* table) + import re as _re + _src = open("ocean/nethack/netlib.h").read() + _body = _re.search(r"nh_obj_armcat\[NH_NUM_OBJECTS\] = \{(.*?)\};", _src, _re.S).group(1) + _tbl = np.array([int(x) for x in _re.findall(r"-?\d+", _body)], dtype=np.int64) + ot = inv_vals.astype(np.int64) - 1906 + 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] + 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) + # 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 + rows = [] + for b in range(B): + row = msg[b] + ids = [] + for t in range(len(row) - 2): + c0, c1, c2 = int(row[t]), int(row[t + 1]), int(row[t + 2]) + 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)) + cnt = len(ids) + s = msg_w[torch.tensor(ids, dtype=torch.long)].sum(dim=0) if cnt else torch.zeros(32, 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 = [] + 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])) + 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) + sk = torch.stack(spkeys, dim=1) # (B,8,16) + spool = torch.relu((sk @ spk2_w.T).max(dim=1).values + spk2_b) + concat = torch.cat([loc, glb, invp, blh, fb, msg_sum, spool], dim=1) + out = torch.relu(concat @ proj_w.T + proj_b) + return out, invh, w, torch.stack(spkeys, dim=1) + + +def run(lib): + B, hidden = 4, 24 + lib.nh_init(B, hidden) + obs_size = lib.nh_obs_size() + grid = lib.nh_grid() + vocab = lib.nh_glyph_vocab() + 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)) + + 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) + obs_d, _ = h2d(obs) + out_d = dev(B * hidden * 4) + + # Fixed upstream grad g_out; loss L = sum(out * g_out). + rng = np.random.default_rng(7) + g_out = rng.standard_normal((B, hidden)).astype(np.float32) + + 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()) + + # 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"] + + # 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 + # unit's sign makes the FD non-smooth and disagree with the (correct) + # subgradient. We detect such kink crossings via the second difference + # (|L+ + L- - 2*L0| is O(eps) at a kink vs O(eps^2) on a smooth region) + # and skip those entries. eps is small so smooth curvature stays negligible. + eps = 1e-3 + kink_tol = 2e-5 # |Lp+Lm-2*L0| above this ⇒ a ReLU flipped; skip entry. + rel_tol = 1.5e-2 + rng = np.random.default_rng(123) + # glb1/inv2/inv1 are FD-unverifiable: their weights feed the max-pool, so + # 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"} + all_ok = True + for name in enc_names: + if name in fd_skip: + print(f" [----] {name:8s} FD skipped (max-pool shared weights); torch-checked below") + continue + get = getattr(lib, f"nh_get_{name}") + setw = getattr(lib, f"nh_set_{name}") + gradf = getattr(lib, f"nh_grad_{name}") + n = getattr(lib, f"nh_numel_{name}")() + w0 = np.empty(n, dtype=np.float32); get(w0.ctypes.data_as(VP)) + ga = np.empty(n, dtype=np.float32); gradf(ga.ctypes.data_as(VP)) + + if name == "embed_w": + used = np.unique(np.concatenate([glyphs.reshape(-1), inv_vals.reshape(-1)])) + D = lib.nh_embed_dim() + cand = np.array([g * D + d for g in used for d in range(D)], dtype=np.int64) + else: + cand = np.arange(n, dtype=np.int64) + cand = cand[np.abs(ga[cand]) > 1e-3] # need signal for a meaningful ratio + if len(cand) == 0: + print(f" [SKIP] {name:8s} n={n:8d} (no entry with |grad|>1e-3)") + continue + rng.shuffle(cand) + + max_rel, checked, skipped = 0.0, 0, 0 + for i in cand: + if checked >= 10: + break + i = int(i) + wp = w0.copy(); wp[i] += eps; setw(wp.ctypes.data_as(VP)); Lp = forward_loss() + wm = w0.copy(); wm[i] -= eps; setw(wm.ctypes.data_as(VP)); Lm = forward_loss() + setw(w0.ctypes.data_as(VP)) # restore + if abs(Lp + Lm - 2 * L0) > kink_tol: # ReLU kink crossing → FD invalid + skipped += 1 + continue + gnum = (Lp - Lm) / (2 * eps) + 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 + all_ok = all_ok and ok + print(f" [{'OK ' if ok else 'FAIL'}] {name:8s} n={n:8d} " + f"checked={checked} kink_skipped={skipped} " + f"max|analytic|={np.abs(ga).max():.4g} max_rel_err={max_rel:.2e}") + + # Exact reference: float64 torch autograd replica of the whole encoder. + # Finite differences can't cleanly verify glb1/inv2 — shared max-pool + # weights flip near-tied argmax winners below the kink detector's + # threshold. + all_ok = torch_check(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg, g_out, enc_names, hidden) and all_ok + all_ok = dec_check(lib, obs_d, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg, B, hidden) and all_ok + return all_ok + + +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() + + 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) + denom = max(1.0, np.abs(gt).max()) + rel = np.abs(ga - gt).max() / denom + good = rel < 1e-3 + ok = ok and good + print(f" [{'OK ' if good else 'FAIL'}] torch {name:8s} max_rel_err={rel:.2e}") + return ok + + +def dec_check(lib, obs_d, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg, B, H): + """float64 torch replica of the 5-head pointer decoder: forward values, + weight grads, keygrad (grad into the encoder's slot features) and + grad_input (grad into the decoder's hidden-state input).""" + import torch + OD = lib.nh_dec_od() + HEADS = lib.nh_heads() + N_ACT = lib.nh_num_actions() + N_DIRS = OD - N_ACT - HEADS * 55 # 8 pre-champion, 48 with per-verb dir heads + PAD = lib.nh_dec_pad() + + # CUDA forward (encoder+decoder) and backward (decoder only) + out_d = dev(B * (OD + 1) * 4) + lib.nh_dec_forward(out_d, obs_d, B) + out_cuda = d2h(out_d, B * (OD + 1)).reshape(B, OD + 1) + rng = np.random.default_rng(11) + g = rng.standard_normal((B, OD + 1)).astype(np.float32) + gl_d, _ = h2d(np.ascontiguousarray(g[:, :OD])) + gv_d, _ = h2d(np.ascontiguousarray(g[:, OD])) + 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) + lib.nh_dec_keygrad(kg_cuda.ctypes.data_as(VP), B) + + # torch replica: hidden state and keys detached so grads are decoder-local, + # matching what the CUDA decoder backward produces (the encoder chain gets + # these via grad_input / keygrad separately). + h_full, invh, _, spkeys = torch_encoder(lib, glyphs, bl_vals, ex_vals, inv_vals, st_vals, itr_vals, msg, H) + h_in = h_full.detach().clone().requires_grad_(True) + 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)) + 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) + 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, :]) + 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 + N_DIRS_LIN = 48 + out = torch.cat([tmp[:, :N_ACT], slot.reshape(B, HEADS * 55), + tmp[:, N_ACT:N_ACT+N_DIRS_LIN], spell, + tmp[:, N_ACT+N_DIRS_LIN:N_ACT+N_DIRS_LIN+1]], dim=1) + (out * torch.tensor(g.astype(np.float64))).sum().backward() + + ok = True + rel = np.abs(out_cuda - out.detach().numpy()).max() / max(1.0, np.abs(out.detach().numpy()).max()) + good = rel < 1e-4 + 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)]: + 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 = ref.grad.numpy().reshape(-1) + denom = max(1.0, np.abs(gt).max()) + rel = np.abs(ga - gt).max() / denom + good = rel < 1e-3 + 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)]: + gt = ref.grad.numpy().reshape(-1) + denom = max(1.0, np.abs(gt).max()) + rel = np.abs(cuda_g.reshape(-1) - gt).max() / denom + good = rel < 1e-3 + ok = ok and good + print(f" [{'OK ' if good else 'FAIL'}] torch dec {name:9s} max_rel_err={rel:.2e}") + return ok + + +if __name__ == "__main__": + if "--no-build" not in sys.argv: + build() + ok = run(load()) + print("\nRESULT:", "PASS" if ok else "FAIL") + sys.exit(0 if ok else 1)