vec_log per-agent population mean - average over agents instead of over completed episodes - #459
vec_log per-agent population mean - average over agents instead of over completed episodes#459eugenevinitsky wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR changes how Drive environment metrics are aggregated in vec_log to remove completion-rate bias by switching from “mean over completed episodes” to a two-stage “per-agent window mean, then population mean across agent slots” approach. It also adds a Drive-specific vec_prepare_log step to drain per-agent accumulators into env->log immediately before vec_log.
Changes:
- Add per-agent log accumulators in Drive and a
prepare_log()drain step to produce per-agent population means. - Add a new C binding
vec_prepare_logand call it fromdrive.pybeforevec_logat eachreport_interval. - Relax
vec_log’s emission gate inenv_binding.hfromaggregate.n < num_agentstoaggregate.n < 1(API arg retained but unused).
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| pufferlib/ocean/env_binding.h | Changes vec_log emission gating to emit with any data (ignores num_agents arg). |
| pufferlib/ocean/drive/drive.py | Calls new binding.vec_prepare_log() prior to binding.vec_log(). |
| pufferlib/ocean/drive/drive.h | Adds per-agent log sum/count buffers; introduces ensure_per_agent_log_capacity() and prepare_log(); rewrites add_log() to accumulate per-agent. |
| pufferlib/ocean/drive/binding.c | Exposes vec_prepare_log Python binding to call prepare_log() across VecEnv envs. |
Comments suppressed due to low confidence (1)
pufferlib/ocean/drive/binding.c:55
- vec_prepare_log_py doesn’t validate its argument count. Most other vec_* entrypoints in env_binding.h explicitly check PyTuple_Size and raise a clear TypeError; without this, accidental extra args will be silently ignored and missing args will lead to confusing errors inside unpack_vecenv.
prepare_log((Drive *) vec->envs[i]);
}
Py_RETURN_NONE;
}
static int my_put(Env *env, PyObject *args, PyObject *kwargs) {
PyObject *obs = PyDict_GetItemString(kwargs, "observations");
if (!PyObject_TypeCheck(obs, &PyArray_Type)) {
PyErr_SetString(PyExc_TypeError, "Observations must be a NumPy array");
return 1;
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| static void ensure_per_agent_log_capacity(Drive *env, int needed) { | ||
| if (env->per_agent_log_capacity >= needed) { | ||
| return; | ||
| } | ||
| int old_cap = env->per_agent_log_capacity; | ||
| env->per_agent_log_sum = (Log *) realloc(env->per_agent_log_sum, needed * sizeof(Log)); | ||
| env->per_agent_log_count = (int *) realloc(env->per_agent_log_count, needed * sizeof(int)); | ||
| memset(&env->per_agent_log_sum[old_cap], 0, (needed - old_cap) * sizeof(Log)); | ||
| memset(&env->per_agent_log_count[old_cap], 0, (needed - old_cap) * sizeof(int)); | ||
| env->per_agent_log_capacity = needed; | ||
| } |
| // Emit whenever any env has data. With Drive's per-agent prepare_log | ||
| // path, aggregate.n is the cross-env count of agents that contributed | ||
| // a window-mean (not completed-episode count), so the meaningful gate | ||
| // is "at least one contribution." Other ocean envs that don't run | ||
| // prepare_log retain completed-episode-count semantics and now emit | ||
| // smaller batches more often — the Python-side mean_and_log | ||
| // (pufferl.py) re-averages across emissions in its rate-limit window. | ||
| if (aggregate.n < 1) { | ||
| return dict; |
eugenevinitsky
changed the title
vcharraut
left a comment
vcharraut
left a comment
There was a problem hiding this comment.
Do you have an example of the same seed with and without this PR ?
|
great question, will launch two runs and shre |
|
@copilot resolve the merge conflicts in this pull request |
Conflicts are resolved in commit |
|
@eugenevinitsky The smoke train test didn't go through, I guess the test itself might require a modification to take your changes into account |
Reworked onto current 3.0 (the original branch predated the goal-system rewrite and metrics refactor). add_log no longer sums completed episodes into env->log — that weighted agents by completion frequency, so short-episode agents (early crashes) dominated the reported means and curves only looked clean after a resample-forced synchronized reset. Each agent slot now keeps an EMA of its completed-episode Log (seeded on first completion, then slot = alpha*slot + (1-alpha)*episode); prepare_log sums the seeded slots into env->log on demand without resetting them, and vec_log's divide-by-n yields a population mean with one weight per agent. binding.vec_prepare_log drains the slots right before each vec_log call. The vec_log gate becomes n >= 1 (the old n >= num_agents throttle made no sense for persistent per-agent state), and the eval-mode branch emits one dict per env with data instead of a fixed-size list with holes. log_ema_alpha is exposed end-to-end (drive.h, binding, drive.py, puffer_drive.yaml; default 0.707 ~ 2-episode half-life). Behavioral tests cover the emission gate, steady-state n == num_agents, and EMA persistence across emissions without new completions. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
63093b1 to
8b543a8
Compare
|
Rebased as a clean re-application onto current 3.0 (the old branch predated the goal-system rewrite #512, the metrics refactor, and the Hydra config switch — a mechanical rebase conflicted on every hunk, so the branch was force-pushed with the change rebuilt on top of Differences from the original:
Local verification: 84/84 unit tests (including the three behavioral logging tests) and the full C suite pass. |
Previously we were logging anytime N agents finished, which biased the metrics towards agents that finished frequently. Now we report an average over all the agent slots.
Summary
We use per-agent EMA-based aggregation in `vec_log` so the cross-agent mean reported in wandb has one weight per agent regardless of completion frequency — no per-emission bias and no multi-emission residual bias.
For each agent slot the env keeps a smoothed `Log` state:
```
slot ← α · slot + (1 − α) · episode_log # on every completion (after first)
slot ← episode_log # on the slot's first completion
```
`prepare_log` then sums slots flagged `has_data`, sets `env->log.n = num_with_data`, and does not reset. `vec_log`'s sum-across-envs / divide-by-`aggregate.n` step then produces a population mean: every agent that has ever completed contributes one term, every emission.
Why EMA?
We need some way of providing a good estimate of the average performance of the agents. We don't want to accumulate forever, as that'd be slow to update, and we could use a fixed size number of logs per agent, but this is a nice compromise.