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

[eugenevinitsky]

What

vec_log reported a mean over completed episodes, which weighted each agent by how often it finished. Agents that crashed early finished more often and so dominated every environment/* curve. This changes the aggregation to a mean over agent slots: one weight per agent, regardless of completion frequency.

[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.

[WaelDLZ]

@eugenevinitsky The smoke train test didn't go through, I guess the test itself might require a modification to take your changes into account

[eugenevinitsky]

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 71269746).

Differences from the original:

• log_ema_alpha lives in puffer_drive.yaml (drive.ini is being removed in Remove drive.ini; puffer_drive.yaml is the single config source #539).
• Dropped the write-only per_agent_log_count array.
• The test moved to tests/unit_tests/ so CI actually runs it.

Local verification: 84/84 unit tests (including the three behavioral logging tests) and the full C suite pass.