Clean pufferl.py, delete Neptune

.dockerignore

@@ -7,7 +7,6 @@ build
 *.egg-info
 experiments
 wandb
-.neptune
 .pytest_cache
 .ruff_cache
 

.gitignore

@@ -149,7 +149,6 @@ checkpoints/
 experiments/
 benchmark*/
 wandb/
-.neptune/
 raylib*/
 box2d*/
 

pufferlib/config/puffer_drive.yaml

@@ -8,9 +8,6 @@ wandb_group: debug
 # Unique run identifier. Doubles as the logger run id, so relaunching with the
 # same run_name and train.data_dir resumes that run instead of starting a new one.
 run_name: default_run
-neptune: false
-neptune_name: pufferai
-neptune_project: ablations
 tb: false
 local_rank: 0
 tag: null

pufferlib/config_schema.py

@@ -449,9 +449,6 @@ class PufferDriveConfig:
     wandb_project: str = _constrained_field(NONEMPTY_STRING_CONSTRAINT)
     wandb_group: str = _constrained_field(NONEMPTY_STRING_CONSTRAINT)
     run_name: str = _constrained_field(NONEMPTY_STRING_CONSTRAINT)
-    neptune: bool = MISSING
-    neptune_name: str = MISSING
-    neptune_project: str = MISSING
     tb: bool = MISSING
     local_rank: int = MISSING
     tag: str | None = MISSING
@@ -540,8 +537,8 @@ def _validate_string_selection(value, context, path, *, allow_none=True):
 
 def _validate_cross_field_constraints(config, context):
     """Validate relationships and context-dependent rules spanning config fields."""
-    if config["load_id"] is not None and not (config["wandb"] or config["neptune"]):
-        _raise_config_error(context, "load_id", "requires wandb or neptune")
+    if config["load_id"] is not None and not config["wandb"]:
+        _raise_config_error(context, "load_id", "requires wandb")
 
     env = config["env"]
     if env["min_agents_per_env"] > env["max_agents_per_env"]:

pufferlib/ocean/drive/drive.h

@@ -187,6 +187,7 @@ struct Drive {
     int *tracks_to_predict;
     // Simulation
     int timestep;
+    int autoreset_pending;
     int init_step;
     float dt;
     float base_max_speed_mps;
@@ -4436,6 +4437,7 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) {
 #include "idm.h"
 
 void c_reset(Drive *env) {
+    env->autoreset_pending = 0;
     if (env->timestep == 0) {
         for (int i = 0; i < env->num_total_agents; i++) {
             copy_pose_to_prev(&env->agents[i]);
@@ -4549,18 +4551,10 @@ void c_step(Drive *env) {
     memset(env->rewards, 0, env->active_agent_count * sizeof(float));
     memset(env->terminals, 0, env->active_agent_count * sizeof(unsigned char));
     memset(env->truncations, 0, env->active_agent_count * sizeof(unsigned char));
-
-    // Update masks for stopped/removed agents
-    for (int i = 0; i < env->active_agent_count; i++) {
-        int agent_idx = env->active_agent_indices[i];
-        Agent *a = &env->agents[agent_idx];
-        if (a->stopped || a->removed || a->is_blind_partner || a->is_phantom_braker) {
-            env->masks[i] = 0;
-        } else {
-            env->masks[i] = 1;
-        }
+    if (env->autoreset_pending) {
+        c_reset(env);
+        return;
     }
-
     env->timestep++;
 
     // -> 1. Apply actions and move agents
@@ -4615,8 +4609,10 @@ void c_step(Drive *env) {
 
     // Mark terminals for stopped or removed agents
     for (int i = 0; i < env->active_agent_count; i++) {
-        int agent_idx = env->active_agent_indices[i];
-        if (env->agents[agent_idx].stopped || env->agents[agent_idx].removed) {
+        Agent *agent = &env->agents[env->active_agent_indices[i]];
+        // Masks describe the next action's eligibility; terminal rewards belong to the previous action.
+        env->masks[i] = !agent->stopped && !agent->removed && !agent->is_blind_partner && !agent->is_phantom_braker;
+        if (agent->stopped || agent->removed) {
             env->terminals[i] = 1;
         }
     }
@@ -4656,7 +4652,10 @@ void c_step(Drive *env) {
             env->eval_episode_done = 1;
             return;
         }
-        c_reset(env);
+        // Expose the terminal state for value bootstrap; the next call resets without applying its action.
+        compute_observations(env);
+        memset(env->masks, 0, env->active_agent_count * sizeof(unsigned char));
+        env->autoreset_pending = 1;
         return;
     }
 
@@ -4699,6 +4698,8 @@ void c_step(Drive *env) {
         if (!regen_ok) {
             invalidate_agent(agent);
             agent->removed = 1;
+            env->masks[i] = 0;
+            env->terminals[i] = 1;
         }
     }
 }

pufferlib/ocean/drive/drive.py

@@ -546,6 +546,7 @@ def reset(self, seed=None):
         else:
             binding.vec_reset(self.c_envs)
         self.tick = 0
+        self._resample_pending = False
         self.truncations[:] = 0
         if self.capture_replay:
             self._initialize_replay_captures()
@@ -560,17 +561,23 @@ def step(self, actions):
         if self.capture_replay:
             self._capture_replay_step()
         self.actions[:] = actions
-        binding.vec_step(self.c_envs)
-        self.tick += 1
+        if not self._resample_pending:
+            binding.vec_step(self.c_envs)
+            self.tick += 1
         info = []
         # vec_log is the training aggregate; it resets env->log, which eval reads
         # per episode, so it must not run in eval mode.
-        if not self.eval_mode and self.tick % self.report_interval == 0:
+        if not self.eval_mode and not self._resample_pending and self.tick % self.report_interval == 0:
             log = binding.vec_log(self.c_envs, self.num_agents)
             if log:
                 info.append(log)
-                # print(log)
         if self.tick > 0 and self.resample_frequency > 0 and self.tick % self.resample_frequency == 0:
+            if not self.eval_mode and not self._resample_pending:
+                self._resample_pending = True
+                self.truncations[:] = 1
+                self.masks[:] = 0
+                return (self.observations, self.rewards, self.terminals, self.truncations, info)
+            self._resample_pending = False
             self.tick = 0
             will_resample = 1
             if will_resample:
@@ -643,8 +650,10 @@ def step(self, actions):
                 binding.vec_reset(self.c_envs)
                 if self.capture_replay:
                     self._initialize_replay_captures()
-                # Map resampling is an external reset boundary (dataset/map switch). Treat as truncation.
-                self.truncations[:] = 1
+                self.truncations[:] = self.eval_mode
+                if not self.eval_mode:
+                    self.rewards[:] = 0
+                    self.terminals[:] = 0
         return (self.observations, self.rewards, self.terminals, self.truncations, info)
 
     def get_global_agent_state(self):

pufferlib/ocean/evaluation_utils/eval_replay.py

@@ -164,7 +164,7 @@ def write_pending(self):
         self.pending_replays = []
 
 
-def _render_eval_replays(episode_summaries, out_dir, keep_zlib_replays):
+def render_eval_replays(episode_summaries, out_dir, keep_zlib_replays):
     """Render captured eval replays as navigable HTML pages plus an index."""
     render_dir = os.path.join(out_dir, "rendered_replays")
     os.makedirs(render_dir, exist_ok=True)

pufferlib/ocean/evaluation_utils/evaluation_utils.py

@@ -56,7 +56,7 @@ def _load_yaml_mapping(path, label):
     return _require_mapping(value, label)
 
 
-def _resolve_map_indices(map_dir, map_names):
+def resolve_map_indices(map_dir, map_names):
     """Map each logged map name back to its index in the sorted .bin map set."""
     if os.path.isfile(map_dir) and str(map_dir).endswith(".bin"):
         map_files = [map_dir]
@@ -270,7 +270,7 @@ def build_benchmark_args(base_args, benchmark, environment_config, cli_overrides
     )
 
 
-def _plan_benchmark_eval_workers(args, num_scenarios, num_workers, scenario_length, capture_replay=False):
+def plan_benchmark_eval_workers(args, num_scenarios, num_workers, scenario_length, capture_replay=False):
     """One disjoint contiguous map window per worker; together they cover the set once."""
     scenarios_per_worker, remainder = divmod(num_scenarios, num_workers)
     worker_env_kwargs = []
@@ -290,7 +290,7 @@ def _plan_benchmark_eval_workers(args, num_scenarios, num_workers, scenario_leng
     return worker_env_kwargs, max_scenarios_per_worker * scenario_length
 
 
-def _plan_failure_replay_workers(args, map_seed_pairs, num_workers, scenario_length):
+def plan_failure_replay_workers(args, map_seed_pairs, num_workers, scenario_length):
     """Split the (map, seed) pairs across workers; each worker cycles through its
     pairs in fit-aware batches (num_agents from config bounds a batch)."""
     pairs_per_worker, remainder = divmod(len(map_seed_pairs), num_workers)
@@ -347,7 +347,7 @@ def _build_eval_report(episode_summaries, num_scenarios):
     return df, summary
 
 
-def _write_eval_reports(episode_summaries, out_dir, num_scenarios):
+def write_eval_reports(episode_summaries, out_dir, num_scenarios):
     """Write a per-episode metrics CSV and a JSON of metric averages to out_dir."""
     import json