| """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) |
| 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), |
| ) |
|
|
| |
| |
| |
|
|
| @property |
| def device(self) -> torch.device: |
| return self.value_table.device |
|
|
| |
| |
| |
|
|
| 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) |
|
|
| |
| |
| |
|
|
| 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()) |
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| |
| |
|
|
| 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", |
| ] |
|
|
