tensor-optix 2.0.0__tar.gz → 2.0.2__tar.gz

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensor-optix
-Version: 2.0.0
+Version: 2.0.2
 Summary: Autonomous training loop for any sequential learning model — PPO, DQN, SAC, TD3, Rainbow DQN, Recurrent PPO for TensorFlow, PyTorch, and JAX/Flax; distributed async actor-learner (IMPALA + V-trace)
 Author: sup3rus3r
 License-Expression: MIT

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensor-optix"
-version = "2.0.0"
+version = "2.0.2"
 description = "Autonomous training loop for any sequential learning model — PPO, DQN, SAC, TD3, Rainbow DQN, Recurrent PPO for TensorFlow, PyTorch, and JAX/Flax; distributed async actor-learner (IMPALA + V-trace)"
 readme = "README.md"
 license = "MIT"

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/adapters/pytorch/torch_agent.py

@@ -2,6 +2,7 @@ import os
 import numpy as np
 
 from tensor_optix.core.base_agent import BaseAgent
+from tensor_optix.core.resume_state import save_torch_resume_state, load_torch_resume_state
 from tensor_optix.core.types import EpisodeData, HyperparamSet
 
 
@@ -147,6 +148,21 @@ class TorchAgent(BaseAgent):
         state = torch.load(os.path.join(path, "model.pt"), map_location=self._device)
         self.model.load_state_dict(state)
 
+    @property
+    def resumable_components(self) -> frozenset:
+        return frozenset({"model", "optimizer", "rng"})
+
+    def save_full_state(self, path: str) -> None:
+        self.save_weights(path)
+        save_torch_resume_state(path, optimizer=self.optimizer)
+
+    def load_full_state(self, path: str) -> None:
+        self.load_weights(path)
+        load_torch_resume_state(path, optimizer=self.optimizer)
+
+    def get_weights(self) -> dict:
+        return {k: v.detach().clone() for k, v in self.model.state_dict().items()}
+
     def teardown(self) -> None:
         """Move model to CPU and free CUDA memory."""
         import torch

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/flax_ppo.py

@@ -38,7 +38,7 @@ import numpy as np
 
 from tensor_optix.core.base_agent import BaseAgent
 from tensor_optix.core.types import EpisodeData, HyperparamSet
-from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches
+from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches, resolve_gae_bootstrap
 
 
 # ---------------------------------------------------------------------------
@@ -195,15 +195,14 @@ class FlaxPPOAgent(BaseAgent):
         dones      = episode_data.dones
 
         # Bootstrap value when trajectory ends mid-episode.
-        last_value = 0.0
-        if not dones[-1] and episode_data.final_obs is not None:
-            obs_final  = jnp.array(
-                np.atleast_2d(episode_data.final_obs), dtype=jnp.float32
-            )
-            last_value = float(self._model.critic(obs_final)[0])
+        def _critic_value(obs):
+            obs_final = jnp.array(np.atleast_2d(obs), dtype=jnp.float32)
+            return float(self._model.critic(obs_final)[0])
+
+        last_value, n_envs = resolve_gae_bootstrap(episode_data, _critic_value)
 
         advantages, returns = compute_gae(
-            rewards, val_arr, dones, gamma, gae_lambda, last_value
+            rewards, val_arr, dones, gamma, gae_lambda, last_value, n_envs=n_envs
         )
         if T > 1:
             advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/tf_ppo.py

@@ -3,7 +3,7 @@ import numpy as np
 import tensorflow as tf
 from tensor_optix.core.base_agent import BaseAgent
 from tensor_optix.core.types import EpisodeData, HyperparamSet
-from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches
+from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches, resolve_gae_bootstrap
 
 
 class TFPPOAgent(BaseAgent):
@@ -149,14 +149,18 @@ class TFPPOAgent(BaseAgent):
                     self._reward_normalizer.reset()
             rewards = list(self._reward_normalizer.normalize(np.array(rewards, dtype=np.float32)))
 
-        # Bootstrap: if the window ended mid-episode, V(s_T) from the critic
+        # Bootstrap: if a window ended mid-episode, V(s_T) from the critic
         # corrects the TD target at the last step. Zero when terminal.
-        last_value = 0.0
-        if not dones[-1] and episode_data.final_obs is not None:
-            final_obs_t = tf.cast(np.atleast_2d(episode_data.final_obs), tf.float32)
-            last_value = float(tf.squeeze(self._critic(final_obs_t, training=False)).numpy())
+        # With VectorBatchPipeline (final_obs is a list), this resolves one
+        # bootstrap value per env, and n_envs tells compute_gae how the
+        # rewards/values array is interleaved across envs.
+        def _critic_value(obs):
+            obs_t = tf.cast(np.atleast_2d(obs), tf.float32)
+            return float(tf.squeeze(self._critic(obs_t, training=False)).numpy())
 
-        advantages, returns = compute_gae(rewards, values_arr, dones, gamma, gae_lambda, last_value)
+        last_value, n_envs = resolve_gae_bootstrap(episode_data, _critic_value)
+
+        advantages, returns = compute_gae(rewards, values_arr, dones, gamma, gae_lambda, last_value, n_envs=n_envs)
 
         # Normalize advantages across the rollout
         if T > 1:

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/tf_ppo_continuous.py

@@ -4,7 +4,7 @@ import tensorflow as tf
 
 from tensor_optix.core.base_agent import BaseAgent
 from tensor_optix.core.types import EpisodeData, HyperparamSet
-from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches
+from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches, resolve_gae_bootstrap
 
 LOG_STD_MIN = -5
 LOG_STD_MAX = 2
@@ -193,12 +193,13 @@ class TFGaussianPPOAgent(BaseAgent):
                     self._reward_normalizer.reset()
             rewards = list(self._reward_normalizer.normalize(np.array(rewards, dtype=np.float32)))
 
-        last_value = 0.0
-        if not dones[-1] and episode_data.final_obs is not None:
-            final_obs_t = tf.cast(np.atleast_2d(episode_data.final_obs), tf.float32)
-            last_value = float(tf.squeeze(self._critic(final_obs_t, training=False)).numpy())
+        def _critic_value(obs):
+            obs_t = tf.cast(np.atleast_2d(obs), tf.float32)
+            return float(tf.squeeze(self._critic(obs_t, training=False)).numpy())
 
-        advantages, returns = compute_gae(rewards, val_arr, dones, gamma, gae_lambda, last_value)
+        last_value, n_envs = resolve_gae_bootstrap(episode_data, _critic_value)
+
+        advantages, returns = compute_gae(rewards, val_arr, dones, gamma, gae_lambda, last_value, n_envs=n_envs)
         if T > 1:
             advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
 

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/tf_sac.py

@@ -151,7 +151,7 @@ class TFSACAgent(BaseAgent):
         gamma             = float(hp.get("gamma",             0.99))
         tau               = float(hp.get("tau",               0.005))
         batch_size        = int(hp.get("batch_size",          256))
-        updates_per_step  = int(hp.get("updates_per_step",    1))
+        updates_per_step  = float(hp.get("updates_per_step",  1))
 
         obs_arr  = np.array(episode_data.observations, dtype=np.float32)
         act_arr  = np.array(episode_data.actions,      dtype=np.float32)
@@ -177,7 +177,10 @@ class TFSACAgent(BaseAgent):
                 "buffer_size": len(self._buffer),
             }
 
-        n_updates = T * updates_per_step
+        # Truncate the product, not updates_per_step itself -- truncating
+        # a fractional updates_per_step (e.g. 0.5) to int first would zero
+        # it out regardless of T.
+        n_updates = int(T * updates_per_step)
         total_actor_loss  = 0.0
         total_critic_loss = 0.0
 
@@ -187,8 +190,8 @@ class TFSACAgent(BaseAgent):
             total_critic_loss += cl
 
         return {
-            "actor_loss":  total_actor_loss  / n_updates,
-            "critic_loss": total_critic_loss / n_updates,
+            "actor_loss":  total_actor_loss  / max(n_updates, 1),
+            "critic_loss": total_critic_loss / max(n_updates, 1),
             "alpha":       float(tf.exp(self._log_alpha).numpy()),
             "buffer_size": len(self._buffer),
         }

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/tf_td3.py

@@ -131,7 +131,7 @@ class TFTDDAgent(BaseAgent):
         gamma            = float(hp.get("gamma",            0.99))
         tau              = float(hp.get("tau",              0.005))
         batch_size       = int(hp.get("batch_size",         256))
-        updates_per_step = int(hp.get("updates_per_step",   1))
+        updates_per_step = float(hp.get("updates_per_step", 1))
         policy_delay     = int(hp.get("policy_delay",       2))
 
         obs_arr  = np.array(episode_data.observations, dtype=np.float32)
@@ -155,7 +155,10 @@ class TFTDDAgent(BaseAgent):
                 "policy_update": 0,
             }
 
-        n_updates = T * updates_per_step
+        # Truncate the product, not updates_per_step itself -- truncating
+        # a fractional updates_per_step (e.g. 0.5) to int first would zero
+        # it out regardless of T.
+        n_updates = int(T * updates_per_step)
         total_al = 0.0
         total_cl = 0.0
         n_actor_updates = 0
@@ -171,7 +174,7 @@ class TFTDDAgent(BaseAgent):
 
         return {
             "actor_loss":    total_al / max(n_actor_updates, 1),
-            "critic_loss":   total_cl / n_updates,
+            "critic_loss":   total_cl / max(n_updates, 1),
             "buffer_size":   len(self._buffer),
             "policy_update": int(n_actor_updates > 0),
         }

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/torch_ppo.py

@@ -2,8 +2,9 @@ import os
 import numpy as np
 
 from tensor_optix.core.base_agent import BaseAgent
+from tensor_optix.core.resume_state import save_torch_resume_state, load_torch_resume_state
 from tensor_optix.core.types import EpisodeData, HyperparamSet
-from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches
+from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches, resolve_gae_bootstrap
 
 
 class TorchPPOAgent(BaseAgent):
@@ -126,17 +127,19 @@ class TorchPPOAgent(BaseAgent):
                     self._reward_normalizer.reset()
             rewards = list(self._reward_normalizer.normalize(np.array(rewards, dtype=np.float32)))
 
-        # Bootstrap: if the window ended mid-episode, V(s_T) from the critic
+        # Bootstrap: if a window ended mid-episode, V(s_T) from the critic
         # corrects the TD target at the last step. Zero when terminal.
-        last_value = 0.0
-        if not dones[-1] and episode_data.final_obs is not None:
+        # With VectorBatchPipeline (final_obs is a list), this resolves one
+        # bootstrap value per env, and n_envs tells compute_gae how the
+        # rewards/values array is interleaved across envs.
+        def _critic_value(obs):
             with torch.no_grad():
-                final_obs_t = torch.as_tensor(
-                    np.atleast_2d(episode_data.final_obs), dtype=torch.float32
-                ).to(self._device)
-                last_value = float(self._critic(final_obs_t).squeeze().item())
+                obs_t = torch.as_tensor(np.atleast_2d(obs), dtype=torch.float32).to(self._device)
+                return float(self._critic(obs_t).squeeze().item())
 
-        advantages, returns = compute_gae(rewards, val_arr, dones, gamma, gae_lambda, last_value)
+        last_value, n_envs = resolve_gae_bootstrap(episode_data, _critic_value)
+
+        advantages, returns = compute_gae(rewards, val_arr, dones, gamma, gae_lambda, last_value, n_envs=n_envs)
         if T > 1:
             advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
 
@@ -248,6 +251,24 @@ class TorchPPOAgent(BaseAgent):
             torch.load(os.path.join(path, "critic.pt"), map_location=self._device)
         )
 
+    @property
+    def resumable_components(self) -> frozenset:
+        return frozenset({"model", "optimizer", "rng"})
+
+    def save_full_state(self, path: str) -> None:
+        self.save_weights(path)
+        save_torch_resume_state(path, optimizer=self._optimizer)
+
+    def load_full_state(self, path: str) -> None:
+        self.load_weights(path)
+        load_torch_resume_state(path, optimizer=self._optimizer)
+
+    def get_weights(self) -> dict:
+        return {
+            "actor":  {k: v.detach().clone() for k, v in self._actor.state_dict().items()},
+            "critic": {k: v.detach().clone() for k, v in self._critic.state_dict().items()},
+        }
+
     def average_weights(self, paths: list) -> None:
         import torch
         modules = ("actor", "critic")

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/torch_ppo_continuous.py

@@ -3,7 +3,7 @@ import numpy as np
 
 from tensor_optix.core.base_agent import BaseAgent
 from tensor_optix.core.types import EpisodeData, HyperparamSet
-from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches
+from tensor_optix.core.trajectory_buffer import compute_gae, make_minibatches, resolve_gae_bootstrap
 
 LOG_STD_MIN = -5
 LOG_STD_MAX = 2
@@ -206,15 +206,14 @@ class TorchGaussianPPOAgent(BaseAgent):
                     self._reward_normalizer.reset()
             rewards = list(self._reward_normalizer.normalize(np.array(rewards, dtype=np.float32)))
 
-        last_value = 0.0
-        if not dones[-1] and episode_data.final_obs is not None:
+        def _critic_value(obs):
             with torch.no_grad():
-                final_obs_t = torch.as_tensor(
-                    np.atleast_2d(episode_data.final_obs), dtype=torch.float32
-                ).to(self._device)
-                last_value = float(self._critic(final_obs_t).squeeze().item())
+                obs_t = torch.as_tensor(np.atleast_2d(obs), dtype=torch.float32).to(self._device)
+                return float(self._critic(obs_t).squeeze().item())
 
-        advantages, returns = compute_gae(rewards, val_arr, dones, gamma, gae_lambda, last_value)
+        last_value, n_envs = resolve_gae_bootstrap(episode_data, _critic_value)
+
+        advantages, returns = compute_gae(rewards, val_arr, dones, gamma, gae_lambda, last_value, n_envs=n_envs)
         if T > 1:
             advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
 

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/torch_sac.py

@@ -4,6 +4,7 @@ import copy
 import numpy as np
 
 from tensor_optix.core.base_agent import BaseAgent
+from tensor_optix.core.resume_state import save_torch_resume_state, load_torch_resume_state
 from tensor_optix.core.types import EpisodeData, HyperparamSet
 from tensor_optix.core.replay_buffer import PrioritizedReplayBuffer
 
@@ -145,7 +146,7 @@ class TorchSACAgent(BaseAgent):
         gamma            = float(hp.get("gamma",            0.99))
         tau              = float(hp.get("tau",              0.005))
         batch_size       = int(hp.get("batch_size",         256))
-        updates_per_step = int(hp.get("updates_per_step",  1))
+        updates_per_step = float(hp.get("updates_per_step", 1))
 
         obs_arr  = np.array(episode_data.observations, dtype=np.float32)
         act_arr  = np.array(episode_data.actions,      dtype=np.float32)
@@ -165,7 +166,10 @@ class TorchSACAgent(BaseAgent):
                 "buffer_size": len(self._buffer),
             }
 
-        n_updates = T * updates_per_step
+        # Truncate the product, not updates_per_step itself -- truncating
+        # a fractional updates_per_step (e.g. 0.5) to int first would zero
+        # it out regardless of T.
+        n_updates = int(T * updates_per_step)
         total_al = 0.0
         total_cl = 0.0
         for _ in range(n_updates):
@@ -174,8 +178,8 @@ class TorchSACAgent(BaseAgent):
             total_cl += cl
 
         return {
-            "actor_loss":  total_al / n_updates,
-            "critic_loss": total_cl / n_updates,
+            "actor_loss":  total_al / max(n_updates, 1),
+            "critic_loss": total_cl / max(n_updates, 1),
             "alpha":       float(self._log_alpha.exp().item()),
             "buffer_size": len(self._buffer),
         }
@@ -224,6 +228,42 @@ class TorchSACAgent(BaseAgent):
         self._c1_tgt.load_state_dict(self._c1.state_dict())
         self._c2_tgt.load_state_dict(self._c2.state_dict())
 
+    @property
+    def resumable_components(self) -> frozenset:
+        # Deliberately does NOT include "replay_buffer" — PrioritizedReplayBuffer
+        # has no state_dict()/load_state_dict(), so resuming from a saved
+        # checkpoint restarts with an EMPTY buffer. For an off-policy algorithm
+        # this is a real gap (see BaseAgent.is_on_policy docstring): the first
+        # batch_size steps after resume have no replay data to sample from.
+        # Claiming "replay_buffer" here would be dishonest about what's
+        # actually restored — see resume_state.py's module docstring.
+        return frozenset({"model", "optimizer", "rng"})
+
+    def save_full_state(self, path: str) -> None:
+        self.save_weights(path)
+        save_torch_resume_state(
+            path,
+            actor_opt=self._actor_opt,
+            critic_opt=self._critic_opt,
+            alpha_opt=self._alpha_opt,
+        )
+
+    def load_full_state(self, path: str) -> None:
+        self.load_weights(path)
+        load_torch_resume_state(
+            path,
+            actor_opt=self._actor_opt,
+            critic_opt=self._critic_opt,
+            alpha_opt=self._alpha_opt,
+        )
+
+    def get_weights(self) -> dict:
+        return {
+            "actor":   {k: v.detach().clone() for k, v in self._actor.state_dict().items()},
+            "critic1": {k: v.detach().clone() for k, v in self._c1.state_dict().items()},
+            "critic2": {k: v.detach().clone() for k, v in self._c2.state_dict().items()},
+        }
+
     def average_weights(self, paths: list) -> None:
         import torch
         n = len(paths)

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/algorithms/torch_td3.py

@@ -174,7 +174,7 @@ class TorchTD3Agent(BaseAgent):
         gamma            = float(hp.get("gamma",            0.99))
         tau              = float(hp.get("tau",              0.005))
         batch_size       = int(hp.get("batch_size",         256))
-        updates_per_step = int(hp.get("updates_per_step",   1))
+        updates_per_step = float(hp.get("updates_per_step", 1))
         policy_delay     = int(hp.get("policy_delay",       2))
 
         obs_arr  = np.array(episode_data.observations, dtype=np.float32)
@@ -198,7 +198,10 @@ class TorchTD3Agent(BaseAgent):
                 "policy_update": 0,
             }
 
-        n_updates = T * updates_per_step
+        # Truncate the product, not updates_per_step itself -- truncating
+        # a fractional updates_per_step (e.g. 0.5, meaning "update every
+        # other step") to int first would zero it out regardless of T.
+        n_updates = int(T * updates_per_step)
         total_al = 0.0
         total_cl = 0.0
         n_actor_updates = 0
@@ -214,7 +217,7 @@ class TorchTD3Agent(BaseAgent):
 
         return {
             "actor_loss":    total_al / max(n_actor_updates, 1),
-            "critic_loss":   total_cl / n_updates,
+            "critic_loss":   total_cl / max(n_updates, 1),
             "buffer_size":   len(self._buffer),
             "policy_update": int(n_actor_updates > 0),
         }

tensor_optix-2.0.2/tensor_optix/core/async_callback.py

@@ -0,0 +1,131 @@
+"""
+AsyncCallbackAdapter — a non-blocking wrapper around any LoopCallback.
+
+LoopController fires callbacks synchronously, inline, on the training
+thread (see LoopController._fire) -- a slow callback (GUI update, network
+I/O, an unbuffered disk write) stalls training for as long as it takes to
+return. This adapter enqueues each event onto a bounded queue and lets a
+single background thread drain it into the wrapped callback, so the
+training thread's call into _fire() always returns immediately regardless
+of how slow the wrapped callback is.
+
+This does NOT change LoopController's own contract -- callbacks are still
+fired synchronously by design (so a broken callback can't silently corrupt
+shared state via unsynchronized concurrent access to the agent/graph). Use
+this adapter only around callbacks that are safe to run on a separate
+thread and don't need to observe training state synchronously.
+"""
+import logging
+import queue
+import threading
+from typing import Optional
+
+from .loop_controller import LoopCallback
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncCallbackAdapter(LoopCallback):
+    """
+    Wraps `inner` so its hooks run on a background thread via a bounded
+    queue, never on the training thread.
+
+    Overflow policy: when the queue is full, the OLDEST pending event is
+    dropped to make room for the newest one -- a consumer showing live
+    state (e.g. a GUI) cares about keeping up with "now," not replaying
+    every historical event. Dropped-event counts are tracked and exposed
+    via `dropped_count` so a consumer can detect and report that it's
+    falling behind.
+
+    Usage:
+        adapter = AsyncCallbackAdapter(MyGuiCallback())
+        controller = LoopController(..., callbacks=[adapter])
+        ...
+        controller.run()
+        adapter.close()  # stop the background thread once training is done
+    """
+
+    def __init__(self, inner: LoopCallback, maxsize: int = 256):
+        self._inner = inner
+        self._queue: "queue.Queue" = queue.Queue(maxsize=maxsize)
+        self._dropped_count = 0
+        self._lock = threading.Lock()
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._drain_loop, daemon=True)
+        self._thread.start()
+
+    def _drain_loop(self) -> None:
+        while not self._stop.is_set() or not self._queue.empty():
+            try:
+                event, args = self._queue.get(timeout=0.1)
+            except queue.Empty:
+                continue
+            try:
+                getattr(self._inner, event)(*args)
+            except Exception as e:
+                logger.warning(
+                    "AsyncCallbackAdapter: inner callback %s raised during %s: %s",
+                    self._inner, event, e,
+                )
+
+    def _enqueue(self, event: str, *args) -> None:
+        try:
+            self._queue.put_nowait((event, args))
+            return
+        except queue.Full:
+            pass
+        # Drop the oldest pending event to make room for this one.
+        try:
+            self._queue.get_nowait()
+            with self._lock:
+                self._dropped_count += 1
+        except queue.Empty:
+            pass
+        try:
+            self._queue.put_nowait((event, args))
+        except queue.Full:
+            with self._lock:
+                self._dropped_count += 1
+
+    def on_loop_start(self) -> None:
+        self._enqueue("on_loop_start")
+
+    def on_loop_stop(self) -> None:
+        self._enqueue("on_loop_stop")
+
+    def on_episode_end(self, episode_id, eval_metrics) -> None:
+        self._enqueue("on_episode_end", episode_id, eval_metrics)
+
+    def on_improvement(self, snapshot) -> None:
+        self._enqueue("on_improvement", snapshot)
+
+    def on_plateau(self, episode_id, state) -> None:
+        self._enqueue("on_plateau", episode_id, state)
+
+    def on_dormant(self, episode_id) -> None:
+        self._enqueue("on_dormant", episode_id)
+
+    def on_degradation(self, episode_id, eval_metrics) -> None:
+        self._enqueue("on_degradation", episode_id, eval_metrics)
+
+    def on_hyperparam_update(self, old_params, new_params) -> None:
+        self._enqueue("on_hyperparam_update", old_params, new_params)
+
+    @property
+    def dropped_count(self) -> int:
+        with self._lock:
+            return self._dropped_count
+
+    @property
+    def pending_count(self) -> int:
+        return self._queue.qsize()
+
+    def close(self, timeout: Optional[float] = None) -> None:
+        """
+        Stop the background thread after draining whatever is already
+        queued. Call once the owning LoopController.run() has returned
+        (e.g. from a different, synchronous callback's on_loop_stop, or
+        by the caller directly) to avoid leaking the thread indefinitely.
+        """
+        self._stop.set()
+        self._thread.join(timeout=timeout)

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/core/base_agent.py

@@ -56,6 +56,57 @@ class BaseAgent(ABC):
     def load_weights(self, path: str) -> None:
         """Restore model weights from the given path."""
 
+    @property
+    def resumable_components(self) -> frozenset:
+        """
+        Which pieces of full training state save_full_state() actually
+        captures, beyond weights — e.g. frozenset({"optimizer", "rng"}).
+
+        Default (empty frozenset): save_full_state()/load_full_state()
+        are equivalent to save_weights()/load_weights() — weights only,
+        no optimizer/scheduler/RNG continuity. CheckpointRegistry records
+        this at save time so a checkpoint can be honestly labeled exact,
+        partial, or weights-only rather than assuming every checkpoint
+        supports exact training continuation — weights alone do not
+        guarantee that.
+
+        Override together with save_full_state()/load_full_state() in a
+        framework-specific subclass. See TorchAgent for a reference
+        implementation built on core.resume_state.
+        """
+        return frozenset()
+
+    def save_full_state(self, path: str) -> None:
+        """
+        Persist everything needed to resume training, not just weights.
+
+        Default: identical to save_weights() — a safe no-op for any agent
+        that hasn't opted into full resumability. Override alongside
+        resumable_components and load_full_state().
+        """
+        self.save_weights(path)
+
+    def load_full_state(self, path: str) -> None:
+        """
+        Restore everything save_full_state() persisted.
+        Default: identical to load_weights().
+        """
+        self.load_weights(path)
+
+    def get_weights(self) -> dict:
+        """
+        Return an in-memory snapshot of the model's learnable parameters,
+        without a disk round-trip through save_weights()/load_weights().
+
+        Default: raises NotImplementedError. Override in framework-specific
+        subclasses that can read parameters directly (currently TorchAgent
+        and its subclasses that reuse TorchAgent's model attribute).
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support get_weights(). "
+            f"Use save_weights()/load_weights() for disk-based access."
+        )
+
     @property
     def is_on_policy(self) -> bool:
         """

{tensor_optix-2.0.0 → tensor_optix-2.0.2}/tensor_optix/core/checkpoint_registry.py

@@ -46,7 +46,13 @@ class CheckpointRegistry:
         weights_path = os.path.join(snapshot_dir, "weights")
         os.makedirs(snapshot_dir, exist_ok=True)
 
-        agent.save_weights(weights_path)
+        # save_full_state() defaults to save_weights() for any agent that
+        # hasn't opted into full resumability. Fall back directly to
+        # save_weights() for duck-typed agents that don't subclass
+        # BaseAgent (and so don't inherit that default) at all.
+        save_fn = getattr(agent, "save_full_state", agent.save_weights)
+        save_fn(weights_path)
+        resumable_components = sorted(getattr(agent, "resumable_components", frozenset()))
 
         snapshot = PolicySnapshot(
             snapshot_id=snapshot_id,
@@ -54,6 +60,7 @@ class CheckpointRegistry:
             hyperparams=hyperparams,
             weights_path=weights_path,
             episode_id=eval_metrics.episode_id,
+            resumable_components=frozenset(resumable_components),
         )
 
         metadata = {
@@ -64,6 +71,7 @@ class CheckpointRegistry:
             "metrics": eval_metrics.metrics,
             "hyperparams": hyperparams.params,
             "weights_path": weights_path,
+            "resumable_components": resumable_components,
         }
         with open(os.path.join(snapshot_dir, "metadata.json"), "w") as f:
             json.dump(metadata, f, indent=2)
@@ -141,7 +149,8 @@ class CheckpointRegistry:
             self._best = self._load_snapshot_from_dir(best_entry["snapshot_dir"])
 
         if self._best is not None:
-            agent.load_weights(self._best.weights_path)
+            load_fn = getattr(agent, "load_full_state", agent.load_weights)
+            load_fn(self._best.weights_path)
         return self._best
 
     @property
@@ -202,4 +211,5 @@ class CheckpointRegistry:
             weights_path=data["weights_path"],
             episode_id=data["episode_id"],
             timestamp=data["timestamp"],
+            resumable_components=frozenset(data.get("resumable_components", [])),
         )