deep-learning-core 0.0.26__tar.gz → 0.0.28__tar.gz

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deep-learning-core
-Version: 0.0.26
+Version: 0.0.28
 Summary: Reusable deep learning framework core.
 Project-URL: Homepage, https://github.com/Blazkowiz47/dl-core
 Project-URL: Documentation, https://github.com/Blazkowiz47/dl-core/tree/master/readme
@@ -76,14 +76,39 @@ across many experiment repositories. It is intended to be the public base
 package, while optional integrations such as Azure are layered on through
 extras and companion extension packages.
 
-Current release: `deep-learning-core==0.0.26`.
+Current public release: `deep-learning-core==0.0.28`.
+Current development version: `0.0.28`.
 
 Compatible companion package floors:
 
 - `deep-learning-azure>=0.0.18,<0.1`
 - `deep-learning-mlflow>=0.0.11,<0.1`
+- `deep-learning-robotics>=0.0.2,<0.1` with
+  `deep-learning-core>=0.0.28,<0.1`
 - `deep-learning-wandb>=0.0.12,<0.1`
 
+## What's New in 0.0.28?
+
+- RL collection, replay insertion, episode persistence, environment creation,
+  and RL component scaffolding now keep one-off logic inline for a more direct
+  implementation
+- public trainer, environment, episode-manager, and scaffold behavior remains
+  unchanged
+
+## What's New in 0.0.27?
+
+- episode managers provide generic RL summaries and selective, complete
+  trajectory capture alongside the existing metric-manager system
+- scalar and same-step vector environments share one collector contract with
+  preserved terminal observations and per-lane episode identity
+- Q-learning, DQN, PPO, and SAC now consume vector collection natively; neural
+  policies perform batched inference and replay/rollout storage preserves the
+  correct algorithm-specific scheduling and boundary semantics
+- replay insertion and PPO rollout/GAE computation operate on real batches,
+  while scalar custom trainers remain compatible through the original hooks
+- extension packages can import and register environments without depending on
+  dl-core import order
+
 ## What's New in 0.0.26?
 
 - Gymnasium-compatible environments can now be registered, discovered, and
@@ -167,6 +192,9 @@ used by the validated Azure packaging stack. The MLflow extra pulls in
 - `deep-learning-core[wandb]`: adds the public
   [`dl-wandb`](https://github.com/Blazkowiz47/dl-wandb)
   package for Weights & Biases integration
+- [`deep-learning-robotics`](https://github.com/Blazkowiz47/dl-robotics):
+  adds fast scalar and vector 2D MAPF environments, metrics, and episode media
+  as a separately installed companion package
 
 The extension packages stay separate so the base package remains reusable and
 vendor-neutral.
@@ -177,6 +205,7 @@ integration without using extras:
 ```bash
 pip install deep-learning-azure
 pip install deep-learning-mlflow
+pip install deep-learning-robotics
 pip install deep-learning-wandb
 ```
 
@@ -367,6 +396,7 @@ scaffold a `wandb` callback block, W&B tracking defaults, and `.env.example`.
 
 - [`dl-azure`](https://github.com/Blazkowiz47/dl-azure)
 - [`dl-mlflow`](https://github.com/Blazkowiz47/dl-mlflow)
+- [`dl-robotics`](https://github.com/Blazkowiz47/dl-robotics)
 - [`dl-wandb`](https://github.com/Blazkowiz47/dl-wandb)
 
 ## Scaffold Commands
@@ -386,6 +416,7 @@ uv run dl-core add trainer MyTrainer
 uv run dl-core add trainer MyPolicy --base rltrainer
 uv run dl-core add callback MyMetrics
 uv run dl-core add metric_manager MyManager
+uv run dl-core add episode_manager MyEpisodeManager
 uv run dl-core add sampler MySampler
 uv run dl-core add optimizer MyOptimizer
 uv run dl-core add scheduler MyScheduler

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/README.md

@@ -7,14 +7,39 @@ across many experiment repositories. It is intended to be the public base
 package, while optional integrations such as Azure are layered on through
 extras and companion extension packages.
 
-Current release: `deep-learning-core==0.0.26`.
+Current public release: `deep-learning-core==0.0.28`.
+Current development version: `0.0.28`.
 
 Compatible companion package floors:
 
 - `deep-learning-azure>=0.0.18,<0.1`
 - `deep-learning-mlflow>=0.0.11,<0.1`
+- `deep-learning-robotics>=0.0.2,<0.1` with
+  `deep-learning-core>=0.0.28,<0.1`
 - `deep-learning-wandb>=0.0.12,<0.1`
 
+## What's New in 0.0.28?
+
+- RL collection, replay insertion, episode persistence, environment creation,
+  and RL component scaffolding now keep one-off logic inline for a more direct
+  implementation
+- public trainer, environment, episode-manager, and scaffold behavior remains
+  unchanged
+
+## What's New in 0.0.27?
+
+- episode managers provide generic RL summaries and selective, complete
+  trajectory capture alongside the existing metric-manager system
+- scalar and same-step vector environments share one collector contract with
+  preserved terminal observations and per-lane episode identity
+- Q-learning, DQN, PPO, and SAC now consume vector collection natively; neural
+  policies perform batched inference and replay/rollout storage preserves the
+  correct algorithm-specific scheduling and boundary semantics
+- replay insertion and PPO rollout/GAE computation operate on real batches,
+  while scalar custom trainers remain compatible through the original hooks
+- extension packages can import and register environments without depending on
+  dl-core import order
+
 ## What's New in 0.0.26?
 
 - Gymnasium-compatible environments can now be registered, discovered, and
@@ -98,6 +123,9 @@ used by the validated Azure packaging stack. The MLflow extra pulls in
 - `deep-learning-core[wandb]`: adds the public
   [`dl-wandb`](https://github.com/Blazkowiz47/dl-wandb)
   package for Weights & Biases integration
+- [`deep-learning-robotics`](https://github.com/Blazkowiz47/dl-robotics):
+  adds fast scalar and vector 2D MAPF environments, metrics, and episode media
+  as a separately installed companion package
 
 The extension packages stay separate so the base package remains reusable and
 vendor-neutral.
@@ -108,6 +136,7 @@ integration without using extras:
 ```bash
 pip install deep-learning-azure
 pip install deep-learning-mlflow
+pip install deep-learning-robotics
 pip install deep-learning-wandb
 ```
 
@@ -298,6 +327,7 @@ scaffold a `wandb` callback block, W&B tracking defaults, and `.env.example`.
 
 - [`dl-azure`](https://github.com/Blazkowiz47/dl-azure)
 - [`dl-mlflow`](https://github.com/Blazkowiz47/dl-mlflow)
+- [`dl-robotics`](https://github.com/Blazkowiz47/dl-robotics)
 - [`dl-wandb`](https://github.com/Blazkowiz47/dl-wandb)
 
 ## Scaffold Commands
@@ -317,6 +347,7 @@ uv run dl-core add trainer MyTrainer
 uv run dl-core add trainer MyPolicy --base rltrainer
 uv run dl-core add callback MyMetrics
 uv run dl-core add metric_manager MyManager
+uv run dl-core add episode_manager MyEpisodeManager
 uv run dl-core add sampler MySampler
 uv run dl-core add optimizer MyOptimizer
 uv run dl-core add scheduler MyScheduler

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "deep-learning-core"
-version = "0.0.26"
+version = "0.0.28"
 description = "Reusable deep learning framework core."
 readme = "README.md"
 requires-python = ">=3.10"

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/readme/README.md

@@ -4,7 +4,14 @@ This documentation is split into quick references, workflow guides, and
 technical notes. The goal is the same as in the original framework repo, but
 focused on the extracted package and the experiment-repo workflow around it.
 
-Current public release: `deep-learning-core==0.0.26`.
+Current public release: `deep-learning-core==0.0.28`.
+Current development version: `0.0.28`.
+
+## What's New in 0.0.28?
+
+- RL collection, replay insertion, episode persistence, environment creation,
+  and component scaffolding use fewer one-off implementation helpers
+- public APIs and runtime behavior remain unchanged
 
 ## Companion Packages
 
@@ -12,6 +19,8 @@ Current public release: `deep-learning-core==0.0.26`.
   Azure dataset foundations
 - [`dl-mlflow`](https://github.com/Blazkowiz47/dl-mlflow): local MLflow
   integration
+- [`dl-robotics`](https://github.com/Blazkowiz47/dl-robotics): fast 2D MAPF
+  environments, episode metrics, and GIF/MP4 artifacts
 - [`dl-wandb`](https://github.com/Blazkowiz47/dl-wandb): Weights & Biases
   integration
 

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/readme/technical/6_reinforcement_learning.md

@@ -105,11 +105,96 @@ simulator state is not serialized; a resumed run begins at a new episode
 boundary. A checkpoint is rejected when its trainer implementation or component
 names do not match the configured trainer.
 
+Training collection can use a Gymnasium vector environment. The built-in
+adapter creates one with same-step autoreset so completed transitions retain
+their real `final_obs` while collection immediately receives the next episode's
+initial observation:
+
+```yaml
+environment:
+  name: gymnasium_vector
+  id: CartPole-v1
+  num_envs: 8
+  vectorization_mode: sync
+
+evaluation_environment:
+  name: gymnasium
+  id: CartPole-v1
+```
+
+`global_step` counts transitions, while `collector_step` counts vector
+environment calls. `select_actions` and `process_transition_batch` expose one
+whole collector step to vector-aware algorithms. Existing custom trainers that
+only implement `select_action` and `process_transition` remain compatible
+through stable environment-index dispatch. Evaluation deliberately uses one
+scalar environment so videos and deterministic episode artifacts have
+unambiguous identity. A transition budget can overshoot by at most
+`num_envs - 1` because a vector step is atomic. An episode budget has the same
+maximum overshoot when several lanes complete in one vector step.
+
+Tabular Q-learning consumes vector steps in stable lane order. DQN performs at
+most one batched action-selection inference per vector step, inserts the complete
+transition batch into replay, and applies every update or target-sync boundary
+crossed by that atomic step in schedule order. PPO performs batched policy and
+next-value inference, stores independent `[time, environment]` streams, and
+flattens them only after per-lane GAE is complete. `rollout_steps` counts
+synchronized collector calls, so an update uses
+`rollout_steps * num_envs` samples. SAC samples continuous actions in one actor
+call, inserts vector transitions directly into replay, and runs the gradient
+work for every update-frequency boundary crossed by the atomic vector step.
+Each crossed boundary emits its own update log, and warm-up is applied to the
+first `learning_starts` transitions even when that boundary falls inside a
+vector step.
+
+Replay storage accepts transition batches directly and preserves configured
+observation/action dtypes. PPO rollout storage is preallocated as
+`[time, environment, ...]`; generalized advantages are propagated only within
+the same environment stream before the rollout is flattened for minibatches.
+
 RL callbacks can implement `on_episode_start`, `on_episode_end`,
 `on_update_end`, and `on_evaluation_end`. The existing run-level
 `on_training_start`, `on_training_end`, and `on_training_finalized` hooks remain
 shared with epoch training.
 
+## Episode Managers
+
+Episode managers are the reinforcement-learning counterpart to the metric
+managers used by `EpochTrainer`. They accumulate environment transitions,
+compute episode summaries, and optionally persist complete trajectories while
+callbacks remain responsible for external logging and side effects.
+
+The built-in `standard` manager always computes return, length, reward
+statistics, termination, truncation, and success when the environment exposes
+`is_success`. Complete trajectories preserve the initial observation followed
+by one observation per transition, so an episode of length `T` contains `T + 1`
+observations and `T` actions, rewards, termination flags, and truncation flags.
+
+```yaml
+episode_managers:
+  standard:
+    capture_phases: [evaluation]
+    capture_every_n_episodes: 10
+    max_captured_episodes: 20
+    info_keys: [is_success, collision]
+    capture_action_info: false
+```
+
+Captured array-valued trajectories are portable compressed NumPy archives under
+`final/episodes/<phase>/`; the episode index and scalar summary streams remain
+JSONL. Array dtypes are preserved and arbitrary Python objects are not pickled.
+Environments with array-valued observations and actions, including nested
+dictionaries and tuples, can persist their environment-boundary trajectory.
+Exact hidden simulator state is available only when a concrete environment
+provides a separate state snapshot capability.
+
+Episode managers use the normal component workflow:
+
+```bash
+dl-core list episode_manager
+dl-core describe episode_manager standard
+dl-core add episode_manager PathAnalysis
+```
+
 The built-in local metric tracker and the MLflow and W&B companion callbacks
 record episode, algorithm-update, and evaluation metrics as well as supervised
 epoch metrics. Training-episode series stay separate from the aggregate metrics
@@ -263,10 +348,18 @@ or continuous `mean` and `log_std` tensors shaped `[batch, action_dimensions]`.
 Policy modules should avoid dropout and other stochastic training-mode layers,
 because PPO must reproduce the behavior-policy probability for stored actions.
 
-The initial PPO implementation collects a single environment stream. This keeps
-episode callbacks and deterministic checkpoint continuation identical to the
-other trainers. A later vector-environment collector can feed the same policy
-and rollout contracts without changing the public trainer configuration.
+PPO accepts scalar or vector training environments. Vector collection evaluates
+all lanes in one policy call, computes GAE independently along each environment
+stream, and flattens the time/environment axes only for minibatch optimization.
+`rollout_steps` counts synchronized collector steps, so a full vector rollout
+contains `rollout_steps * num_envs` training samples. Evaluation remains scalar
+to preserve deterministic episode-level reporting.
+
+Partial rollouts are stored in checkpoints. Because environment state is not
+part of an RL checkpoint, resuming marks the final stored transition in every
+unfinished lane as truncated. Its one-step value target is retained, while GAE
+cannot propagate from the newly reset environment into the earlier rollout
+fragment.
 
 ## Soft Actor-Critic
 

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/src/dl_core/__init__.py

@@ -22,6 +22,7 @@ _BUILTIN_COMPONENT_MODULES = (
     "datasets",
     "environments",
     "executors",
+    "episode_managers",
     "metrics_sources",
     "metric_managers",
     "metrics",
@@ -33,7 +34,7 @@ _BUILTIN_COMPONENT_MODULES = (
     "trainers",
 )
 
-__version__ = "0.0.26"
+__version__ = "0.0.28"
 
 _LOCAL_MODULES: dict[str, ModuleType] = {}
 _LOCAL_REGISTRATIONS: list[tuple[Any, str, type[Any]]] = []

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/src/dl_core/component_describer.py

@@ -17,6 +17,7 @@ from dl_core.core import (
     CALLBACK_REGISTRY,
     CRITERION_REGISTRY,
     DATASET_REGISTRY,
+    EPISODE_MANAGER_REGISTRY,
     ENVIRONMENT_REGISTRY,
     EXECUTOR_REGISTRY,
     FACE_DETECTOR_REGISTRY,
@@ -39,6 +40,7 @@ _COMPONENT_REGISTRIES: dict[str, ComponentRegistry] = {
     "callback": CALLBACK_REGISTRY,
     "criterion": CRITERION_REGISTRY,
     "dataset": DATASET_REGISTRY,
+    "episode_manager": EPISODE_MANAGER_REGISTRY,
     "environment": ENVIRONMENT_REGISTRY,
     "executor": EXECUTOR_REGISTRY,
     "face_detector": FACE_DETECTOR_REGISTRY,
@@ -67,6 +69,10 @@ _DESCRIBE_TYPE_ALIASES = {
     "criterions": "criterion",
     "dataset": "dataset",
     "datasets": "dataset",
+    "episode_manager": "episode_manager",
+    "episode_managers": "episode_manager",
+    "episodemanager": "episode_manager",
+    "episodemanagers": "episode_manager",
     "environment": "environment",
     "environments": "environment",
     "executor": "executor",
@@ -478,6 +484,8 @@ def _build_config_example(
         payload = {"criterions": {component_name: field_values}}
     elif component_type == "metric_manager":
         payload = {"metric_managers": {component_name: field_values}}
+    elif component_type == "episode_manager":
+        payload = {"episode_managers": {component_name: field_values}}
     elif component_type == "callback":
         payload = {"callbacks": {component_name: field_values}}
     elif component_type == "optimizer":

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/src/dl_core/component_scaffold.py

@@ -15,6 +15,7 @@ from dl_core.core import (
     AUGMENTATION_REGISTRY,
     CALLBACK_REGISTRY,
     CRITERION_REGISTRY,
+    EPISODE_MANAGER_REGISTRY,
     EXECUTOR_REGISTRY,
     METRIC_MANAGER_REGISTRY,
     METRIC_REGISTRY,
@@ -97,6 +98,13 @@ _COMPONENT_SPECS = {
         init_docstring="Local executor extensions.",
         register_name="register_executor",
     ),
+    "episode_manager": ComponentSpec(
+        canonical_name="episode_manager",
+        package_dir="episode_managers",
+        class_suffix="EpisodeManager",
+        init_docstring="Local reinforcement-learning episode managers.",
+        register_name="register_episode_manager",
+    ),
     "metric": ComponentSpec(
         canonical_name="metric",
         package_dir="metrics",
@@ -153,6 +161,7 @@ _COMPONENT_BASE_REGISTRIES: dict[str, ComponentRegistry] = {
     "callback": CALLBACK_REGISTRY,
     "criterion": CRITERION_REGISTRY,
     "executor": EXECUTOR_REGISTRY,
+    "episode_manager": EPISODE_MANAGER_REGISTRY,
     "metric": METRIC_REGISTRY,
     "metric_manager": METRIC_MANAGER_REGISTRY,
     "model": MODEL_REGISTRY,
@@ -187,6 +196,14 @@ _DEFAULT_COMPONENT_BASE_SPECS = {
         base_class="BaseExecutor",
         class_docstring="Local executor scaffold based on BaseExecutor.",
     ),
+    "episode_manager": ComponentBaseSpec(
+        canonical_name="episode_manager",
+        import_path="dl_core.core",
+        base_class="BaseEpisodeManager",
+        class_docstring=(
+            "Local episode manager scaffold based on BaseEpisodeManager."
+        ),
+    ),
     "metric": ComponentBaseSpec(
         canonical_name="metric",
         import_path="dl_core.core",
@@ -244,6 +261,10 @@ _COMPONENT_TYPE_ALIASES = {
     "datasets": "dataset",
     "executor": "executor",
     "executors": "executor",
+    "episode_manager": "episode_manager",
+    "episode_managers": "episode_manager",
+    "episodemanager": "episode_manager",
+    "episodemanagers": "episode_manager",
     "metric": "metric",
     "metrics": "metric",
     "metricmanager": "metric_manager",
@@ -608,6 +629,33 @@ def _render_component(
                 base_class=metric_manager_base.base_class,
                 class_docstring=metric_manager_base.class_docstring,
             )
+    if spec.canonical_name == "episode_manager":
+        episode_manager_base = _component_base_required(component_base)
+        if _is_default_component_base(spec.canonical_name, episode_manager_base):
+            return f'''"""Local episode manager scaffold."""
+
+from __future__ import annotations
+
+from dl_core.core import EpisodeRecord, EpisodeResult, register_episode_manager
+from {episode_manager_base.import_path} import {episode_manager_base.base_class}
+
+
+@register_episode_manager({registry_literal})
+class {class_name}({episode_manager_base.base_class}):
+    """{episode_manager_base.class_docstring}"""
+
+    def summarize_episode(
+        self,
+        record: EpisodeRecord,
+        result: EpisodeResult,
+        **statistics: float | int,
+    ) -> dict[str, float]:
+        """Compute scalar metrics for one completed episode."""
+        return {{
+            "episode/return": float(statistics["episode_return"]),
+            "episode/length": float(statistics["length"]),
+        }}
+'''
     if spec.canonical_name == "model":
         model_base = _component_base_required(component_base)
         if _is_default_component_base(spec.canonical_name, model_base):
@@ -1036,33 +1084,7 @@ def _trainer_component(
             class_docstring=class_docstring,
         )
     if base_class == "RLTrainer":
-        return _rl_trainer_component(
-            registry_literal=registry_literal,
-            class_name=class_name,
-            import_path=import_path,
-            base_class=base_class,
-            class_docstring=class_docstring,
-        )
-    return _epoch_trainer_component(
-        registry_literal=registry_literal,
-        class_name=class_name,
-        import_path=import_path,
-        base_class=base_class,
-        class_docstring=class_docstring,
-    )
-
-
-def _rl_trainer_component(
-    *,
-    registry_literal: str,
-    class_name: str,
-    import_path: str,
-    base_class: str,
-    class_docstring: str,
-) -> str:
-    """Render an episode-oriented reinforcement-learning trainer scaffold."""
-
-    return f'''"""Local reinforcement-learning trainer scaffold."""
+        return f'''"""Local reinforcement-learning trainer scaffold."""
 
 from __future__ import annotations
 
@@ -1105,6 +1127,13 @@ class {class_name}({base_class}):
         if state:
             raise ValueError("TODO: validate and restore algorithm state.")
 '''
+    return _epoch_trainer_component(
+        registry_literal=registry_literal,
+        class_name=class_name,
+        import_path=import_path,
+        base_class=base_class,
+        class_docstring=class_docstring,
+    )
 
 
 def _epoch_trainer_component(

{deep_learning_core-0.0.26 → deep_learning_core-0.0.28}/src/dl_core/core/__init__.py

@@ -16,6 +16,7 @@ from .base_executor import BaseExecutor
 from .base_metric import BaseMetric
 from .base_metrics_source import BaseMetricsSource
 from .base_metric_manager import BaseMetricManager
+from .base_episode_manager import BaseEpisodeManager
 from .base_biometric_model import BaseBiometricModel
 from .base_model import BaseModel
 from .base_sampler import BaseSampler
@@ -26,12 +27,23 @@ from .adaptive_computation_trainer import (
     CarryState,
 )
 from .base_callback import Callback
+from .batched_environment import BatchedEnvironment
 from .base_trainer import BaseTrainer
 from .epoch_trainer import EpochTrainer
 from .sequence_trainer import SequenceStepOutput, SequenceTrainer
 from .base_transform import BaseTransform
 from .rl_trainer import RLTrainer
-from .rl_types import ActionOutput, Environment, EpisodeResult, Transition
+from .rl_types import (
+    ActionOutput,
+    BatchActionOutput,
+    Environment,
+    EpisodeContext,
+    EpisodeRecord,
+    EpisodeResult,
+    Transition,
+    TransitionBatch,
+    VectorEnvironment,
+)
 from .replay_buffer import ReplayBatch, ReplayBuffer
 from .rollout_buffer import RolloutBatch, RolloutBuffer
 from .config_metadata import ConfigFieldSpec, config_field
@@ -44,6 +56,7 @@ from .registry import (
     CRITERION_REGISTRY,
     DATASET_REGISTRY,
     ENVIRONMENT_REGISTRY,
+    EPISODE_MANAGER_REGISTRY,
     EXECUTOR_REGISTRY,
     FACE_DETECTOR_REGISTRY,
     METRIC_MANAGER_REGISTRY,
@@ -65,6 +78,7 @@ from .registry import (
     register_criterion,
     register_dataset,
     register_environment,
+    register_episode_manager,
     register_executor,
     register_face_detector,
     register_metric,
@@ -94,8 +108,10 @@ __all__ = [
     "BaseTracker",
     "BaseMetricsSource",
     "BaseMetricManager",
+    "BaseEpisodeManager",
     "BaseSampler",
     "Callback",
+    "BatchedEnvironment",
     "BaseTrainer",
     "EpochTrainer",
     "SequenceTrainer",
@@ -109,8 +125,13 @@ __all__ = [
     "LandmarkDetection",
     "ConfigFieldSpec",
     "Environment",
+    "VectorEnvironment",
     "ActionOutput",
+    "BatchActionOutput",
     "Transition",
+    "TransitionBatch",
+    "EpisodeContext",
+    "EpisodeRecord",
     "EpisodeResult",
     "RLTrainer",
     "ReplayBatch",
@@ -123,6 +144,7 @@ __all__ = [
     "TRAINER_REGISTRY",
     "DATASET_REGISTRY",
     "ENVIRONMENT_REGISTRY",
+    "EPISODE_MANAGER_REGISTRY",
     "CRITERION_REGISTRY",
     "METRIC_REGISTRY",
     "METRIC_MANAGER_REGISTRY",
@@ -141,6 +163,7 @@ __all__ = [
     "register_trainer",
     "register_dataset",
     "register_environment",
+    "register_episode_manager",
     "register_face_detector",
     "register_criterion",
     "register_metric",