runtime/src/resynthesis/causal_exploration_tensor.py

"""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",
]