Skip to content

Cleaning up craftax #670

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
daphne-cornelisse wants to merge 3 commits into from
Closed

Cleaning up craftax #670

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c785e2c
Ensure light mechanics match original exactly.
daphne-cornelisse Aug 27, 2026
44c252d
Delete remaining profiling code.
daphne-cornelisse Aug 27, 2026
f8d2324
Better config?
daphne-cornelisse Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 18 additions & 24 deletions config/craftax_clean.ini
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,52 +2,46 @@
env_name = craftax_clean

[vec]
total_agents = 2048
total_agents = 4096
num_buffers = 16
num_threads = 16

[env]
reset_pool_size = 32
action_mask = 0
action_mask = 1

[train]
total_timesteps = 1_872_097_150
learning_rate = 0.00179971103
ent_coef = 0.0000203385534
gamma = 0.997127712
gae_lambda = 0.960012317
horizon = 64
total_timesteps = 429_130_004
learning_rate = 0.00883168
ent_coef = 0.000194094
gamma = 0.986552
gae_lambda = 0.861994
horizon = 128
minibatch_size = 65536
replay_ratio = 3.26567292
clip_coef = 0.207467377
vf_coef = 2.66733265
vf_clip_coef = 0.01
max_grad_norm = 0.496576279
momentum = 0.961853921
replay_ratio = 1.39199
clip_coef = 0.172376
vf_coef = 2.46606
vf_clip_coef = 1.37661
max_grad_norm = 0.239896
momentum = 0.908417

[policy]
hidden_size = 1024
num_layers = 4
hidden_size = 512
num_layers = 2.88843

[sweep]
metric = score
metric_distribution = linear
goal = maximize
max_suggestion_cost = 1200
max_suggestion_cost = 3600
max_runs = 1200
gpus = -1
downsample = 15
prune_pareto = True
early_stop_quantile = 0.3

[sweep.env.action_mask]
distribution = int_uniform
min = 0
max = 1
scale = auto

[sweep.train.total_timesteps]
distribution = log_normal
min = 3e7
max = 1e10
max = 1e11
scale = time
112 changes: 21 additions & 91 deletions ocean/craftax_clean/craftax_clean.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,55 +14,6 @@ typedef float obs_t;
#include <stdio.h>
#include <stdlib.h>

#ifdef CRAFTAX_CLEAN_PROFILE
#define CLEAN_PROF_ZONES 9
static const char* clean_prof_names[CLEAN_PROF_ZONES] = {
"floor+craft",
"do",
"place+shoot+potion",
"book+enchant+move",
"update_mobs",
"spawn",
"plants+intrinsics",
"reward+reset",
"obs",
};
static uint64_t clean_prof_ticks[CLEAN_PROF_ZONES];
static uint64_t clean_prof_count[CLEAN_PROF_ZONES];

static inline uint64_t clean_prof_now(void) {
return __builtin_ia32_rdtsc();
}

static inline void clean_prof_add(int zone, uint64_t t0) {
clean_prof_ticks[zone] += clean_prof_now() - t0;
clean_prof_count[zone]++;
}

void clean_prof_report(void) {
uint64_t total = 0;
for (int i = 0; i < CLEAN_PROF_ZONES; i++) {
total += clean_prof_ticks[i];
}
fprintf(stderr, "\n=== puf_step zones ===\n");
for (int i = 0; i < CLEAN_PROF_ZONES; i++) {
double pct = total ? 100.0 * (double)clean_prof_ticks[i] / (double)total : 0.0;
fprintf(stderr, "%-22s %6.2f%% %lu calls ticks=%lu\n",
clean_prof_names[i], pct, (unsigned long)clean_prof_count[i],
(unsigned long)clean_prof_ticks[i]);
}
}

#define CLEAN_PROF_START() uint64_t _z0 = 0
#define CLEAN_ZONE(n) do { _z0 = clean_prof_now(); } while (0)
#define CLEAN_ZONE_END(n) clean_prof_add((n), _z0)
#else
#define CLEAN_PROF_START()
#define CLEAN_ZONE(n)
#define CLEAN_ZONE_END(n)
#define clean_prof_report()
#endif

#define ACT_SIZES {ATN_DIM}
#define NUM_ATNS 1
#ifdef PUFFERCPU_EVAL_MAIN
Expand Down Expand Up @@ -1837,7 +1788,6 @@ void puf_step(Craftax* env) {
if (craftax_clean_human_controls(env) < 0) {
return;
}
CLEAN_PROF_START();
env->agents[0].rewards[0] = 0.0f;
env->agents[0].terminals[0] = 0.0f;
int action = env->agents[0].actions[0];
Expand All @@ -1863,7 +1813,6 @@ void puf_step(Craftax* env) {
action = ACTION_NOOP;
}

CLEAN_ZONE(0);
int level = state->player_level;
int row = state->player_position[0];
int col = state->player_position[1];
Expand Down Expand Up @@ -2007,9 +1956,7 @@ void puf_step(Craftax* env) {
inv->armour[idx] = 2;
state->achievements[ACH_MAKE_DIAMOND_ARMOUR] = 1;
}
CLEAN_ZONE_END(0);

CLEAN_ZONE(1);
Rng interact_rng = rng_key(&step_rng);
int direction[2];

Expand Down Expand Up @@ -2168,9 +2115,7 @@ void puf_step(Craftax* env) {
state->chests_opened[level] |= block == BLOCK_CHEST;
}
}
CLEAN_ZONE_END(1);

CLEAN_ZONE(2);
action_to_direction(state->player_direction, direction);
row = state->player_position[0] + direction[0];
col = state->player_position[1] + direction[1];
Expand Down Expand Up @@ -2294,9 +2239,7 @@ void puf_step(Craftax* env) {
}
state->achievements[ACH_DRINK_POTION] = 1;
}
CLEAN_ZONE_END(2);

CLEAN_ZONE(3);
Rng book_rng = rng_key(&step_rng);

bool reading = action == ACTION_READ_BOOK && state->inventory.books > 0;
Expand Down Expand Up @@ -2407,9 +2350,7 @@ void puf_step(Craftax* env) {
if (direction[0] != 0 || direction[1] != 0) {
state->player_direction = action;
}
CLEAN_ZONE_END(3);

CLEAN_ZONE(4);
Rng mobs_rng = rng_key(&step_rng);

level = state->player_level;
Expand All @@ -2428,9 +2369,7 @@ void puf_step(Craftax* env) {
update_projectile_set(state, false);
rng_key(&mobs_rng);
update_projectile_set(state, true);
CLEAN_ZONE_END(4);

CLEAN_ZONE(5);
Rng spawn_rng = rng_key(&step_rng);

level = state->player_level;
Expand Down Expand Up @@ -2519,9 +2458,7 @@ void puf_step(Craftax* env) {
}
}
}
CLEAN_ZONE_END(5);

CLEAN_ZONE(6);
for (int plant = 0; plant < MAX_GROWING_PLANTS; plant++) {
if (!state->growing_plants_mask[plant]) {
continue;
Expand Down Expand Up @@ -2657,7 +2594,6 @@ void puf_step(Craftax* env) {
state->achievements[ACH_MAKE_IRON_SWORD] |= state->inventory.sword >= 3;
state->achievements[ACH_MAKE_DIAMOND_SWORD] |= state->inventory.sword >= 4;
update_log_state(env);
CLEAN_ZONE_END(6);

store_rng(state, rng_key(&step_rng));
state->timestep += 1;
Expand All @@ -2667,7 +2603,6 @@ void puf_step(Craftax* env) {
done = state->player_health <= 0.0f || state->timestep >= DEFAULT_MAX_TIMESTEPS;
} while (!done && (state->is_sleeping || state->is_resting));

CLEAN_ZONE(7);
float reward = 0.0f;
for (int i = 0; i < NUM_ACHIEVEMENTS; i++) {
int delta = state->achievements[i] - initial_achievements[i];
Expand Down Expand Up @@ -2717,11 +2652,8 @@ void puf_step(Craftax* env) {
generate_world_from_key(&env->state, world_key);
}
}
CLEAN_ZONE_END(7);

CLEAN_ZONE(8);
compute_observations(env);
CLEAN_ZONE_END(8);
}

void puf_init(Env* env, Dict* kwargs) {
Expand Down Expand Up @@ -2833,6 +2765,11 @@ static void draw_tile(int tex_id, int x, int y, int px) {
DrawTexturePro(textures, src, dst, (Vector2){0, 0}, 0.0f, WHITE);
}

static int map_lit(const State* state, int level, int wr, int wc) {
return (unsigned)wr < MAP_SIZE && (unsigned)wc < MAP_SIZE
&& state->light_map[level][wr][wc] > 12;
}

static int projectile_tex(int ptype, int dr, int dc) {
if (ptype == PROJECTILE_DAGGER) {
return TEX_PROJ_DAGGER;
Expand Down Expand Up @@ -2885,10 +2822,8 @@ static void draw_agent_obs(Craftax* env, int panel_x, int panel_y,
int wc = pc + (vc - rc);
int dst_x = grid_x + vc * px;
int dst_y = grid_y + vr * px;
int lit = (unsigned)wr < MAP_SIZE && (unsigned)wc < MAP_SIZE
&& state->light_map[level][wr][wc] > 12;
if (!lit) {
DrawRectangle(dst_x, dst_y, px, px, BLACK);
if (!map_lit(state, level, wr, wc)) {
draw_tile(BLOCK_DARKNESS, dst_x, dst_y, px);
continue;
}
int block = state->map[level][wr][wc];
Expand Down Expand Up @@ -2922,9 +2857,8 @@ static void draw_agent_obs(Craftax* env, int panel_x, int panel_y,
int col = projectiles->position[i][1];
int vr = row - pr + rr;
int vc = col - pc + rc;
int lit = (unsigned)row < MAP_SIZE && (unsigned)col < MAP_SIZE
&& state->light_map[level][row][col] > 12;
if (vr < 0 || vr >= OBS_ROWS || vc < 0 || vc >= OBS_COLS || !lit) {
if (vr < 0 || vr >= OBS_ROWS || vc < 0 || vc >= OBS_COLS
|| !map_lit(state, level, row, col)) {
continue;
}
int dr = directions[level][i][0];
Expand Down Expand Up @@ -3091,18 +3025,15 @@ void puf_render(Craftax* env) {
int dst_x = origin_x + vc * TEX_DRAW_PX;
int dst_y = vr * TEX_DRAW_PX;

int in_map = wr >= 0 && wr < MAP_SIZE && wc >= 0 && wc < MAP_SIZE;
int block = BLOCK_OUT_OF_BOUNDS;
if (in_map) {
block = env->state.map[level][wr][wc];
if (!map_lit(&env->state, level, wr, wc)) {
draw_tile(BLOCK_DARKNESS, dst_x, dst_y, TEX_DRAW_PX);
continue;
}
int block = env->state.map[level][wr][wc];
if (block < 0 || block >= NUM_BLOCK_TYPES) {
block = BLOCK_INVALID;
}
draw_tile(block, dst_x, dst_y, TEX_DRAW_PX);
if (!in_map) {
continue;
}
int item = env->state.item_map[level][wr][wc];
if (item > ITEM_NONE) {
draw_tile(TEX_ITEM_BASE + item, dst_x, dst_y, TEX_DRAW_PX);
Expand All @@ -3112,7 +3043,8 @@ void puf_render(Craftax* env) {
if (find_mob_at(&env->state, level, wr, wc, &mob_class, &slot)) {
int type_id = mobs_for_class(&env->state, level, mob_class)
->type_id[slot];
draw_tile(mob_tex_base[mob_class] + type_id, dst_x, dst_y, TEX_DRAW_PX);
draw_tile(mob_tex_base[mob_class] + type_id, dst_x, dst_y,
TEX_DRAW_PX);
}
}
}
Expand All @@ -3128,9 +3060,12 @@ void puf_render(Craftax* env) {
if (!projectiles->mask[i]) {
continue;
}
int vr = projectiles->position[i][0] - top_row;
int vc = projectiles->position[i][1] - left_col;
if (vr < 0 || vr >= RENDER_ROWS || vc < 0 || vc >= RENDER_COLS) {
int wr = projectiles->position[i][0];
int wc = projectiles->position[i][1];
int vr = wr - top_row;
int vc = wc - left_col;
if (vr < 0 || vr >= RENDER_ROWS || vc < 0 || vc >= RENDER_COLS
|| !map_lit(&env->state, level, wr, wc)) {
continue;
}
int dr = directions[level][i][0];
Expand All @@ -3153,11 +3088,6 @@ void puf_render(Craftax* env) {
draw_tile(player_tex, origin_x + half_c * TEX_DRAW_PX,
half_r * TEX_DRAW_PX, TEX_DRAW_PX);

if (env->state.light_level < 1.0f) {
unsigned char alpha = (unsigned char)((1.0f - env->state.light_level) * 140.0f);
DrawRectangle(origin_x, 0, view_w, view_h, (Color){0, 0, 40, alpha});
}

int floor_bar_h = 16;
int cell_w = view_w / NUM_LEVELS;
DrawRectangle(origin_x, 0, view_w, floor_bar_h, (Color){18, 18, 18, 230});
Expand Down