[RFC] Add a general monitoring and logger attachment API for collectors and replay buffers

[vmoens]

Summary

Add a general monitoring API that lets users attach a TorchRL logger, such as Weights & Biases, TensorBoard, CSV, MLflow, or TrackIO, to any collector or replay buffer and periodically log lightweight operational statistics.

The API should work across:

• in-process collectors and replay buffers;
• multiprocessing collectors and shared replay buffers;
• RayCollector and RayReplayBuffer;
• iterator-driven collection and autonomous .start() collection;
• local loggers and logger services created with use_ray_service=True.

A tentative direction is to separate:

1. the object that produces a cheap statistics snapshot;
2. the monitor that schedules, aggregates, and namespaces snapshots;
3. the existing Logger that consumes metrics.

The proposed default would be an external pull-based monitor. Event-driven observers could be added for metrics that require access to a collected batch or exact operation boundaries.

Motivation

Collectors and replay buffers increasingly run independently of the training loop. In particular:

• a collector may run asynchronously through .start();
• a distributed collector may contain several remote workers;
• a replay buffer may be a Ray actor receiving writes from several collectors;
• the training loop may sample the buffer without directly observing collection batches.

In these configurations, logging only from the trainer is insufficient. It should be possible to monitor questions such as:

• Is the collector still producing frames?
• What is its recent collection throughput?
• Are all workers alive?
• How many frames or batches has each worker produced?
• How full is the replay buffer?
• How quickly is the buffer being written and sampled?
• Is the sample rate keeping up with the write rate?
• Is an asynchronous queue or prefetch queue growing?
• Which policy version is being used by collectors?
• Are remote services healthy?

This should not require users to manually poll each implementation-specific attribute or initialize a separate W&B run inside every worker.

Existing building blocks

TorchRL already has several pieces that make this feasible:

• Logger.log_metrics() accepts scalar mappings and TensorDict inputs.
• Logger implementations can run as a Ray actor with use_ray_service=True.
• Replay buffers expose useful state such as len(buffer) and write_count.
• RayReplayBuffer already delegates method and property access to its actor.
• Collectors expose pre- and post-collection hooks.
• Distributed collectors already contain worker fan-out and aggregation machinery.

There are still important gaps:

• There is no common lightweight stats() or metrics() surface.
• state_dict() is not suitable for monitoring. It can contain policy state, environment state, or the replay-buffer storage.
• Collector hooks are batch-oriented and do not form a universal monitoring surface across autonomous collection and replay-buffer write paths.
• The current hook API is a single slot, so independent features such as profiling and logging are not naturally composable.
• Replay buffers do not expose general add, extend, sample, or empty observers.
• Calling the current Ray logger service from a hot path is synchronous because the proxy waits with ray.get.
• Passing an owning logger-service proxy into several workers requires clear ownership and shutdown semantics.
• Distributed collection needs explicit aggregate-versus-per-worker behavior.

Goals

• Attach one logger to any number of collectors and replay buffers.
• Support wall-clock and logical-counter intervals.
• Work without putting remote logger calls in collection or sampling hot paths by default.
• Return lightweight, serializable, side-effect-free statistics.
• Support local, multiprocessing, and Ray implementations with the same user-facing API.
• Namespace metrics so several monitored objects can share one logger.
• Support aggregate and per-worker views.
• Allow custom statistics providers.
• Provide a future path for batch-derived and event-derived metrics.
• Keep monitoring lifecycle and logger lifecycle explicit.
• Avoid unbounded logging queues and backpressure on collection.

Non-goals

• Logging a collector or replay buffer's complete state_dict().
• Automatically deciding which arbitrary internal state is meaningful.
• Scanning the full replay-buffer contents on every interval.
• Guaranteeing an exact globally synchronized "every N frames" event across independent workers without a centralized counter.
• Making logging part of compiled replay-buffer or collector hot paths.

Design principle: separate source, scheduling, and sink

Three concepts appear useful:

1. Statistics source

Collectors and replay buffers expose a cheap snapshot:

collector.stats()
replay_buffer.stats()

The result should contain cumulative counters and gauges, not heavyweight state.

2. Monitor

A monitor decides:

• when to request a snapshot;
• which counter is the logical step;
• how to compute rates from counter deltas;
• how to namespace metrics;
• how to aggregate worker snapshots;
• how to handle errors and lifecycle.

3. Logger

The existing logger remains the output sink. It should not need collector- or replay-buffer-specific knowledge.

API alternatives

Alternative A: constructor arguments

collector = RayCollector(
    ...,
    logger=logger,
    log_interval=10_000,
)

buffer = RayReplayBuffer(
    ...,
    logger=logger,
    log_interval=10_000,
)

Advantages:

• Simple and discoverable.
• Minimal setup for common cases.

Drawbacks:

• Adds parameters to every collector and replay-buffer constructor.
• log_interval is ambiguous: seconds, frames, writes, samples, or batches.
• Custom metrics, aggregation, namespace, and lifecycle introduce many more parameters.
• Logger serialization and ownership become implicit.
• Config companions would need to mirror the new constructor arguments.

This could eventually be convenience sugar, but it is probably not the best foundational abstraction.

Alternative B: event observers

collector.add_observer(
    LoggerObserver(
        logger,
        schedule=Every.counter("frames", 10_000),
    )
)

buffer.add_observer(
    LoggerObserver(
        logger,
        schedule=Every.counter("write_count", 100_000),
    )
)

Possible events include:

• collector: pre_collect, post_collect, shutdown;
• replay buffer: post_add, post_extend, post_sample, empty, shutdown.

Advantages:

• Precise event-based triggering.
• Access to the collected or sampled batch.
• Extensible beyond logging.

Drawbacks:

• Adds work to hot paths.
• Requires composable observer lists rather than one callback property.
• Requires clear worker, coordinator, and aggregation semantics.
• A synchronous service call would block the source actor.
• Custom observers must be serialized into workers.
• Exact global thresholds across workers require centralized coordination.

This seems appropriate for batch-dependent metrics, but more complex than necessary for structural monitoring.

Alternative C: wrappers

buffer = MonitoredReplayBuffer(
    buffer,
    logger=logger,
    schedule=Every.seconds(5),
)

Advantages:

• Requires fewer core changes.
• Useful for prototyping.

Drawbacks:

• Forwarding __iter__, __len__, properties, and implementation-specific methods is fragile.
• Changes identity and type expectations.
• Becomes awkward with shared objects and Ray proxies.

Alternative D: external pull-based monitor

monitor = LoggerMonitor(logger, poll_interval=1.0)

monitor.watch(
    collector,
    name="collector",
    schedule=Every.counter("frames", 10_000),
)

monitor.watch(
    replay_buffer,
    name="replay_buffer",
    schedule=Every.seconds(5),
    step="write_count",
)

with monitor:
    collector.start()
    run_training()

Advantages:

• Logging is not called from collector or replay-buffer hot paths.
• The logger remains in the driver or in its own logger-service actor.
• One monitor can watch several objects.
• Local and remote objects can share the same API.
• A Ray object can return a coherent snapshot in a single RPC.
• Rates can be computed from cumulative counter deltas.
• Works naturally with autonomous .start() collection.

Drawbacks:

• Counter thresholds are observed rather than executed at the exact operation.
• Batch contents are not available after the fact.
• Requires a background thread/task or explicit monitor.step() calls.

Tentative recommendation

Use an external pull-based LoggerMonitor as the core abstraction and add event observers only for metrics that require batch data or exact operation timing.

A complete example could look like:

from torchrl.record import WandbLogger
from torchrl.record.monitoring import Every, LoggerMonitor

logger = WandbLogger(
    exp_name="experiment",
    project="torchrl",
    use_ray_service=True,
)

monitor = LoggerMonitor(
    logger,
    poll_interval=1.0,
)

monitor.watch(
    collector,
    name="collector",
    schedule=Every.counter("frames", 10_000),
    workers="aggregate",
)

monitor.watch(
    replay_buffer,
    name="replay_buffer",
    schedule=Every.seconds(5),
    step="write_count",
)

with monitor:
    collector.start()
    run_training()

A convenience method could later be added:

handle = replay_buffer.attach_logger(
    logger,
    schedule=Every.seconds(5),
)

This should be thin sugar around the monitor abstraction rather than a separate implementation.

Proposed statistics contract

A stats() call should be:

• cheap;
• side-effect free;
• safe to call concurrently;
• serializable;
• limited to scalar tensors, Python numeric values, booleans, and small metadata;
• cumulative where possible;
• independent from state_dict();
• available before, during, and after iteration when meaningful.

The initial return type could be a mapping accepted by Logger.log_metrics(). A scalar TensorDict may also be worth considering for consistency with TorchRL's TensorDict-first APIs.

The source should expose cumulative counters. The monitor should compute rates. This avoids placing timers and rolling windows in every source implementation.

Counter decreases may occur after a reset, empty(empty_write_count=True), or load_state_dict(). The monitor should treat a decrease as a new baseline rather than report a negative rate.

Candidate collector statistics

Structural metrics could include:

frames
batches
requested_frames_per_batch
total_frames
completed
workers
workers_alive
pending_batches
policy_version

Monitor-derived metrics could include:

frames_per_second
batches_per_second
seconds_since_last_progress

Optional per-worker metrics could be namespaced as:

collector/worker_0/frames
collector/worker_0/frames_per_second
collector/worker_0/policy_version

Data-dependent metrics such as reward mean, episode return, done count, and trajectory length should be optional batch observers rather than part of the default structural snapshot.

Candidate replay-buffer statistics

Structural metrics could include:

size
capacity
utilization
write_count
sample_calls
samples_returned
prefetch_queue_size
initialized

Monitor-derived metrics could include:

writes_per_second
samples_per_second
sample_to_write_ratio
seconds_since_last_write

Sampler-specific statistics could be provided when they are cheap, for example current beta or an O(1) priority summary. The default must not walk the storage or priority tree.

Adding sampling counters needs care around torch.compile and compilable replay buffers. If a Python-side counter would introduce a graph break, it may need to remain optional, live outside the compiled region, or use a compile-friendly representation.

Scheduling semantics

The API should distinguish wall-clock and logical-counter schedules:

Every.seconds(5)
Every.counter("frames", 10_000)
Every.counter("write_count", 100_000)
Every.counter("sample_calls", 100)

Suggested semantics:

• Counter schedules log when the latest snapshot crosses the next threshold.
• If a counter jumps across several thresholds, log the latest snapshot once by default rather than invent historical values.
• Optional catch-up behavior can be considered, but it cannot reconstruct snapshots at thresholds that were not observed.
• log_on_start and log_on_close should be explicit options.
• Wall-clock polling and logging intervals should be separable when counter schedules are used.
• A manual mode should be available for deterministic loops and tests:

monitor = LoggerMonitor(logger, background=False)
...
monitor.step()

Distributed and Ray semantics

RayReplayBuffer

RayReplayBuffer.stats() should perform one RPC to the replay-buffer actor. It should not make separate remote calls for length, write count, capacity, and sampling counters.

The underlying actor can collect a coherent snapshot under its existing synchronization mechanism.

RayCollector

RayCollector.stats() should query remote collectors concurrently.

Possible worker modes:

collector.stats(workers="aggregate")
collector.stats(workers="per_worker")
collector.stats(workers="both")

Aggregation must be metric-aware:

• frame and batch counters are summed;
• worker health is counted;
• throughput may be summed or recomputed from aggregate deltas;
• gauges should not be blindly averaged without defined semantics.

When a distributed collector writes directly to a shared RayReplayBuffer, the replay buffer's write_count is likely the canonical global production counter. Coordinator-side collector counters may lag or may not be updated in free-running modes.

Multiprocessing collectors

The same user-facing API should apply. The coordinator may aggregate worker snapshots through existing pipes, while autonomous free-running collection may be monitored through the shared replay buffer when that is the authoritative counter.

Logger-service considerations

Pull monitoring avoids passing the logger proxy into collection workers. Event-driven source-side logging will require additional service semantics.

Nonblocking submission

Calling ray.get from a collection, write, or sample path would introduce cross-actor latency. A future API could distinguish synchronous logging from submission:

logger.submit_metrics(metrics, step=step)
logger.flush()

Alternatively, the monitor can own a bounded queue and perform existing synchronous logger calls from its own thread or actor.

Ownership

The object that creates a Ray logger service should own its lifecycle. Serialized handles passed to workers should be explicitly non-owning. Shutdown should not depend on destruction of arbitrary proxy copies.

An owning service plus a lightweight picklable handle may be clearer than passing the same owning proxy everywhere.

Backpressure

Monitoring must not create an unbounded backlog if the backend is slow. For state snapshots, a bounded queue with coalescing is appropriate: the most recent state usually supersedes older pending state.

Error behavior should also be configurable, for example "warn", "raise", or "ignore", with "warn" as a reasonable background-monitor default.

Step handling

Collector frames, replay-buffer writes, and replay-buffer samples may advance independently. They should not compete for a single global W&B step.

Namespaced metric groups such as:

collector/...
replay_buffer/...
replay_buffer_samples/...

fit the existing per-prefix step handling in WandbLogger.

Custom statistics

Users should be able to add cheap domain-specific gauges without subclassing every collector or replay buffer.

Possible forms include:

monitor.watch(
    replay_buffer,
    name="replay_buffer",
    stats=ReplayBufferStats(
        extra={
            "policy_lag": policy_lag_provider,
        }
    ),
)

or:

monitor.watch(
    replay_buffer,
    stats_fn=my_stats_fn,
)

For remote objects, the execution location must be explicit:

• a driver-side function can only inspect the public proxy;
• a source-side provider must be serializable and execute in the actor or worker;
• only the resulting scalar mapping should cross the process boundary.

Batch-derived metrics

Polling is not sufficient for metrics derived from the latest collection batch. A later observer API could support:

collector.add_observer(
    LoggerObserver(
        logger,
        schedule=Every.counter("frames", 10_000),
        metrics=CollectorBatchMetrics(
            reward_key=("next", "reward"),
            episode_return=True,
        ),
        scope="aggregate",
    )
)

The reduction should happen next to the data. Workers should send only reduced scalars to an aggregation monitor or logger service.

The observer API should support multiple registered observers so profiling, logging, and user callbacks can coexist.

Possible implementation phases

Phase 1: statistics surfaces

• Add stats() to BaseCollector and ReplayBuffer.
• Define the initial stable structural metrics.
• Add one-RPC delegation for RayReplayBuffer.
• Add concurrent worker fan-out for RayCollector.
• Add tests for local and remote snapshots.

Phase 2: pull monitor

• Add LoggerMonitor.watch().
• Support wall-clock and counter schedules.
• Support background and manual .step() modes.
• Compute rates from cumulative counters.
• Add namespacing, start/final snapshots, and reset detection.
• Add bounded/coalescing logging.
• Support aggregate and per-worker views.

Phase 3: richer counters and providers

• Add replay-buffer sample counters where compile-friendly.
• Add collector health and queue metrics.
• Add custom statistics providers.
• Add explicit logger service flush() and close() behavior if needed.

Phase 4: event observers

• Replace single-purpose callback slots with composable observers where appropriate.
• Add replay-buffer operation events.
• Add batch-derived collector metrics.
• Add nonblocking source-side submission and non-owning logger-service handles.

Testing considerations

The behavior should be exercised with:

• Collector;
• synchronous and asynchronous multiprocessing collectors;
• RayCollector;
• base and TensorDict replay buffers;
• prioritized replay buffers where applicable;
• RayReplayBuffer;
• iterator-driven collection;
• autonomous .start() collection;
• a collector writing directly into a replay buffer;
• local CSV logging;
• a logger running through use_ray_service=True;
• counter reset and state restoration;
• slow logger backpressure;
• monitor shutdown while targets are still running;
• target shutdown while the monitor is still running.

Monitoring tests should also verify that disabled monitoring adds no hot-path work and that enabled pull monitoring does not alter collected or sampled data.

Open questions

1. Should the public method be called stats(), metrics(), get_stats(), or snapshot_metrics()?
2. Should the return type be a scalar TensorDict or a plain scalar mapping?
3. Should background monitoring be opt-in, with manual .step() as the default?
4. Should counter schedules poll at a configurable frequency or receive lightweight counter notifications?
5. What is the default worker view: aggregate, per-worker, or both?
6. Which collector counter should be canonical in autonomous distributed modes?
7. Which replay-buffer sample counters can be added without harming compilation?
8. Should attach_logger() exist as convenience sugar, or is LoggerMonitor.watch() sufficient?
9. Should a logger service expose separate owning and non-owning handle types?
10. Should event observers be part of the same monitor API or a separate callback system?
11. Which metrics should be stable public API versus best-effort implementation-specific extras?
12. Should monitor state be checkpointable so interval baselines survive resume?

Initial acceptance criteria

A useful first milestone would allow the following pattern for both local and Ray objects:

logger = WandbLogger(
    exp_name="experiment",
    project="torchrl",
    use_ray_service=True,
)

with LoggerMonitor(logger) as monitor:
    monitor.watch(
        collector,
        name="collector",
        schedule=Every.counter("frames", 10_000),
    )
    monitor.watch(
        replay_buffer,
        name="replay_buffer",
        schedule=Every.seconds(5),
        step="write_count",
    )

    collector.start()
    run_training()

with these guarantees:

• the logger backend is initialized once;
• the collector and replay buffer do not perform synchronous logger RPCs;
• Ray snapshots use bounded RPCs rather than one RPC per metric;
• metrics are namespaced and use independent logical steps;
• monitoring can be stopped cleanly without shutting down user-owned targets;
• logger shutdown is explicit;
• default statistics are cheap and never include full object state.

[vmoens]

cc @theap06
Do you want to give this a shot?

[coder-jayp]

@vmoens I'd like to take this up.

Alternative D (External pull-based monitor) looks like the best approach.

I plan to start with Phase 1 by adding a stats() method to BaseCollector, ReplayBuffer, and their Ray counterparts to expose the basic structural metrics (frames, batches, capacity, etc).

[theap06]

@coder-jayp I think there is a lot of existing infrastructure that can be abstracted away with logging here. These are the reference PRs for phase 1 and 2: #4030, #4031, #4032, #4033. I think this can collapse the three per-backend stats() overrides into one. also, would like to see a wandb and ray collector example + results. Feel free to ping me if you want to discuss design.

[coder-jayp]

@theap06 I'll start by putting together the WandB + RayCollector script to validate the baseline results.

Regarding collapsing the three stats() overrides—could you clarify what you mean by abstracting it away with logging? Do you mean workers should log their own metrics directly to a logger service, completely removing the need for the coordinator to poll them? Or did you have a different approach in mind?