"""Tensor-native companion to :mod:`resynthesis.causal_exploration`. This module re-expresses the visited state-action exploration graph as dense ``torch`` tensors so the frontier / UCB / entropy / merge queries become single vectorized kernel launches instead of Python loops over ``dict`` state nodes. Layout ------ Visit counts and outcome sums live in two dense ``[table_size, num_actions]`` float tensors, indexed by a stable SHA-256 hash of the string state key: index = int.from_bytes(sha256(state_key).digest()[:8], 'little') % table_size Collisions are acceptable (hash-bucketed, like a hash table) -- two keys landing in the same bucket simply share a row, exactly as they would share a dict slot. The hash is identical to the one used by the other ``*_tensor.py`` companions so a state addressed in :mod:`value_function_tensor` resolves to the same row here. The ``TensorExplorationGraph`` is a :class:`torch.nn.Module`: * It carries a *learnable* state-value embedding ``state_value`` of shape ``[table_size]`` (``requires_grad=True``) used by the UCB exploit term and the optional value-prior blend. Gradients flow through ``recommend_next_scores`` / ``ucb_scores`` / ``state_value``. * Visit counts / outcome sums are buffers (``register_buffer``) -- they are exploration statistics, not learned parameters, so they do not take gradients but they DO move with ``.to(device)`` / ``.cuda()`` and serialize in ``state_dict()``. * All computation is tensor ops -- no Python loops over states/actions in the hot path. ``index_add_`` / masked select / ``torch.where`` replace the dict walks of the Python reference. The original :mod:`resynthesis.causal_exploration` module is preserved as the torch-free reference; this companion is additive and importable independently. """ from __future__ import annotations import hashlib import math from collections.abc import Sequence import torch from torch import Tensor, nn CAUSAL_EXPLORATION_TENSOR_SCHEMA = "nnf.resynthesis.causal_exploration_tensor.v1" # Strategy names mirror the Python reference module. STRATEGY_SHORTEST_TO_UNTESTED = "shortest_to_untested" STRATEGY_LEAST_SAMPLED = "least_sampled" STRATEGY_BEST_OUTCOME = "best_outcome" STRATEGY_UCB = "ucb" STRATEGY_HYPOTHESIS_DRIVEN = "hypothesis_driven" UCB_DEFAULT_EXPLORATION = math.sqrt(2.0) RHAE_DEFAULT_CAP = 1.15 DEFAULT_TABLE_SIZE = 4096 DEFAULT_NUM_ACTIONS = 8 def state_hash_index( state_key: str, *, table_size: int = DEFAULT_TABLE_SIZE, ) -> int: """Stable SHA-256 -> integer table row index. Same formula across every ``*_tensor.py`` companion so a state addressed in one module resolves to the same row in the others. """ if table_size <= 0: raise ValueError("table_size must be positive") digest = hashlib.sha256(state_key.encode("utf-8")).digest()[:8] return int.from_bytes(digest, "little") % table_size def state_hash_indices( state_keys: Sequence[str], *, table_size: int = DEFAULT_TABLE_SIZE, ) -> Tensor: """Vectorized :func:`state_hash_index` over a batch of keys -> ``int64[N]``.""" if table_size <= 0: raise ValueError("table_size must be positive") return torch.tensor( [state_hash_index(k, table_size=table_size) for k in state_keys], dtype=torch.long, ) class TensorExplorationGraph(nn.Module): """Dense tensor representation of the visited ``(state, action)`` graph. Visit statistics live in two ``[table_size, num_actions]`` float tensors (counts and outcome sums); a learned per-state value embedding of shape ``[table_size]`` provides the UCB exploit term / value-prior seam. All hot-path queries are single tensor ops. """ # Class-level annotations make mypy strict happy (register_buffer / # nn.Parameter assignments are otherwise typed as Tensor | Module). visit_counts: Tensor outcome_sums: Tensor state_value: Tensor def __init__( self, *, table_size: int = DEFAULT_TABLE_SIZE, num_actions: int = DEFAULT_NUM_ACTIONS, device: str | torch.device | None = None, dtype: torch.dtype = torch.float32, ) -> None: super().__init__() if table_size <= 0: raise ValueError("table_size must be positive") if num_actions <= 0: raise ValueError("num_actions must be positive") self.table_size = int(table_size) self.num_actions = int(num_actions) self.dtype = dtype # Exploration statistics (NOT learned): buffers move with .to(device). self.register_buffer( "visit_counts", torch.zeros((self.table_size, self.num_actions), dtype=dtype, device=device), ) self.register_buffer( "outcome_sums", torch.zeros((self.table_size, self.num_actions), dtype=dtype, device=device), ) # Learned per-state value embedding (gradients flow). self.state_value = nn.Parameter( torch.zeros(self.table_size, dtype=dtype, device=device) ) # ------------------------------------------------------------------ # device / dtype helpers # ------------------------------------------------------------------ @property def device(self) -> torch.device: return self.visit_counts.device # ------------------------------------------------------------------ # indexing # ------------------------------------------------------------------ def _row(self, state_key: str) -> int: return state_hash_index(state_key, table_size=self.table_size) def _rows(self, state_keys: Sequence[str]) -> Tensor: return state_hash_indices(state_keys, table_size=self.table_size).to(self.device) # ------------------------------------------------------------------ # core mutation (tensor in-place ops; no Python loop over actions) # ------------------------------------------------------------------ def record( self, *, state_key: str, action_index: int, outcome: float, ) -> None: """Record one ``(state, action_index, outcome)`` visit tensorially. ``action_index`` is the integer column in ``[0, num_actions)`` (the tensor module addresses actions by index, unlike the dict-keyed Python reference). Visit counts and outcome sums accumulate in place. """ if not 0 <= action_index < self.num_actions: raise ValueError( f"action_index must be in [0, {self.num_actions}), got {action_index}" ) row = self._row(state_key) # In-place tensor updates on a single cell -- still a tensor op. self.visit_counts[row, action_index] += 1.0 self.outcome_sums[row, action_index] += float(outcome) def record_batch( self, *, state_keys: Sequence[str], action_indices: Tensor, outcomes: Tensor, ) -> None: """Vectorized batch record via :func:`index_add_`. ``action_indices`` and ``outcomes`` are 1-D tensors of length ``N``; ``state_keys`` is the matching length-``N`` sequence of state strings (hashed to row indices). All ``N`` updates happen in one kernel. """ if len(state_keys) != int(action_indices.shape[0]): raise ValueError("state_keys and action_indices length mismatch") if int(action_indices.shape[0]) != int(outcomes.shape[0]): raise ValueError("action_indices and outcomes length mismatch") rows = self._rows(state_keys) actions = action_indices.to(self.device).to(torch.long) outs = outcomes.to(self.device).to(self.dtype) flat_index = rows * self.num_actions + actions ones = torch.ones_like(outs) self.visit_counts.view(-1).index_add_(0, flat_index, ones) self.outcome_sums.view(-1).index_add_(0, flat_index, outs) # ------------------------------------------------------------------ # queries # ------------------------------------------------------------------ def visit_count(self, state_key: str, action_index: int) -> Tensor: """Scalar tensor visit count for one ``(state, action)`` pair.""" return self.visit_counts[self._row(state_key), action_index] def mean_outcome(self, state_key: str, action_index: int) -> Tensor: """Scalar tensor mean outcome (0 where unvisited).""" count = self.visit_count(state_key, action_index) total = self.outcome_sums[self._row(state_key), action_index] return torch.where(count > 0, total / count, torch.zeros_like(total)) def row_visits(self, state_key: str) -> Tensor: """Per-action visit counts ``[num_actions]`` for ``state_key``.""" return self.visit_counts[self._row(state_key)] def row_mean_outcomes(self, state_key: str) -> Tensor: """Per-action mean outcomes ``[num_actions]`` (0 where unvisited).""" row = self._row(state_key) counts = self.visit_counts[row] sums = self.outcome_sums[row] return torch.where(counts > 0, sums / counts, torch.zeros_like(sums)) def total_visits(self, state_key: str) -> Tensor: """Scalar tensor sum of visits across all actions at ``state_key``.""" return self.row_visits(state_key).sum() def frontier_mask(self, state_key: str) -> Tensor: """Boolean ``[num_actions]`` mask: ``True`` where action is untested.""" return self.row_visits(state_key) == 0 def has_frontier(self, state_key: str) -> Tensor: """Scalar boolean tensor: any untested action at ``state_key``?""" return self.frontier_mask(state_key).any() def recommend_next( self, state_key: str, *, available_actions: Tensor | None = None, strategy: str = STRATEGY_SHORTEST_TO_UNTESTED, hypothesis_prior: Tensor | None = None, exploration: float = UCB_DEFAULT_EXPLORATION, ) -> Tensor: """Pick the next action index under the chosen strategy (scalar tensor). ``available_actions`` (optional ``[K]`` long tensor of action indices in ``[0, num_actions)``) restricts selection; if omitted all actions are eligible. Returns the chosen action index as a 0-D long tensor. """ if available_actions is None: available = torch.arange(self.num_actions, device=self.device) else: available = available_actions.to(self.device).to(torch.long) row = self._row(state_key) counts = self.visit_counts[row] sums = self.outcome_sums[row] means = torch.where(counts > 0, sums / counts, torch.zeros_like(sums)) untested = counts == 0 if strategy == STRATEGY_LEAST_SAMPLED: # anti-Thompson: prefer untested first, then fewest-sampled. score = torch.where(untested, torch.full_like(counts, -1.0), counts) chosen = available[torch.argmin(score[available])] return chosen.to(torch.long) if strategy == STRATEGY_BEST_OUTCOME: score = torch.where(untested, torch.full_like(means, torch.finfo(self.dtype).max), means) chosen = available[torch.argmax(score[available])] return chosen.to(torch.long) if strategy == STRATEGY_HYPOTHESIS_DRIVEN: prior = ( torch.zeros(self.num_actions, dtype=means.dtype, device=self.device) if hypothesis_prior is None else hypothesis_prior.to(self.device).to(self.dtype) ) scores = self.ucb_scores_for_row(row, exploration=exploration) + prior scores = torch.where(untested, torch.full_like(scores, torch.finfo(self.dtype).max), scores) chosen = available[torch.argmax(scores[available])] return chosen.to(torch.long) if strategy == STRATEGY_UCB: scores = self.ucb_scores_for_row(row, exploration=exploration) scores = torch.where(untested, torch.full_like(scores, torch.finfo(self.dtype).max), scores) chosen = available[torch.argmax(scores[available])] return chosen.to(torch.long) # default: shortest-to-untested -- if any untested action is available, # take the first; otherwise fall back to least-sampled. avail_untested = untested[available] if avail_untested.any(): chosen = available[torch.argmax(avail_untested.to(torch.long))] return chosen.to(torch.long) score = counts chosen = available[torch.argmin(score[available])] return chosen.to(torch.long) def ucb_scores_for_row( self, row: int, *, exploration: float = UCB_DEFAULT_EXPLORATION, ) -> Tensor: """UCB scores ``[num_actions]`` for one hashed row. Exploit = the learned ``state_value[row]`` broadcast as the per-action mean-outcome baseline (so gradients flow through the value embedding), plus the empirical mean outcome. Explore = the standard ``c * sqrt(log(N) / n_a)`` bonus, zero where ``n_a == 0``. """ counts = self.visit_counts[row] sums = self.outcome_sums[row] means = torch.where(counts > 0, sums / counts, torch.zeros_like(sums)) total = counts.sum() log_total = torch.log(torch.clamp(total, min=1.0)) explore_bonus = exploration * torch.sqrt( log_total / torch.clamp(counts, min=1.0) ) explore_bonus = torch.where(counts > 0, explore_bonus, torch.zeros_like(explore_bonus)) # Exploit term mixes the learned per-state value with the empirical mean. exploit = self.state_value[row] + means return exploit + explore_bonus def ucb_scores( self, state_key: str, *, exploration: float = UCB_DEFAULT_EXPLORATION, ) -> Tensor: """UCB scores ``[num_actions]`` for ``state_key`` (gradient-flowing).""" return self.ucb_scores_for_row(self._row(state_key), exploration=exploration) def softmax_recommend( self, state_key: str, *, temperature: float = 1.0, ) -> Tensor: """Softmax sampling distribution ``[num_actions]`` over UCB scores. Differentiable sampling distribution (caller may ``torch.multinomial`` or take the expectation). Useful as a stochastic exploration policy. """ scores = self.ucb_scores(state_key) return torch.softmax(scores / max(temperature, 1e-6), dim=0) def action_distribution_entropy(self, state_key: str) -> Tensor: """Shannon entropy (nats) of the visit distribution at ``state_key``. High entropy = broadly explored; low = focused. Differentiable through the count tensor (counts are buffers, but the math is tensor-native). """ counts = self.row_visits(state_key) total = counts.sum() probs = counts / torch.clamp(total, min=1.0) log_probs = torch.log(torch.clamp(probs, min=1e-12)) entropy = -(probs * log_probs).sum() return torch.where(total > 0, entropy, torch.zeros_like(entropy)) # ------------------------------------------------------------------ # federation merge # ------------------------------------------------------------------ def merge(self, other: "TensorExplorationGraph") -> "TensorExplorationGraph": """Federation merge: visit counts and outcome sums add (elementwise). Returns a fresh module (does not mutate ``self`` or ``other``). The learned ``state_value`` becomes the mean of the two (or self where other is zero) -- a stable federation average for the value embedding. """ if self.table_size != other.table_size or self.num_actions != other.num_actions: raise ValueError("cannot merge graphs of differing shape") merged = TensorExplorationGraph( table_size=self.table_size, num_actions=self.num_actions, device=self.device, dtype=self.dtype, ) merged.visit_counts = (self.visit_counts + other.visit_counts).clone() merged.outcome_sums = (self.outcome_sums + other.outcome_sums).clone() # Average learned values; fall back to self where other has none. self_has = self.state_value != 0 other_has = other.state_value != 0 both = self_has & other_has summed = self.state_value + other.state_value averaged = torch.where(both, summed / 2.0, self.state_value + other.state_value) with torch.no_grad(): merged.state_value.copy_(averaged) return merged.to(self.device) # ------------------------------------------------------------------ # RHAE efficiency metric (tensor-native, differentiable) # ------------------------------------------------------------------ def relative_human_action_efficiency( self, human_actions: Tensor, ai_actions: Tensor, *, cap: float = RHAE_DEFAULT_CAP, ) -> Tensor: """RHAE = (human / ai) ** 2, clamped to ``cap``. Differentiable. Tensor inputs make this usable as a training signal: gradients flow through ``ai_actions`` (e.g. a soft action count produced by the model). """ raw = (human_actions / ai_actions) ** 2 return torch.clamp(raw, max=cap) def action_efficiency_penalty( self, ai_actions: Tensor, *, reference_actions: Tensor, cap: float = RHAE_DEFAULT_CAP, ) -> Tensor: """Differentiable RHAE-style penalty in ``[0, 1]`` for training signals.""" rhae = self.relative_human_action_efficiency( reference_actions, ai_actions, cap=cap ) return torch.clamp(1.0 - rhae, min=0.0) def ucb_score_tensor( *, mean_outcome: Tensor, visit_count: Tensor, total_visits: Tensor, exploration: float = UCB_DEFAULT_EXPLORATION, ) -> Tensor: """Vectorized UCB score. ``mean_outcome``, ``visit_count`` are ``[...]`` (any shape, broadcastable); ``total_visits`` is scalar or broadcastable. Untested actions (``visit_count <= 0``) score ``+inf`` so they are tried first -- matching :func:`resynthesis.causal_exploration.ucb_score`. """ inf = torch.full_like(mean_outcome, float("inf")) safe_count = torch.clamp(visit_count, min=1.0).to(mean_outcome.dtype) safe_total = torch.clamp(total_visits, min=1.0).to(mean_outcome.dtype) bonus = exploration * torch.sqrt(torch.log(safe_total) / safe_count) score = mean_outcome + bonus return torch.where(visit_count > 0, score, inf) __all__ = [ "CAUSAL_EXPLORATION_TENSOR_SCHEMA", "DEFAULT_NUM_ACTIONS", "DEFAULT_TABLE_SIZE", "RHAE_DEFAULT_CAP", "STRATEGY_BEST_OUTCOME", "STRATEGY_HYPOTHESIS_DRIVEN", "STRATEGY_LEAST_SAMPLED", "STRATEGY_SHORTEST_TO_UNTESTED", "STRATEGY_UCB", "TensorExplorationGraph", "UCB_DEFAULT_EXPLORATION", "state_hash_index", "state_hash_indices", "ucb_score_tensor", ]