"""Tensor-native companion to :mod:`resynthesis.value_function`. Re-expresses ``V(state)`` -- the expected discounted intent-shaped return -- as a dense ``torch.nn.Parameter`` table of shape ``[table_size]`` so the TD(0) and Monte-Carlo updates become single tensor ops, and gradients flow through ``V`` to whatever outer loss consumes it. Layout ------ The value table is indexed by the same SHA-256 -> integer row hash used by the other ``*_tensor.py`` companions (see :mod:`causal_exploration_tensor`): index = int.from_bytes(sha256(state_key).digest()[:8], 'little') % table_size Collisions are acceptable (hash-bucketed). Visit counts are a non-learned buffer so they move with ``.to(device)`` but take no gradients; the value table itself is a learnable :class:`torch.nn.Parameter`. Updates ------- * :meth:`TensorValueFunction.update_td_value` -- the TD(0) rule as a differentiable tensor op: ``V[s] + alpha * (r + gamma * V[s'] - V[s])``. Returns the new ``V[s]`` tensor (caller may use it as a loss target or for the TD error). * :meth:`TensorValueFunction.update_mc_value` -- the running-mean MC rule. * :meth:`TensorValueFunction.update_mc_trajectory_value` -- fully vectorized discounted return-to-go via ``torch.cumprod`` / ``torch.cumsum`` (no Python loops), returning a length-``T`` tensor of per-step MC targets. The original :mod:`resynthesis.value_function` module is preserved as the torch-free reference; this companion is additive and importable independently. """ from __future__ import annotations from collections.abc import Sequence import torch from torch import Tensor, nn from resynthesis.causal_exploration_tensor import state_hash_index, state_hash_indices VALUE_FUNCTION_TENSOR_SCHEMA = "nnf.resynthesis.value_function_tensor.v1" DEFAULT_VALUE_DISCOUNT = 0.9 DEFAULT_TD_ALPHA = 0.1 DEFAULT_TABLE_SIZE = 4096 def discounted_returns(rewards: Tensor, *, discount: float) -> Tensor: """Vectorized discounted return-to-go ``G_t = sum_k gamma^(k-t) r_k``. ``rewards`` is a 1-D ``[T]`` tensor. Returns ``[T]`` where entry ``t`` is the discounted sum of rewards from ``t`` onward. Built via an upper-triangular weight matrix ``M[t, k] = gamma^(k-t)`` (looked up from the discount-power vector by integer offset ``k - t``) for ``k >= t``; the returns are ``M @ rewards`` -- single matmul, fully differentiable, no Python loops. Correct when ``discount == 0`` (only the immediate reward counts at each step). """ if rewards.dim() != 1: raise ValueError("rewards must be 1-D") length = int(rewards.shape[0]) if length == 0: return rewards powers = torch.tensor( [discount ** i for i in range(length)], dtype=rewards.dtype, device=rewards.device, ) idx = torch.arange(length, device=rewards.device, dtype=torch.long) exponent = idx.unsqueeze(0) - idx.unsqueeze(1) # [t, k] lookup = torch.clamp(exponent, min=0) m = powers[lookup] upper = exponent >= 0 m = torch.where(upper, m, torch.zeros_like(m)) return m @ rewards class TensorValueFunction(nn.Module): """Per-state expected discounted intent return as a dense value table. The value table ``V`` is a learnable ``[table_size]`` parameter; visit counts are a buffer (no gradients). All updates and lookups are tensor ops. """ value_table: Tensor visit_counts: Tensor def __init__( self, *, table_size: int = DEFAULT_TABLE_SIZE, discount: float = DEFAULT_VALUE_DISCOUNT, 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") self.table_size = int(table_size) self.discount = float(discount) self.dtype = dtype self.value_table = nn.Parameter( torch.zeros(self.table_size, dtype=dtype, device=device) ) self.register_buffer( "visit_counts", torch.zeros(self.table_size, dtype=dtype, device=device), ) # ------------------------------------------------------------------ # device helpers # ------------------------------------------------------------------ @property def device(self) -> torch.device: return self.value_table.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) # ------------------------------------------------------------------ # evaluation # ------------------------------------------------------------------ def evaluate(self, state_key: str) -> Tensor: """Scalar tensor ``V(state_key)`` (gradient-flowing).""" return self.value_table[self._row(state_key)] def evaluate_many(self, state_keys: Sequence[str]) -> Tensor: """Vectorized ``[N]`` evaluation (one ``index_select`` kernel).""" rows = self._rows(state_keys) return self.value_table[rows] def visits(self, state_key: str) -> Tensor: return self.visit_counts[self._row(state_key)] def best_state(self, state_keys: Sequence[str]) -> int: """Index (into ``state_keys``) of the highest-``V`` candidate. Returns ``-1`` for an empty sequence. Returns an ``int`` (``argmax`` is not differentiable through the index choice itself, but the underlying values are gradient-flowing -- use :meth:`evaluate_many` directly if you need a differentiable reduction). """ if not state_keys: return -1 values = self.evaluate_many(state_keys) return int(torch.argmax(values).item()) # ------------------------------------------------------------------ # TD(0) update (differentiable tensor op) # ------------------------------------------------------------------ def td_error( self, *, reward: float | Tensor, state_key: str, next_state_key: str, discount: float | None = None, ) -> Tensor: gamma = self.discount if discount is None else float(discount) target = torch.as_tensor(reward, dtype=self.dtype, device=self.device) + gamma * self.evaluate(next_state_key) return target - self.evaluate(state_key) def update_td_value( self, *, reward: float | Tensor, state_key: str, next_state_key: str, alpha: float = DEFAULT_TD_ALPHA, discount: float | None = None, ) -> Tensor: """Differentiable TD(0) update op. Returns the new ``V(state_key)`` *as a tensor* (so the caller can form a loss). Note: because ``V`` is a parameter, the actual on-table update is applied via ``index_add_`` on the underlying parameter data inside ``torch.no_grad()`` -- this module is meant to be the *target* of an outer optimizer step, not an in-place learning rule. For the classic in-place rule use :meth:`apply_td_update`. """ error = self.td_error( reward=reward, state_key=state_key, next_state_key=next_state_key, discount=discount, ) return self.evaluate(state_key) + alpha * error def apply_td_update( self, *, reward: float | Tensor, state_key: str, next_state_key: str, alpha: float = DEFAULT_TD_ALPHA, discount: float | None = None, ) -> Tensor: """In-place classic TD(0) update on the parameter; returns TD error. Mirrors :meth:`resynthesis.value_function.ValueFunction.update_td`: ``V[s] += alpha * (r + gamma*V[s'] - V[s])`` and the visit count bumps. Done under ``no_grad`` because this is a tabular learning rule, not a gradient-step target. """ error = self.td_error( reward=reward, state_key=state_key, next_state_key=next_state_key, discount=discount, ) row = self._row(state_key) with torch.no_grad(): self.value_table[row] += alpha * error self.visit_counts[row] += 1.0 return error def apply_td_update_batch( self, *, state_keys: Sequence[str], next_state_keys: Sequence[str], rewards: Tensor, alpha: float = DEFAULT_TD_ALPHA, discount: float | None = None, ) -> Tensor: """Vectorized in-place TD(0) over a batch of transitions. ``rewards`` is ``[N]``; ``state_keys`` / ``next_state_keys`` are length-``N``. Returns the ``[N]`` TD errors. One ``index_add_`` per tensor -- no Python loop over the batch. """ n = len(state_keys) if len(next_state_keys) != n or int(rewards.shape[0]) != n: raise ValueError("state_keys / next_state_keys / rewards length mismatch") gamma = self.discount if discount is None else float(discount) rows = self._rows(state_keys) next_rows = self._rows(next_state_keys) rew = rewards.to(self.device).to(self.dtype) v = self.value_table targets = rew + gamma * v[next_rows] errors = targets - v[rows] with torch.no_grad(): self.value_table.index_add_( 0, rows, alpha * errors ) ones = torch.ones(n, dtype=self.dtype, device=self.device) self.visit_counts.index_add_(0, rows, ones) return errors # ------------------------------------------------------------------ # Monte-Carlo update # ------------------------------------------------------------------ def update_mc_value( self, *, state_key: str, return_value: float | Tensor, ) -> Tensor: """Differentiable running-mean MC update op (returns the new V[s]).""" count = self.visits(state_key) prior = self.evaluate(state_key) return (prior * count + torch.as_tensor(return_value, dtype=self.dtype, device=self.device)) / (count + 1.0) def apply_mc_update( self, *, state_key: str, return_value: float | Tensor, ) -> Tensor: """In-place running-mean MC update; returns the new ``V[state_key]``.""" row = self._row(state_key) count = self.visit_counts[row] prior = self.value_table[row] ret = torch.as_tensor(return_value, dtype=self.dtype, device=self.device) new_value = (prior * count + ret) / (count + 1.0) with torch.no_grad(): self.value_table[row] = new_value self.visit_counts[row] += 1.0 return new_value def update_mc_trajectory_value( self, state_keys: Sequence[str], rewards: Tensor, *, discount: float | None = None, ) -> Tensor: """Differentiable MC return-to-go over a trajectory. Returns the ``[T]`` tensor of per-step MC returns ``G_t = sum_k gamma^(k-t) rewards[k]`` (the *targets* for an outer optimizer step on ``V``). Fully vectorized via :func:`discounted_returns` -- no Python loops. """ if len(state_keys) != int(rewards.shape[0]): raise ValueError("state_keys and rewards length mismatch") gamma = self.discount if discount is None else float(discount) rewards = rewards.to(self.device).to(self.dtype) return discounted_returns(rewards, discount=gamma) def apply_mc_trajectory( self, state_keys: Sequence[str], rewards: Tensor, *, discount: float | None = None, ) -> None: """In-place MC running-mean update for every state on the trajectory.""" targets = self.update_mc_trajectory_value( state_keys, rewards, discount=discount ) rows = self._rows(state_keys) counts = self.visit_counts[rows] priors = self.value_table[rows] new_values = (priors * counts + targets) / (counts + 1.0) with torch.no_grad(): self.value_table[rows] = new_values self.visit_counts[rows] += 1.0 # ------------------------------------------------------------------ # Endstate-as-target framing (tensor-native, differentiable) # ------------------------------------------------------------------ def value_toward_endstate( self, state_key: str, is_endstate: bool | Tensor, *, terminal_value: float = 1.0, ) -> Tensor: """Value expressed as progress toward the merit endstate. ``is_endstate`` true -> ``terminal_value`` (the endstate is the goal, not a rung); otherwise the learned ``V`` clamped into ``[0, terminal_value]`` so it reads as a fraction of the way to the goal. Differentiable through ``V`` (the clamp is a soft-ish gate via ``torch.clamp``). """ if terminal_value < 0.0: raise ValueError("terminal_value must be non-negative") raw = self.evaluate(state_key) is_end = torch.as_tensor(is_endstate, dtype=self.dtype, device=self.device) terminal = torch.as_tensor(terminal_value, dtype=self.dtype, device=self.device) clamped = torch.clamp(raw, min=0.0, max=terminal_value) if terminal_value > 0.0 else torch.zeros_like(raw) return torch.where(is_end > 0, terminal, clamped) def goal_progress( self, state_key: str, *, is_endstate: bool | Tensor = False, ) -> Tensor: """Shorthand for ``value_toward_endstate(..., terminal_value=1.0)``.""" return self.value_toward_endstate(state_key, is_endstate, terminal_value=1.0) __all__ = [ "DEFAULT_TABLE_SIZE", "DEFAULT_TD_ALPHA", "DEFAULT_VALUE_DISCOUNT", "TensorValueFunction", "VALUE_FUNCTION_TENSOR_SCHEMA", "discounted_returns", ]