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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensor-optix
-Version: 2.0.2
+Version: 2.0.3
 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
@@ -35,6 +35,8 @@ Requires-Dist: gymnasium[atari]>=1.0.0; extra == "atari"
 Requires-Dist: ale-py>=0.8; extra == "atari"
 Provides-Extra: mujoco
 Requires-Dist: gymnasium[mujoco]>=1.0.0; extra == "mujoco"
+Provides-Extra: pettingzoo
+Requires-Dist: pettingzoo>=1.24.0; extra == "pettingzoo"
 Provides-Extra: all
 Requires-Dist: tensorflow>=2.18.0; extra == "all"
 Requires-Dist: torch>=2.11.0; extra == "all"
@@ -69,13 +71,22 @@ Requires-Dist: torch>=2.11.0; extra == "dev"
 
 # tensor-optix
 
-tensor-optix is a training loop framework with statistical convergence control, online hyperparameter optimisation, and an optional neuroevolution subsystem for dynamic topology.
+tensor-optix is a training framework for reinforcement learning and for supervised/unsupervised machine learning. Both have statistical convergence control, online hyperparameter optimisation, and real checkpointing.
 
-The core loop runs your agent against a pipeline, maintains a separate validation signal, and manages four states: ACTIVE, COOLING, DORMANT, and watchdog shutdown. Convergence is detected using a corrected t-test on the smoothed score slope plus lag-1 autocorrelation, not a fixed patience counter. Hyperparameters are updated every episode via SPSA gradient estimates, with automatic routing to momentum-based or sign-only updates depending on the autocorrelation structure of the score landscape. Checkpointing and rollback are driven by the validation signal only, never training score. On DORMANT, a MetaController evaluates the generalization gap and its slope to decide between spawning a policy variant, pruning the ensemble, or stopping.
+It has two parallel cores:
 
-The neuroevo subsystem (`pip install tensor-optix[neuroevo]`) represents the policy as a `NeuronGraph`: a mutable directed graph of heterogeneous scalar neurons (point, GRU, LSTM, trainable-GRU, trainable-LSTM) with variable-delay edges. Weights for excitatory and inhibitory neurons follow softplus Dale's Law: raw parameter θ maps to `softplus(θ)` (excitatory) or `-softplus(θ)` (inhibitory), eliminating gradient dead zones at weight boundaries. `TopologyController` runs as a loop callback and evaluates three independent signals per episode: improvement slope significance, residual autocorrelation structure, and gradient utilization across hidden neurons. All three must cross their thresholds before a grow operation fires. Pruning is by importance score (incident edge weight magnitude times mean absolute activation). Merging is by Pearson correlation of per-episode activation histories. After every structural mutation the controller calls `graph.invalidate_compile()` to reset for the new topology. `TopologyAwareAdam` resets momentum state for parameters affected by any structural change. `BrainNetwork` composes multiple named `NeuronGraph` regions with sparse learnable inter-region edges — when a neuron is pruned from a region, all inter-region edges referencing it are automatically cleaned up. `HebbianHook` accumulates co-activation products across each episode and applies an Oja-style weight update after the PPO gradient step. `NeuromodulatorSignal` maps a `RegimeDetector` output (trending / ranging / volatile) to simultaneous changes in Hebbian learning rate, entropy coefficient, and topology grow/prune thresholds.
+- **RL core**. An episode-driven loop (`LoopController`/`RLOptimizer`) built around `BaseAgent`'s `act()`/`learn()` contract. Statistical convergence control, SPSA-based online hyperparameter tuning, and an optional neuroevolution subsystem for dynamic topology.
+- **ML core** (`tensor_optix.ml_core`). An epoch/batch-driven loop (`Trainer`, wrapped by `BaseMLModel.fit()`) built around a sklearn-style `fit()`/`predict()`/`evaluate()` contract. Real classification metrics (accuracy, precision/recall/F1, confusion matrix) and regression metrics (MAE, RMSE, R²), validation-based early stopping, and checkpointing on the best metric.
 
-The entire system, including neuroevo, is accessed through a six-method `BaseAgent` interface:
+Both cores reuse the same lower-level, algorithm-agnostic machinery: `CheckpointRegistry`, `BackoffScheduler`'s convergence/plateau detection, `SPSAOptimizer`. Each core has its own contract on top because an RL episode and a training epoch aren't the same shape. See [Concepts](https://sup3rus3r.github.io/tensor-optix/getting-started/concepts/).
+
+## RL core
+
+The RL loop runs your agent against a pipeline, maintains a separate validation signal, and manages four states: ACTIVE, COOLING, DORMANT, and watchdog shutdown. Convergence is detected via a linear trend fit over the recent score window compared against an adaptive, noise-scaled floor, not a fixed patience counter (see [Concepts](https://sup3rus3r.github.io/tensor-optix/getting-started/concepts/) for the exact math; `BackoffScheduler` is the implementation). Hyperparameters are updated every episode via SPSA gradient estimates, with automatic routing to momentum-based or sign-only updates depending on the autocorrelation structure of the score landscape (`AdaptiveOptimizer`). Checkpointing and rollback are driven by the validation signal only, never training score. On DORMANT, a MetaController evaluates the generalization gap and its slope to decide between spawning a policy variant, pruning the ensemble, or stopping.
+
+The neuroevo subsystem (`pip install tensor-optix[neuroevo]`) represents the policy as a `NeuronGraph`: a mutable directed graph of heterogeneous scalar neurons (point, GRU, LSTM, trainable-GRU, trainable-LSTM) with variable-delay edges. Weights for excitatory and inhibitory neurons follow softplus Dale's Law: raw parameter θ maps to `softplus(θ)` (excitatory) or `-softplus(θ)` (inhibitory), eliminating gradient dead zones at weight boundaries. `TopologyController` runs as a loop callback and evaluates three independent signals per episode: improvement slope significance, residual autocorrelation structure, and gradient utilization across hidden neurons. All three must cross their thresholds before a grow operation fires. Pruning is by importance score (incident edge weight magnitude times mean absolute activation). Merging is by Pearson correlation of per-episode activation histories. After every structural mutation the controller calls `graph.invalidate_compile()` to reset for the new topology. `TopologyAwareAdam` resets momentum state for parameters affected by any structural change. `BrainNetwork` composes multiple named `NeuronGraph` regions with sparse learnable inter-region edges. When a neuron is pruned from a region, all inter-region edges referencing it are automatically cleaned up. `HebbianHook` accumulates co-activation products across each episode and applies an Oja-style weight update after the PPO gradient step. `NeuromodulatorSignal` maps a `RegimeDetector` output (trending / ranging / volatile) to simultaneous changes in Hebbian learning rate, entropy coefficient, and topology grow/prune thresholds.
+
+The RL core (including neuroevo) is accessed through a six-method `BaseAgent` interface:
 
 ```python
 class BaseAgent(ABC):
@@ -89,6 +100,40 @@ class BaseAgent(ABC):
 
 Fifteen agents are included across PyTorch, TensorFlow, and JAX/Flax. Bring your own by implementing the interface above.
 
+## ML core
+
+The ML loop trains your model for a number of epochs, evaluates on a held-out validation set each epoch, and stops early once the validation metric plateaus. It restores the best checkpoint automatically. It is accessed through a `BaseMLModel` interface, with `fit()` as the one-call entry point built on top:
+
+```python
+class BaseMLModel(ABC):
+    def predict(self, X) -> Any: ...
+    def train_step(self, batch) -> float: ...
+    def evaluate(self, data) -> MLMetrics: ...
+    def iter_batches(self, data, batch_size: int): ...
+    def save_weights(self, path: str) -> None: ...
+    def load_weights(self, path: str) -> None: ...
+    # fit(train_data, val_data, epochs, early_stopping_patience, ...) is provided
+```
+
+```python
+import torch.nn as nn
+from tensor_optix.ml_core import TorchMLModel
+
+model = TorchMLModel(
+    nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)),
+    loss_fn=nn.CrossEntropyLoss(),
+)
+history = model.fit(
+    (X_train, y_train), (X_val, y_val),
+    epochs=50, early_stopping_patience=5, checkpoint_dir="./checkpoints",
+)
+print(history.val_metrics[-1].metrics)  # {'accuracy': 0.94, 'f1_macro': 0.93, ...}
+```
+
+Metrics are auto-detected from the target dtype. Integer/class-index targets get classification metrics (accuracy, per-class and macro precision/recall/F1, confusion matrix). Float targets get regression metrics (MAE, RMSE, R²). Pass `metrics="classification"`, `metrics="regression"`, or a list of custom `(y_true, y_pred) -> float` callables to override. `TorchMLModel`, `TFMLModel`, and `FlaxMLModel` are provided. The Torch and TF adapters accept a raw `(X, y)` array pair or a native `Dataset`/`DataLoader`/`tf.data.Dataset`; the Flax adapter accepts a raw `(X, y)` array pair only. `fit()` also takes `lr_schedule=` (see `tensor_optix.ml_core.lr_schedules`), `augment=` for per-batch data augmentation, and there are `freeze_layers()`/`unfreeze_layers()`/`load_pretrained_backbone()` helpers for transfer learning. See [Train an ML model](https://sup3rus3r.github.io/tensor-optix/guides/ml-core/) for the full guide.
+
+There is also an older, RL-loop-based way to train an arbitrary `nn.Module`: `Optimizer(model, dataset)`, which runs the model through the RL loop as a fake episode. See [ML support](#ml-support) below. It predates `ml_core`, is kept for backward compatibility, and has no real classification/regression metrics and no validation-based early stopping. Use `ml_core` for new code.
+
 ---
 
 ## Install
@@ -605,6 +650,8 @@ graph.invalidate_compile()
 
 ## ML support
 
+> **Use [`tensor_optix.ml_core`](#ml-core) for new code.** The `Optimizer(model, dataset)` path below is the original way to train an arbitrary `nn.Module`. It runs the model through the RL loop's `BaseAgent`/`RLOptimizer` machinery by treating each mini-batch as a fake RL episode (`reward = -loss`). It is deprecated but kept working for backward compatibility — `MLAgent` emits a `DeprecationWarning` pointing to `ml_core` on construction, and there is no removal planned. It has no real classification/regression metrics (checkpointing and rollback compare mean "episode return," which is just negative loss) and no validation-split-based early stopping. `ml_core`'s `BaseMLModel`/`Trainer` reuse the same convergence/checkpointing machinery through a real `fit()`/`evaluate()` contract with real metrics.
+
 tensor-optix also supports general ML training via `Optimizer`. Pass any `nn.Module` and a PyTorch `Dataset` or `DataLoader` — the same SPSA hyperparameter tuning, rollback on degradation, convergence detection, and checkpointing apply automatically.
 
 ```python

{tensor_optix-2.0.2 → tensor_optix-2.0.3}/README.md

@@ -1,12 +1,21 @@
 # tensor-optix
 
-tensor-optix is a training loop framework with statistical convergence control, online hyperparameter optimisation, and an optional neuroevolution subsystem for dynamic topology.
+tensor-optix is a training framework for reinforcement learning and for supervised/unsupervised machine learning. Both have statistical convergence control, online hyperparameter optimisation, and real checkpointing.
 
-The core loop runs your agent against a pipeline, maintains a separate validation signal, and manages four states: ACTIVE, COOLING, DORMANT, and watchdog shutdown. Convergence is detected using a corrected t-test on the smoothed score slope plus lag-1 autocorrelation, not a fixed patience counter. Hyperparameters are updated every episode via SPSA gradient estimates, with automatic routing to momentum-based or sign-only updates depending on the autocorrelation structure of the score landscape. Checkpointing and rollback are driven by the validation signal only, never training score. On DORMANT, a MetaController evaluates the generalization gap and its slope to decide between spawning a policy variant, pruning the ensemble, or stopping.
+It has two parallel cores:
 
-The neuroevo subsystem (`pip install tensor-optix[neuroevo]`) represents the policy as a `NeuronGraph`: a mutable directed graph of heterogeneous scalar neurons (point, GRU, LSTM, trainable-GRU, trainable-LSTM) with variable-delay edges. Weights for excitatory and inhibitory neurons follow softplus Dale's Law: raw parameter θ maps to `softplus(θ)` (excitatory) or `-softplus(θ)` (inhibitory), eliminating gradient dead zones at weight boundaries. `TopologyController` runs as a loop callback and evaluates three independent signals per episode: improvement slope significance, residual autocorrelation structure, and gradient utilization across hidden neurons. All three must cross their thresholds before a grow operation fires. Pruning is by importance score (incident edge weight magnitude times mean absolute activation). Merging is by Pearson correlation of per-episode activation histories. After every structural mutation the controller calls `graph.invalidate_compile()` to reset for the new topology. `TopologyAwareAdam` resets momentum state for parameters affected by any structural change. `BrainNetwork` composes multiple named `NeuronGraph` regions with sparse learnable inter-region edges — when a neuron is pruned from a region, all inter-region edges referencing it are automatically cleaned up. `HebbianHook` accumulates co-activation products across each episode and applies an Oja-style weight update after the PPO gradient step. `NeuromodulatorSignal` maps a `RegimeDetector` output (trending / ranging / volatile) to simultaneous changes in Hebbian learning rate, entropy coefficient, and topology grow/prune thresholds.
+- **RL core**. An episode-driven loop (`LoopController`/`RLOptimizer`) built around `BaseAgent`'s `act()`/`learn()` contract. Statistical convergence control, SPSA-based online hyperparameter tuning, and an optional neuroevolution subsystem for dynamic topology.
+- **ML core** (`tensor_optix.ml_core`). An epoch/batch-driven loop (`Trainer`, wrapped by `BaseMLModel.fit()`) built around a sklearn-style `fit()`/`predict()`/`evaluate()` contract. Real classification metrics (accuracy, precision/recall/F1, confusion matrix) and regression metrics (MAE, RMSE, R²), validation-based early stopping, and checkpointing on the best metric.
 
-The entire system, including neuroevo, is accessed through a six-method `BaseAgent` interface:
+Both cores reuse the same lower-level, algorithm-agnostic machinery: `CheckpointRegistry`, `BackoffScheduler`'s convergence/plateau detection, `SPSAOptimizer`. Each core has its own contract on top because an RL episode and a training epoch aren't the same shape. See [Concepts](https://sup3rus3r.github.io/tensor-optix/getting-started/concepts/).
+
+## RL core
+
+The RL loop runs your agent against a pipeline, maintains a separate validation signal, and manages four states: ACTIVE, COOLING, DORMANT, and watchdog shutdown. Convergence is detected via a linear trend fit over the recent score window compared against an adaptive, noise-scaled floor, not a fixed patience counter (see [Concepts](https://sup3rus3r.github.io/tensor-optix/getting-started/concepts/) for the exact math; `BackoffScheduler` is the implementation). Hyperparameters are updated every episode via SPSA gradient estimates, with automatic routing to momentum-based or sign-only updates depending on the autocorrelation structure of the score landscape (`AdaptiveOptimizer`). Checkpointing and rollback are driven by the validation signal only, never training score. On DORMANT, a MetaController evaluates the generalization gap and its slope to decide between spawning a policy variant, pruning the ensemble, or stopping.
+
+The neuroevo subsystem (`pip install tensor-optix[neuroevo]`) represents the policy as a `NeuronGraph`: a mutable directed graph of heterogeneous scalar neurons (point, GRU, LSTM, trainable-GRU, trainable-LSTM) with variable-delay edges. Weights for excitatory and inhibitory neurons follow softplus Dale's Law: raw parameter θ maps to `softplus(θ)` (excitatory) or `-softplus(θ)` (inhibitory), eliminating gradient dead zones at weight boundaries. `TopologyController` runs as a loop callback and evaluates three independent signals per episode: improvement slope significance, residual autocorrelation structure, and gradient utilization across hidden neurons. All three must cross their thresholds before a grow operation fires. Pruning is by importance score (incident edge weight magnitude times mean absolute activation). Merging is by Pearson correlation of per-episode activation histories. After every structural mutation the controller calls `graph.invalidate_compile()` to reset for the new topology. `TopologyAwareAdam` resets momentum state for parameters affected by any structural change. `BrainNetwork` composes multiple named `NeuronGraph` regions with sparse learnable inter-region edges. When a neuron is pruned from a region, all inter-region edges referencing it are automatically cleaned up. `HebbianHook` accumulates co-activation products across each episode and applies an Oja-style weight update after the PPO gradient step. `NeuromodulatorSignal` maps a `RegimeDetector` output (trending / ranging / volatile) to simultaneous changes in Hebbian learning rate, entropy coefficient, and topology grow/prune thresholds.
+
+The RL core (including neuroevo) is accessed through a six-method `BaseAgent` interface:
 
 ```python
 class BaseAgent(ABC):
@@ -20,6 +29,40 @@ class BaseAgent(ABC):
 
 Fifteen agents are included across PyTorch, TensorFlow, and JAX/Flax. Bring your own by implementing the interface above.
 
+## ML core
+
+The ML loop trains your model for a number of epochs, evaluates on a held-out validation set each epoch, and stops early once the validation metric plateaus. It restores the best checkpoint automatically. It is accessed through a `BaseMLModel` interface, with `fit()` as the one-call entry point built on top:
+
+```python
+class BaseMLModel(ABC):
+    def predict(self, X) -> Any: ...
+    def train_step(self, batch) -> float: ...
+    def evaluate(self, data) -> MLMetrics: ...
+    def iter_batches(self, data, batch_size: int): ...
+    def save_weights(self, path: str) -> None: ...
+    def load_weights(self, path: str) -> None: ...
+    # fit(train_data, val_data, epochs, early_stopping_patience, ...) is provided
+```
+
+```python
+import torch.nn as nn
+from tensor_optix.ml_core import TorchMLModel
+
+model = TorchMLModel(
+    nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)),
+    loss_fn=nn.CrossEntropyLoss(),
+)
+history = model.fit(
+    (X_train, y_train), (X_val, y_val),
+    epochs=50, early_stopping_patience=5, checkpoint_dir="./checkpoints",
+)
+print(history.val_metrics[-1].metrics)  # {'accuracy': 0.94, 'f1_macro': 0.93, ...}
+```
+
+Metrics are auto-detected from the target dtype. Integer/class-index targets get classification metrics (accuracy, per-class and macro precision/recall/F1, confusion matrix). Float targets get regression metrics (MAE, RMSE, R²). Pass `metrics="classification"`, `metrics="regression"`, or a list of custom `(y_true, y_pred) -> float` callables to override. `TorchMLModel`, `TFMLModel`, and `FlaxMLModel` are provided. The Torch and TF adapters accept a raw `(X, y)` array pair or a native `Dataset`/`DataLoader`/`tf.data.Dataset`; the Flax adapter accepts a raw `(X, y)` array pair only. `fit()` also takes `lr_schedule=` (see `tensor_optix.ml_core.lr_schedules`), `augment=` for per-batch data augmentation, and there are `freeze_layers()`/`unfreeze_layers()`/`load_pretrained_backbone()` helpers for transfer learning. See [Train an ML model](https://sup3rus3r.github.io/tensor-optix/guides/ml-core/) for the full guide.
+
+There is also an older, RL-loop-based way to train an arbitrary `nn.Module`: `Optimizer(model, dataset)`, which runs the model through the RL loop as a fake episode. See [ML support](#ml-support) below. It predates `ml_core`, is kept for backward compatibility, and has no real classification/regression metrics and no validation-based early stopping. Use `ml_core` for new code.
+
 ---
 
 ## Install
@@ -536,6 +579,8 @@ graph.invalidate_compile()
 
 ## ML support
 
+> **Use [`tensor_optix.ml_core`](#ml-core) for new code.** The `Optimizer(model, dataset)` path below is the original way to train an arbitrary `nn.Module`. It runs the model through the RL loop's `BaseAgent`/`RLOptimizer` machinery by treating each mini-batch as a fake RL episode (`reward = -loss`). It is deprecated but kept working for backward compatibility — `MLAgent` emits a `DeprecationWarning` pointing to `ml_core` on construction, and there is no removal planned. It has no real classification/regression metrics (checkpointing and rollback compare mean "episode return," which is just negative loss) and no validation-split-based early stopping. `ml_core`'s `BaseMLModel`/`Trainer` reuse the same convergence/checkpointing machinery through a real `fit()`/`evaluate()` contract with real metrics.
+
 tensor-optix also supports general ML training via `Optimizer`. Pass any `nn.Module` and a PyTorch `Dataset` or `DataLoader` — the same SPSA hyperparameter tuning, rollback on degradation, convergence detection, and checkpointing apply automatically.
 
 ```python

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensor-optix"
-version = "2.0.2"
+version = "2.0.3"
 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"
@@ -60,6 +60,10 @@ atari = [
 mujoco = [
     "gymnasium[mujoco]>=1.0.0",
 ]
+# Multi-agent environments (tensor_optix.multi_agent)
+pettingzoo = [
+    "pettingzoo>=1.24.0",
+]
 # Convenience bundles
 all = [
     "tensorflow>=2.18.0",

{tensor_optix-2.0.2 → tensor_optix-2.0.3}/tensor_optix/algorithms/tf_dqn.py

@@ -7,6 +7,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.replay_buffer import PrioritizedReplayBuffer
+from tensor_optix.core.obs_utils import to_batch_numpy, to_tf_batch, already_batched_to_tf, batch_index
 
 
 class TFDQNAgent(BaseAgent):
@@ -97,7 +98,7 @@ class TFDQNAgent(BaseAgent):
         epsilon = float(self._hyperparams.params.get("epsilon", 0.1))
         if np.random.random() < epsilon:
             return np.random.randint(self._n_actions)
-        obs = tf.cast(np.atleast_2d(observation), tf.float32)
+        obs = to_tf_batch(observation)
         q_values = self._q(obs, training=False)
         return int(tf.argmax(q_values, axis=-1).numpy()[0])
 
@@ -109,7 +110,10 @@ class TFDQNAgent(BaseAgent):
         epsilon_decay      = float(hp.get("epsilon_decay",      0.995))
         epsilon_min        = float(hp.get("epsilon_min",        0.05))
 
-        obs_arr  = np.array(episode_data.observations, dtype=np.float32)
+        # DQN pushes the RAW observation into the buffer (not pre-encoded)
+        # -- same single-network constraint as TorchDQNAgent, see
+        # torch_dqn.py's learn() comment for the full rationale.
+        obs_arr  = to_batch_numpy(list(episode_data.observations))
         act_arr  = episode_data.actions
         rew_arr  = episode_data.rewards
         done_arr = episode_data.dones
@@ -117,8 +121,8 @@ class TFDQNAgent(BaseAgent):
         T = len(rew_arr)
         for t in range(T - 1):
             self._buffer.push(
-                obs_arr[t], int(act_arr[t]), float(rew_arr[t]),
-                obs_arr[t + 1], bool(done_arr[t]),
+                batch_index(obs_arr, t), int(act_arr[t]), float(rew_arr[t]),
+                batch_index(obs_arr, t + 1), bool(done_arr[t]),
             )
         self._buffer.flush_episode()
 
@@ -138,10 +142,10 @@ class TFDQNAgent(BaseAgent):
 
         for _ in range(n_updates):
             obs_b, act_b, rew_b, next_obs_b, done_b, weights, indices, n_steps = self._buffer.sample(batch_size)
-            obs_b      = tf.cast(obs_b,      tf.float32)
+            obs_b      = already_batched_to_tf(obs_b)
             act_b      = tf.cast(act_b,      tf.int32)
             rew_b      = tf.cast(rew_b,      tf.float32)
-            next_obs_b = tf.cast(next_obs_b, tf.float32)
+            next_obs_b = already_batched_to_tf(next_obs_b)
             done_b     = tf.cast(done_b,     tf.float32)
             weights    = tf.cast(weights,    tf.float32)
             gammas_n   = tf.cast(gamma ** n_steps.astype(np.float32), tf.float32)

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

@@ -4,6 +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, resolve_gae_bootstrap
+from tensor_optix.core.obs_utils import to_batch_numpy, to_tf_batch, already_batched_to_tf
 
 
 class TFPPOAgent(BaseAgent):
@@ -89,7 +90,7 @@ class TFPPOAgent(BaseAgent):
         Sample action from π(·|obs), cache log_prob and V(obs) for learn().
         Returns a Python int (action index).
         """
-        obs = tf.cast(np.atleast_2d(observation), tf.float32)
+        obs = to_tf_batch(observation)
 
         logits = self._actor(obs, training=False)                 # [1, n_actions]
         log_probs_all = tf.nn.log_softmax(logits)                 # [1, n_actions]
@@ -98,7 +99,9 @@ class TFPPOAgent(BaseAgent):
 
         value = float(tf.squeeze(self._critic(obs, training=False)).numpy())
 
-        self._cache_obs.append(np.squeeze(observation))
+        # Cache the raw observation -- learn() re-batches via
+        # to_batch_numpy(), which handles both array and dict observations.
+        self._cache_obs.append(np.squeeze(observation) if not isinstance(observation, dict) else observation)
         self._cache_log_probs.append(log_prob)
         self._cache_values.append(value)
 
@@ -134,7 +137,7 @@ class TFPPOAgent(BaseAgent):
                 f"Cache underflow: expected >= {T} entries but got {len(self._cache_obs)}. "
                 "Ensure act() is called exactly once per environment step."
             )
-        obs_arr      = np.array(self._cache_obs[:T],      dtype=np.float32)
+        obs_arr      = to_batch_numpy(self._cache_obs[:T])
         old_lp_arr   = np.array(self._cache_log_probs[:T], dtype=np.float32)
         values_arr   = np.array(self._cache_values[:T],    dtype=np.float32)
         rewards      = list(episode_data.rewards)
@@ -155,7 +158,7 @@ class TFPPOAgent(BaseAgent):
         # 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)
+            obs_t = to_tf_batch(obs)
             return float(tf.squeeze(self._critic(obs_t, training=False)).numpy())
 
         last_value, n_envs = resolve_gae_bootstrap(episode_data, _critic_value)
@@ -183,7 +186,7 @@ class TFPPOAgent(BaseAgent):
 
         for _ in range(n_epochs):
             for batch in make_minibatches(rollout, mb_size):
-                obs_b   = tf.cast(batch["obs"],        tf.float32)
+                obs_b   = already_batched_to_tf(batch["obs"])
                 act_b   = tf.cast(batch["actions"],    tf.int32)
                 old_b   = tf.cast(batch["old_lp"],     tf.float32)
                 adv_b   = tf.cast(batch["advantages"], tf.float32)
@@ -248,7 +251,7 @@ class TFPPOAgent(BaseAgent):
         Used by PolicyManager.ensemble_action() to average distributions
         across ensemble members rather than averaging sampled actions.
         """
-        obs = tf.cast(np.atleast_2d(observation), tf.float32)
+        obs = to_tf_batch(observation)
         logits = self._actor(obs, training=False)
         return tf.nn.softmax(logits)[0].numpy()
 

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

@@ -5,6 +5,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, resolve_gae_bootstrap
+from tensor_optix.core.obs_utils import to_batch_numpy, to_tf_batch, already_batched_to_tf
 
 LOG_STD_MIN = -5
 LOG_STD_MAX = 2
@@ -133,7 +134,7 @@ class TFGaussianPPOAgent(BaseAgent):
         Sample action ∈ (-1, 1)^action_dim; cache log_prob and V(obs) for learn().
         Returns a float numpy array of shape [action_dim].
         """
-        obs = tf.cast(np.atleast_2d(observation), tf.float32)
+        obs = to_tf_batch(observation)
 
         out     = self._actor(obs, training=False)             # [1, 2*action_dim]
         mean    = out[:, :self._action_dim]
@@ -147,7 +148,7 @@ class TFGaussianPPOAgent(BaseAgent):
         log_prob = self._gaussian_log_prob(pre_tanh, mean, log_std)  # [1]
         value    = float(tf.squeeze(self._critic(obs, training=False)).numpy())
 
-        self._cache_obs.append(np.squeeze(observation))
+        self._cache_obs.append(np.squeeze(observation) if not isinstance(observation, dict) else observation)
         self._cache_log_probs.append(float(tf.squeeze(log_prob).numpy()))
         self._cache_values.append(value)
 
@@ -179,7 +180,7 @@ class TFGaussianPPOAgent(BaseAgent):
                 "Ensure act() is called exactly once per environment step."
             )
 
-        obs_arr    = np.array(self._cache_obs[:T],       dtype=np.float32)
+        obs_arr    = to_batch_numpy(self._cache_obs[:T])
         old_lp_arr = np.array(self._cache_log_probs[:T], dtype=np.float32)
         val_arr    = np.array(self._cache_values[:T],    dtype=np.float32)
         act_arr    = np.array(episode_data.actions,      dtype=np.float32)
@@ -194,7 +195,7 @@ class TFGaussianPPOAgent(BaseAgent):
             rewards = list(self._reward_normalizer.normalize(np.array(rewards, dtype=np.float32)))
 
         def _critic_value(obs):
-            obs_t = tf.cast(np.atleast_2d(obs), tf.float32)
+            obs_t = to_tf_batch(obs)
             return float(tf.squeeze(self._critic(obs_t, training=False)).numpy())
 
         last_value, n_envs = resolve_gae_bootstrap(episode_data, _critic_value)
@@ -220,7 +221,7 @@ class TFGaussianPPOAgent(BaseAgent):
 
         for _ in range(n_epochs):
             for batch in make_minibatches(rollout, mb_size):
-                obs_b = tf.cast(batch["obs"],        tf.float32)
+                obs_b = already_batched_to_tf(batch["obs"])
                 act_b = tf.cast(batch["actions"],    tf.float32)
                 old_b = tf.cast(batch["old_lp"],     tf.float32)
                 adv_b = tf.cast(batch["advantages"], tf.float32)

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

@@ -7,6 +7,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.replay_buffer import PrioritizedReplayBuffer
+from tensor_optix.core.obs_utils import to_batch_numpy, to_tf_batch, already_batched_to_tf
 
 
 class TFSACAgent(BaseAgent):
@@ -132,14 +133,28 @@ class TFSACAgent(BaseAgent):
     def is_on_policy(self) -> bool:
         return False  # replay buffer — rollback without buffer clear is harmful
 
+    def _encode(self, observation_batched):
+        """Run the actor's encoder alone, if it has one (for non-flat
+        observations); identity pass-through for the flat-Box case,
+        exactly matching today's behavior. Used to pre-encode observations
+        before they enter the replay buffer -- see tensor_optix.encoders'
+        module docstring for why SAC's critic only ever sees already-
+        encoded flat vectors, never raw observations. self._actor.encoder
+        is the standalone encoder Model, attached by the factory alongside
+        head_model (see _make_sac_tf), so this never has to split the
+        composed actor Model's graph apart at call time."""
+        encoder = getattr(self._actor, "encoder", None)
+        if encoder is None:
+            return observation_batched
+        return encoder(observation_batched, training=False)
+
     def act(self, observation) -> np.ndarray:
         """
         Sample action from the squashed Gaussian policy.
         Returns np.ndarray of shape [action_dim], values in (-1, 1).
         """
-        obs = tf.cast(np.atleast_2d(observation), tf.float32)
+        obs = to_tf_batch(observation)
         action, _ = self._sample_action(obs, training=False)
-        self._last_obs = np.array(observation, dtype=np.float32)
         return action.numpy()[0]
 
     def learn(self, episode_data: EpisodeData) -> dict:
@@ -153,7 +168,8 @@ class TFSACAgent(BaseAgent):
         batch_size        = int(hp.get("batch_size",          256))
         updates_per_step  = float(hp.get("updates_per_step",  1))
 
-        obs_arr  = np.array(episode_data.observations, dtype=np.float32)
+        obs_batched = to_tf_batch(list(episode_data.observations))
+        obs_arr  = self._encode(obs_batched).numpy()
         act_arr  = np.array(episode_data.actions,      dtype=np.float32)
         rew_arr  = episode_data.rewards
         done_arr = episode_data.dones
@@ -267,15 +283,25 @@ class TFSACAgent(BaseAgent):
     # Internal helpers
     # ------------------------------------------------------------------
 
-    def _sample_action(self, obs_tf, training: bool):
+    def _sample_action(self, obs_tf, training: bool, pre_encoded: bool = False):
         """
         Reparameterized squashed Gaussian sample.
 
+        pre_encoded=True: obs_tf is already a flat, encoder-output vector
+        (buffer-sampled data) -- route through the actor's head_model only,
+        skipping its encoder a second time (see tensor_optix.encoders'
+        module docstring). pre_encoded=False (default): obs_tf is a raw
+        observation (from act()) -- route through the full composed actor
+        (encoder + head). For the flat-Box case (actor has no head_model
+        attribute, it IS the head) this falls back to calling it directly,
+        identical to today's behavior.
+
         Returns:
             action:   tanh(mean + std * ε),  shape [batch, action_dim]
             log_prob: log π(action|obs) with tanh correction, shape [batch]
         """
-        out = self._actor(obs_tf, training=training)        # [batch, 2*action_dim]
+        net = getattr(self._actor, "head_model", self._actor) if pre_encoded else self._actor
+        out = net(obs_tf, training=training)        # [batch, 2*action_dim]
         mean, log_std = tf.split(out, 2, axis=-1)
         log_std = tf.clip_by_value(log_std, self.LOG_STD_MIN, self.LOG_STD_MAX)
         std = tf.exp(log_std)
@@ -306,18 +332,23 @@ class TFSACAgent(BaseAgent):
     def _update_step(self, batch_size: int, gamma: float, tau: float):
         """One full SAC gradient step (critic + actor + alpha + soft update)."""
         obs_b, act_b, rew_b, next_obs_b, done_b, weights, indices, n_steps = self._buffer.sample(batch_size)
-        obs_b      = tf.cast(obs_b,      tf.float32)
+        obs_b      = already_batched_to_tf(obs_b)
         act_b      = tf.cast(act_b,      tf.float32)
         rew_b      = tf.cast(rew_b,      tf.float32)
-        next_obs_b = tf.cast(next_obs_b, tf.float32)
+        next_obs_b = already_batched_to_tf(next_obs_b)
         done_b     = tf.cast(done_b,     tf.float32)
         weights    = tf.cast(weights,    tf.float32)
         gammas_n   = tf.cast(gamma ** n_steps.astype(np.float32), tf.float32)
         alpha      = tf.exp(self._log_alpha)
 
+        # Buffer-sampled obs_b/next_obs_b are already flat, pre-encoded
+        # vectors -- actor updates below train only the head's variables,
+        # never the (frozen, collection-time-only) encoder's.
+        actor_trainable_vars = getattr(self._actor, "head_model", self._actor).trainable_variables
+
         # ---- Critic update ----
         with tf.GradientTape() as tape:
-            next_a, next_lp = self._sample_action(next_obs_b, training=False)
+            next_a, next_lp = self._sample_action(next_obs_b, training=False, pre_encoded=True)
             ci_next = self._critic_input(next_obs_b, next_a)
             q1_tgt  = tf.squeeze(self._c1_tgt(ci_next, training=False), -1)
             q2_tgt  = tf.squeeze(self._c2_tgt(ci_next, training=False), -1)
@@ -339,19 +370,19 @@ class TFSACAgent(BaseAgent):
 
         # ---- Actor update ----
         with tf.GradientTape() as tape:
-            new_a, new_lp = self._sample_action(obs_b, training=True)
+            new_a, new_lp = self._sample_action(obs_b, training=True, pre_encoded=True)
             ci_new = self._critic_input(obs_b, new_a)
             q1_new = tf.squeeze(self._c1(ci_new, training=False), -1)
             q2_new = tf.squeeze(self._c2(ci_new, training=False), -1)
             a_loss = tf.reduce_mean(alpha * new_lp - tf.minimum(q1_new, q2_new))
 
-        a_grads = tape.gradient(a_loss, self._actor.trainable_variables)
+        a_grads = tape.gradient(a_loss, actor_trainable_vars)
         a_grads, _ = tf.clip_by_global_norm(a_grads, 10.0)
-        self._actor_opt.apply_gradients(zip(a_grads, self._actor.trainable_variables))
+        self._actor_opt.apply_gradients(zip(a_grads, actor_trainable_vars))
 
         # ---- Alpha (entropy temperature) update ----
         with tf.GradientTape() as tape:
-            _, new_lp2 = self._sample_action(obs_b, training=False)
+            _, new_lp2 = self._sample_action(obs_b, training=False, pre_encoded=True)
             alpha_loss = -tf.reduce_mean(
                 self._log_alpha * tf.stop_gradient(new_lp2 + self._target_entropy)
             )

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

@@ -62,6 +62,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.replay_buffer import PrioritizedReplayBuffer
+from tensor_optix.core.obs_utils import to_tf_batch, already_batched_to_tf
 
 
 class TFTDDAgent(BaseAgent):
@@ -97,6 +98,20 @@ class TFTDDAgent(BaseAgent):
         self._c1_tgt.set_weights(critic1.get_weights())
         self._c2_tgt.set_weights(critic2.get_weights())
 
+        # clone_model() reconstructs the Model graph/config but does not
+        # carry over head_model/encoder (plain Python attributes the
+        # factory attaches post-construction, not part of the Keras
+        # graph/config clone_model reads from) -- re-derive them the same
+        # way so _update_step()'s getattr(self._actor_tgt, "head_model",
+        # ...) lookups actually find the real head-only submodel instead
+        # of silently falling back to the full (encoder+head) actor_tgt.
+        if hasattr(actor, "head_model"):
+            self._actor_tgt.head_model = tf.keras.models.clone_model(actor.head_model)
+            self._actor_tgt.head_model.set_weights(actor.head_model.get_weights())
+        if hasattr(actor, "encoder"):
+            self._actor_tgt.encoder = tf.keras.models.clone_model(actor.encoder)
+            self._actor_tgt.encoder.set_weights(actor.encoder.get_weights())
+
         self._action_dim  = action_dim
         self._actor_opt   = actor_optimizer
         self._critic_opt  = critic_optimizer
@@ -120,9 +135,17 @@ class TFTDDAgent(BaseAgent):
     def is_on_policy(self) -> bool:
         return False
 
+    def _encode(self, observation_batched):
+        """Run the actor's encoder alone, if it has one -- see
+        tf_sac.py's _encode() docstring for the full rationale."""
+        encoder = getattr(self._actor, "encoder", None)
+        if encoder is None:
+            return observation_batched
+        return encoder(observation_batched, training=False)
+
     def act(self, observation) -> np.ndarray:
         """Deterministic action: a = π_θ(s) = tanh(actor(s))."""
-        obs = tf.cast(np.atleast_2d(observation), tf.float32)
+        obs = to_tf_batch(observation)
         action = self._actor(obs, training=False)
         return action.numpy()[0]
 
@@ -134,7 +157,8 @@ class TFTDDAgent(BaseAgent):
         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)
+        obs_batched = to_tf_batch(list(episode_data.observations))
+        obs_arr  = self._encode(obs_batched).numpy()
         act_arr  = np.array(episode_data.actions,      dtype=np.float32)
         rew_arr  = episode_data.rewards
         done_arr = episode_data.dones
@@ -265,14 +289,21 @@ class TFTDDAgent(BaseAgent):
         obs_b, act_b, rew_b, next_b, done_b, weights, indices, n_steps = \
             self._buffer.sample(batch_size)
 
-        obs_b    = tf.cast(obs_b,    tf.float32)
+        obs_b    = already_batched_to_tf(obs_b)
         act_b    = tf.cast(act_b,    tf.float32)
         rew_b    = tf.cast(rew_b,    tf.float32)
-        next_b   = tf.cast(next_b,   tf.float32)
+        next_b   = already_batched_to_tf(next_b)
         done_b   = tf.cast(done_b,   tf.float32)
         weights  = tf.cast(weights,  tf.float32)
         gammas_n = tf.cast(gamma ** n_steps.astype(np.float32), tf.float32)
 
+        # Buffer-sampled obs_b/next_b are already flat, pre-encoded vectors
+        # -- route through head_model only, skipping the encoder a second
+        # time (see tf_sac.py's _update_step comment for the full
+        # rationale). Falls back to the full actor/actor_tgt for flat-Box.
+        actor_head     = getattr(self._actor, "head_model", self._actor)
+        actor_tgt_head = getattr(self._actor_tgt, "head_model", self._actor_tgt)
+
         # ---- Critic update ----
         with tf.GradientTape() as tape:
             # Fix 3: target policy smoothing
@@ -282,7 +313,7 @@ class TFTDDAgent(BaseAgent):
                 target_noise_clip,
             )
             next_action = tf.clip_by_value(
-                self._actor_tgt(next_b, training=False) + noise, -1.0, 1.0
+                actor_tgt_head(next_b, training=False) + noise, -1.0, 1.0
             )
 
             # Fix 1: twin critics
@@ -311,20 +342,20 @@ class TFTDDAgent(BaseAgent):
         a_loss = 0.0
         if update_actor:
             with tf.GradientTape() as tape:
-                new_a  = self._actor(obs_b, training=True)
+                new_a  = actor_head(obs_b, training=True)
                 ci_new = tf.concat([obs_b, new_a], axis=-1)
                 # Deterministic policy gradient: maximise Q1(s, π(s))
                 a_loss_tf = -tf.reduce_mean(
                     tf.squeeze(self._c1(ci_new, training=False), -1)
                 )
 
-            a_grads = tape.gradient(a_loss_tf, self._actor.trainable_variables)
+            a_grads = tape.gradient(a_loss_tf, actor_head.trainable_variables)
             a_grads, _ = tf.clip_by_global_norm(a_grads, 10.0)
             self._actor_opt.apply_gradients(
-                zip(a_grads, self._actor.trainable_variables)
+                zip(a_grads, actor_head.trainable_variables)
             )
             a_loss = float(a_loss_tf.numpy())
-            self._soft_update(self._actor, self._actor_tgt, tau)
+            self._soft_update(actor_head, actor_tgt_head, tau)
 
         self._soft_update(self._c1, self._c1_tgt, tau)
         self._soft_update(self._c2, self._c2_tgt, tau)