vec_log per-agent population mean - average over agents instead of over completed episodes

[eugenevinitsky]

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.

[Copilot]

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_log and call it from drive.py before vec_log at each report_interval.
• Relax vec_log’s emission gate in env_binding.h from aggregate.n < num_agents to aggregate.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.

[vcharraut]

Do you have an example of the same seed with and without this PR ?

[eugenevinitsky]

great question, will launch two runs and shre

[eugenevinitsky]

@copilot resolve the merge conflicts in this pull request

[Copilot]

@copilot resolve the merge conflicts in this pull request

Conflicts are resolved in commit 63093b1 by merging emerge/temp_training into this branch and resolving the two modify/delete conflicts (.github/workflows/utest.yml, tests/test_drive_train.py) in favor of the base branch file moves.