Logging¤
jaxnasium.algorithms.core.scan_callback(func: Callable | None = None, callback_fn: Callable | Literal['tqdm', 'simple'] | None = None, callback_interval: float = 20, n: int | None = None, reduce_ys_fn: Callable | Literal['mean'] | None = None) -> Callable
¤
Wrap a scan body so its per-iteration metrics can be logged and shrunk.
Arguments:
func: function to wrap
callback_fn: function to call on the ys of the scan
callback_interval: how often to call the callback in iterations. Can be a fraction of n.
n: total number of iterations, required if callback_interval 0.0 < 1.0.
reduce_ys_fn: optional function to reduce the metrics after the callback is called,
but before they are returned to the scan. May be used to aggregate metrics to reduce
memory consumption. Set this to mean (default) to average LogWrapper returned
returned returns across the current iteration of data.
jaxnasium.algorithms.core.mean_episode_returns(metrics: dict[str, Any]) -> PyTree[Float[Array, '']]
¤
mean episode return across all vectorized environments in a single batch of metrics produces a single scalar (per agent) per scan iteration. Requires LogWrapper on the environment.