DDPG & TD3

"""DDPG - the deterministic policy gradient, transcribed.

Two networks and one chain rule. The critic is fitted by Bellman regression.
The actor is trained by pushing its output uphill on the critic's surface,
which means dQ/da is computed explicitly and then chained into the actor.

Watch what the actor's objective IS: there is no target in it and no quantity
derived from data. It says "make the critic rate this action higher". That is
an optimizer pointed at a learned function, and it is the whole difficulty.
"""

import math
import random


def make_critic(n_state, n_action, n_hidden):
    """Q(s, a) over the concatenation [s; a]. It must be differentiable in a -
    that requirement is why this branch exists only for continuous control."""
    scale = 1.0 / math.sqrt(n_state + n_action)
    return {
        "w1": [[random.uniform(-scale, scale) for _ in range(n_state + n_action)]
               for _ in range(n_hidden)],
        "b1": [0.0] * n_hidden,
        "w2": [random.uniform(-scale, scale) for _ in range(n_hidden)],
        "b2": 0.0,
    }


def critic_forward(critic, state, action):
    joint = list(state) + list(action)

    hidden = []
    for row, bias in zip(critic["w1"], critic["b1"]):
        total = bias
        for weight, feature in zip(row, joint):
            total += weight * feature
        hidden.append(total if total > 0.0 else 0.0)        # ReLU

    q = critic["b2"]
    for weight, activation in zip(critic["w2"], hidden):
        q += weight * activation

    return q, hidden, joint


def critic_backward(critic, joint, hidden, td_error, lr, n_state):
    """One SGD step on the squared TD error, and dQ/da on the way out.

    The action gradient is the piece the actor needs, and it costs nothing
    extra: it is the same backward pass, read off at the input layer instead of
    stopping at the weights.
    """
    grad_q = -td_error

    grad_hidden = [0.0] * len(hidden)
    for index in range(len(hidden)):
        grad_hidden[index] = grad_q * critic["w2"][index]
        critic["w2"][index] -= lr * grad_q * hidden[index]
    critic["b2"] -= lr * grad_q

    # dQ/d(input), accumulated across hidden units, then sliced to the action
    # half of the concatenation.
    grad_input = [0.0] * len(joint)
    for unit in range(len(hidden)):
        if hidden[unit] <= 0.0:
            continue                                        # ReLU is closed here
        delta = grad_hidden[unit]
        for index in range(len(joint)):
            grad_input[index] += delta * critic["w1"][unit][index]
            critic["w1"][unit][index] -= lr * delta * joint[index]
        critic["b1"][unit] -= lr * delta

    return grad_input[n_state:]


def actor_forward(actor, state, action_limit):
    """A DETERMINISTIC policy: one action, not a distribution.

    tanh bounds the output to the actuator range. It is also a real failure
    mode - a saturated tanh has a near-zero derivative, so an actor driven to
    the bound stops learning while the critic keeps training.
    """
    pre_activation = []
    for row in actor["w"]:
        total = 0.0
        for weight, feature in zip(row, state):
            total += weight * feature
        pre_activation.append(total)

    action = [action_limit * math.tanh(z) for z in pre_activation]
    return action, pre_activation


def actor_backward(actor, state, pre_activation, dq_daction, lr, action_limit):
    """Chain dQ/da into the actor. ASCENT: the actor maximizes the critic.

    da_k/dw[k][i] = action_limit * (1 - tanh(z_k)^2) * s[i], so the whole
    update is dQ/da_k times that. Two networks joined by one chain rule, and
    the actor inherits every error in the surface it is climbing.
    """
    for index in range(len(pre_activation)):
        squashed = math.tanh(pre_activation[index])
        local = action_limit * (1.0 - squashed * squashed)
        coefficient = dq_daction[index] * local
        for feature_index, feature in enumerate(state):
            actor["w"][index][feature_index] += lr * coefficient * feature


def soft_update(target, online, tau):
    """Polyak averaging: the target drifts continuously rather than jumping.

    DQN copies its target network periodically. Here the actor is moving too,
    so a smoothly drifting target suits the setting better than a step change.
    """
    for key in target:
        if isinstance(target[key], list) and target[key] and isinstance(target[key][0], list):
            for row_index in range(len(target[key])):
                for index in range(len(target[key][row_index])):
                    target[key][row_index][index] = (
                        tau * online[key][row_index][index]
                        + (1.0 - tau) * target[key][row_index][index]
                    )
        elif isinstance(target[key], list):
            for index in range(len(target[key])):
                target[key][index] = tau * online[key][index] + (1.0 - tau) * target[key][index]
        else:
            target[key] = tau * online[key] + (1.0 - tau) * target[key]


def train(env, n_state, n_action, steps=100_000, n_hidden=64, lr=1e-3, gamma=0.99,
          tau=0.005, action_limit=1.0, noise_scale=0.1, buffer_size=50_000,
          warmup=1_000, batch_size=32):
    scale = 1.0 / math.sqrt(n_state)
    actor = {"w": [[random.uniform(-scale, scale) for _ in range(n_state)]
                   for _ in range(n_action)]}
    critic = make_critic(n_state, n_action, n_hidden)

    target_actor = {"w": [row[:] for row in actor["w"]]}
    target_critic = {
        "w1": [row[:] for row in critic["w1"]],
        "b1": critic["b1"][:],
        "w2": critic["w2"][:],
        "b2": critic["b2"],
    }

    buffer = []
    state = env.reset()

    for step in range(steps):
        # The policy is deterministic, so it supplies NO exploration. Noise has
        # to be added from outside, and its scale is in action units - which is
        # why this hyperparameter does not transfer between environments.
        if step < warmup:
            action = [random.uniform(-action_limit, action_limit) for _ in range(n_action)]
        else:
            action, _ = actor_forward(actor, state, action_limit)
            action = [
                max(-action_limit, min(action_limit, a + random.gauss(0.0, noise_scale)))
                for a in action
            ]

        next_state, reward, done = env.step(action)
        buffer.append((state, action, reward, next_state, done))
        if len(buffer) > buffer_size:
            buffer.pop(0)

        state = env.reset() if done else next_state
        if len(buffer) < warmup:
            continue

        for _ in range(batch_size):
            s, a, r, s_next, terminal = random.choice(buffer)

            # The target: bootstrap through the TARGET actor and TARGET critic,
            # both drifting slowly, so the regression target is approximately
            # stationary while the online networks move.
            if terminal:
                bootstrap = 0.0                             # no future past a terminal
            else:
                next_action, _ = actor_forward(target_actor, s_next, action_limit)
                bootstrap, _, _ = critic_forward(target_critic, s_next, next_action)

            q, hidden, joint = critic_forward(critic, s, a)
            td_error = r + gamma * bootstrap - q

            # One backward pass serves both networks: it fits the critic and
            # hands back dQ/da for the actor.
            dq_daction = critic_backward(critic, joint, hidden, td_error, lr, n_state)

            # The actor is evaluated at ITS OWN current output, not at the
            # stored action - the objective is "what would the critic say about
            # what I would do now", which is what makes this off-policy.
            current_action, pre_activation = actor_forward(actor, s, action_limit)
            _, actor_hidden, actor_joint = critic_forward(critic, s, current_action)
            dq_dcurrent = critic_backward(critic, actor_joint, actor_hidden, 0.0, 0.0, n_state)

            actor_backward(actor, s, pre_activation, dq_dcurrent, lr, action_limit)

            soft_update(target_critic, critic, tau)
            soft_update(target_actor, actor, tau)

    return actor, critic

"""TD3 - twin critics, delayed actor, smoothed targets, typed throughout."""

from __future__ import annotations

from dataclasses import dataclass
from typing import NamedTuple

import numpy as np
import torch
from numpy.typing import NDArray
from torch import Tensor, nn


class Batch(NamedTuple):
    states: Tensor
    actions: Tensor
    rewards: Tensor
    next_states: Tensor
    # Terminal, not done: a time-limit truncation still has a future and must
    # still be bootstrapped. Conflating them teaches the agent the world ends
    # at the episode limit.
    terminals: Tensor


@dataclass(frozen=True)
class Td3Config:
    n_state: int
    n_action: int
    action_limit: float
    hidden: tuple[int, ...] = (256, 256)
    learning_rate: float = 3e-4
    gamma: float = 0.99
    tau: float = 0.005
    # Exploration, in ACTION UNITS. The policy is deterministic and supplies
    # none of its own, so this is the entire mechanism - and it does not
    # transfer between environments with different action scales.
    exploration_noise: float = 0.1
    # Target policy smoothing: noise on the target action, clipped so it stays
    # local. It asserts that nearby actions have similar values, which is what
    # makes a critic artifact narrower than the noise unexploitable.
    target_noise: float = 0.2
    target_noise_clip: float = 0.5
    # One actor step per two critic steps, so the actor climbs a surface that
    # has had time to settle between its own moves.
    policy_delay: int = 2
    warmup_steps: int = 10_000

    def __post_init__(self) -> None:
        if self.action_limit <= 0.0:
            raise ValueError("action limit must be positive")
        if not 0.0 < self.tau <= 1.0:
            raise ValueError(f"tau must be in (0, 1], got {self.tau}")
        if self.policy_delay < 1:
            raise ValueError("policy delay must be at least one")
        if self.target_noise_clip < self.target_noise:
            raise ValueError("clipping below the noise scale makes the smoothing one-sided")


class Actor(nn.Module):
    """Deterministic: one action, not a distribution."""

    def __init__(self, config: Td3Config) -> None:
        super().__init__()
        layers: list[nn.Module] = []
        width = config.n_state
        for size in config.hidden:
            layers += [nn.Linear(width, size), nn.ReLU()]
            width = size
        layers += [nn.Linear(width, config.n_action), nn.Tanh()]
        self.net = nn.Sequential(*layers)
        self.action_limit = config.action_limit

    def forward(self, states: Tensor) -> Tensor:
        return self.action_limit * self.net(states)


class TwinCritic(nn.Module):
    """Two independently initialized critics over the same (s, a) input.

    Independent initialization is the point: their errors have to be
    uncorrelated for the min to remove anything.
    """

    def __init__(self, config: Td3Config) -> None:
        super().__init__()
        self.q1 = self._build(config)
        self.q2 = self._build(config)

    @staticmethod
    def _build(config: Td3Config) -> nn.Module:
        layers: list[nn.Module] = []
        width = config.n_state + config.n_action
        for size in config.hidden:
            layers += [nn.Linear(width, size), nn.ReLU()]
            width = size
        layers.append(nn.Linear(width, 1))
        return nn.Sequential(*layers)

    def forward(self, states: Tensor, actions: Tensor) -> tuple[Tensor, Tensor]:
        joint = torch.cat([states, actions], dim=1)
        return self.q1(joint).squeeze(-1), self.q2(joint).squeeze(-1)


class ReplayBuffer:
    """Fixed-capacity ring over pre-allocated arrays, one per field."""

    def __init__(self, capacity: int, n_state: int, n_action: int) -> None:
        if capacity < 1:
            raise ValueError("capacity must be positive")
        self._states = np.zeros((capacity, n_state), dtype=np.float32)
        self._next_states = np.zeros((capacity, n_state), dtype=np.float32)
        self._actions = np.zeros((capacity, n_action), dtype=np.float32)
        self._rewards = np.zeros(capacity, dtype=np.float32)
        self._terminals = np.zeros(capacity, dtype=np.bool_)
        self._capacity = capacity
        self._cursor = 0
        self._size = 0

    def __len__(self) -> int:
        return self._size

    def add(self, state: NDArray[np.float32], action: NDArray[np.float32],
            reward: float, next_state: NDArray[np.float32], terminal: bool) -> None:
        index = self._cursor
        self._states[index] = state
        self._actions[index] = action
        self._rewards[index] = reward
        self._next_states[index] = next_state
        self._terminals[index] = terminal
        self._cursor = (self._cursor + 1) % self._capacity
        self._size = min(self._size + 1, self._capacity)

    def sample(self, batch_size: int, rng: np.random.Generator) -> Batch:
        if self._size < batch_size:
            raise ValueError(f"buffer holds {self._size}, cannot sample {batch_size}")
        index = rng.integers(0, self._size, size=batch_size)
        return Batch(
            torch.from_numpy(self._states[index]),
            torch.from_numpy(self._actions[index]),
            torch.from_numpy(self._rewards[index]),
            torch.from_numpy(self._next_states[index]),
            torch.from_numpy(self._terminals[index]),
        )


@torch.no_grad()
def smoothed_pessimistic_target(
    target_actor: Actor, target_critic: TwinCritic, batch: Batch, config: Td3Config
) -> Tensor:
    """The TD3 target: smoothed action, minimum over two critics.

    The min is a deliberate change of BIAS DIRECTION, not a reduction in its
    magnitude. Underestimating an action's value makes the actor avoid it,
    which costs a bounded amount of performance. Overestimating makes the actor
    SEEK it, which compounds - because the actor is a maximizer and a spurious
    peak is a maximum.
    """
    noise = (torch.randn_like(batch.actions) * config.target_noise).clamp(
        -config.target_noise_clip, config.target_noise_clip
    )
    next_actions = (target_actor(batch.next_states) + noise).clamp(
        -config.action_limit, config.action_limit
    )

    q1, q2 = target_critic(batch.next_states, next_actions)
    next_value = torch.min(q1, q2).masked_fill(batch.terminals, 0.0)
    return batch.rewards + config.gamma * next_value


def update(
    actor: Actor, critic: TwinCritic,
    target_actor: Actor, target_critic: TwinCritic,
    actor_optimizer: torch.optim.Optimizer, critic_optimizer: torch.optim.Optimizer,
    batch: Batch, config: Td3Config, step: int,
) -> dict[str, float]:
    target = smoothed_pessimistic_target(target_actor, target_critic, batch, config)
    q1, q2 = critic(batch.states, batch.actions)
    critic_loss = nn.functional.mse_loss(q1, target) + nn.functional.mse_loss(q2, target)

    critic_optimizer.zero_grad(set_to_none=True)
    critic_loss.backward()
    critic_optimizer.step()

    stats = {"critic_loss": float(critic_loss), "mean_q": float(q1.mean())}

    # Delayed: the actor moves only every policy_delay critic updates, so it is
    # climbing a surface that has settled rather than chasing one that moves
    # underneath it every step.
    if step % config.policy_delay != 0:
        return stats

    # The actor maximizes q1 alone, not the min. Using the min here would make
    # the actor climb the more pessimistic of two surfaces for no benefit; the
    # pessimism belongs in the TARGET, where the compounding happens.
    actor_loss = -critic.q1(
        torch.cat([batch.states, actor(batch.states)], dim=1)
    ).mean()

    actor_optimizer.zero_grad(set_to_none=True)
    actor_loss.backward()
    actor_optimizer.step()

    with torch.no_grad():
        for online, target_parameter in zip(critic.parameters(), target_critic.parameters()):
            target_parameter.mul_(1 - config.tau).add_(online, alpha=config.tau)
        for online, target_parameter in zip(actor.parameters(), target_actor.parameters()):
            target_parameter.mul_(1 - config.tau).add_(online, alpha=config.tau)

    stats["actor_loss"] = float(actor_loss)
    return stats

"""TD3 - fused twin critics, foreach Polyak, actor work skipped when delayed.

None of this changes the algorithm. All of it removes work that a direct
implementation does per step and that adds up at a million environment steps.
"""

from __future__ import annotations

import numpy as np
import torch
from torch import Tensor, nn


class FusedTwinCritic(nn.Module):
    """Both critics as one batched matrix multiply.

    Two separate networks mean two GEMMs per layer, each half the size and each
    too small to saturate anything. Stacking them into a single weight tensor
    with a leading pair dimension turns every layer into one baddbmm - identical
    arithmetic, half the kernel launches, and the pair dimension keeps the two
    sets of parameters genuinely independent, which is what the min requires.
    """

    def __init__(self, n_input: int, hidden: int) -> None:
        super().__init__()
        # Leading dim 2: index 0 is critic one, index 1 is critic two.
        self.w1 = nn.Parameter(torch.empty(2, n_input, hidden))
        self.b1 = nn.Parameter(torch.zeros(2, 1, hidden))
        self.w2 = nn.Parameter(torch.empty(2, hidden, 1))
        self.b2 = nn.Parameter(torch.zeros(2, 1, 1))
        for tensor in (self.w1, self.w2):
            nn.init.orthogonal_(tensor, gain=1.0)

    def forward(self, states: Tensor, actions: Tensor) -> Tensor:
        joint = torch.cat([states, actions], dim=1)
        # (2, batch, n_input) without copying: expand is a stride trick.
        batched = joint.unsqueeze(0).expand(2, -1, -1)

        hidden = torch.baddbmm(self.b1, batched, self.w1).relu_()   # ReLU in place
        return torch.baddbmm(self.b2, hidden, self.w2).squeeze(-1)  # (2, batch)


class ContiguousReplay:
    """One flat allocation per field; a batch is one fancy-index per field.

    A list of per-transition tuples allocates twice per environment step and
    scatters every sample, which at a million steps is the largest avoidable
    cost in the whole loop.
    """

    def __init__(self, capacity: int, n_state: int, n_action: int, device: torch.device) -> None:
        # Single dtype throughout, so a sampled batch needs no conversion and
        # transfers as one contiguous block per field.
        self.states = torch.zeros((capacity, n_state), dtype=torch.float32, device=device)
        self.next_states = torch.zeros((capacity, n_state), dtype=torch.float32, device=device)
        self.actions = torch.zeros((capacity, n_action), dtype=torch.float32, device=device)
        self.rewards = torch.zeros(capacity, dtype=torch.float32, device=device)
        self.terminals = torch.zeros(capacity, dtype=torch.bool, device=device)
        self.capacity = capacity
        self.cursor = 0
        self.size = 0

    def add(self, state: Tensor, action: Tensor, reward: float,
            next_state: Tensor, terminal: bool) -> None:
        index = self.cursor
        # copy_ writes into storage that already exists; assignment would not.
        self.states[index].copy_(state)
        self.actions[index].copy_(action)
        self.next_states[index].copy_(next_state)
        self.rewards[index] = reward
        self.terminals[index] = terminal
        self.cursor = (self.cursor + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)

    def sample(self, batch_size: int, generator: torch.Generator):
        index = torch.randint(0, self.size, (batch_size,), generator=generator,
                              device=self.states.device)
        return (
            self.states[index], self.actions[index], self.rewards[index],
            self.next_states[index], self.terminals[index],
        )


@torch.no_grad()
def polyak_foreach(target_params: list[Tensor], online_params: list[Tensor], tau: float) -> None:
    """Soft-update every parameter tensor in two fused calls.

    A Python loop with mul_ and add_ per tensor runs hundreds of tiny kernel
    launches on EVERY step - and unlike DQN's periodic copy, this happens
    continuously. The foreach ops batch the whole parameter list into one
    launch each, which is a real fraction of wall-clock in a direct TD3.
    """
    torch._foreach_mul_(target_params, 1.0 - tau)
    torch._foreach_add_(target_params, online_params, alpha=tau)


@torch.no_grad()
def target_value(
    target_actor: nn.Module, target_critic: FusedTwinCritic,
    next_states: Tensor, rewards: Tensor, terminals: Tensor,
    gamma: float, action_limit: float, noise: float, noise_clip: float,
) -> Tensor:
    """Smoothed pessimistic target, with the min taken across the pair axis.

    Because both critics live in one tensor, the minimum is a reduction over
    dimension 0 rather than a comparison between two separate results - no
    second forward pass and no intermediate tensors to materialize.
    """
    smoothing = (torch.randn_like(rewards.unsqueeze(-1)) * noise).clamp(-noise_clip, noise_clip)
    next_actions = (target_actor(next_states) + smoothing).clamp(-action_limit, action_limit)

    pair = target_critic(next_states, next_actions)          # (2, batch)
    next_value = pair.min(dim=0).values.masked_fill(terminals, 0.0)
    return rewards + gamma * next_value


def critic_step(critic: FusedTwinCritic, optimizer: torch.optim.Optimizer,
                states: Tensor, actions: Tensor, target: Tensor) -> Tensor:
    """One loss covering both critics: the pair axis broadcasts against the
    shared target, so there is no second loss term to add."""
    pair = critic(states, actions)                           # (2, batch)
    loss = nn.functional.mse_loss(pair, target.unsqueeze(0).expand_as(pair))

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    return loss.detach()