feat: add async policy trainer primitives - #480
Conversation
Add the standalone building blocks for the experimental asynchronous policy trainer under areno.experimental.async_policy. - AsyncTrainBatch: batch metadata carrying the rollout policy version, timing, and size fields used for staleness control and metrics. - AsyncTrainBatchQueue: a bounded queue with backpressure, put/get timeouts, close, and error propagation. - staleness_delta/is_stale: policy-version staleness helpers. All three are pure Python and covered by 19 CPU tests; algorithm registration lands together with the async trainer itself.
|
@xsuler Hi, I'd really appreciate a review on this PR when you have a chance. Quick context so this PR doesn't look like a one-off: I'm applying to OSPP 2026 with AReno as my target project, and I'm landing that work as a sequence of small, independently-reviewable PRs so each one stays easy to audit. This first Let me know if this structure works for you, or if you'd prefer I bundle these primitives into the trainer PR instead. |
What does this PR do?
Adds the standalone building blocks for an experimental asynchronous policy trainer under
areno.experimental.async_policy. The async trainer will overlaprollout/reward with training on single-node multi-GPU; these primitives — a bounded batch queue, batch metadata, and staleness control — are the first self-contained piece, landable and testable before any trainer integration. No algorithms are registered yet.The change adds:
AsyncTrainBatch(batch.py): a slotted dataclass holding the materializedTrainSequences plus the rollout policy version and per-stage timing/size fields used for staleness control and metrics.AsyncTrainBatchQueue(queue.py): a bounded FIFO queue with backpressure, per-operationput/gettimeouts,close()to unblock waiters, andreport_error()to propagate a worker failure to the next operation. It usesthreading.Conditionrather thanqueue.Queuebecause the stdlib queue has no close/unblock or error-propagation semantics.staleness_delta/is_stale(staleness.py): helpers that compare a batch's rollout policy version against the current training policy version.tests/test_async_policy_cpu.py: 19 CPU tests covering FIFO ordering, backpressure, timeouts, close, error propagation, and staleness.A runnable example and an async-vs-sync evaluation report will accompany the async trainer itself; this PR is intentionally limited to the queue, batch, and staleness primitives.
Related issue
Fixes #487
Type of change
How was it tested?
Local CPU validation on the final commit:
Checklist
Breaking change details
None — purely additive code under areno.experimental/; no areno/api/, CLI, or stable-surface changes.