Blame for lecture_7/branched.py

"""Two bridges that change the shape of the problem: branching, and interacting particles.

A population leaving one source does not always end in one place, and Branched Schrodinger Bridge
Matching gives each outcome its own controlled diffusion and its own share of the mass, with a
learned growth rate moving mass between branches. Entangled Schrodinger Bridge Matching keeps one
population and changes the reference, taking the molecular dynamics of the system itself and
tilting it at the final time by a target density. A bias force on every particle reproduces that
tilt.

Chapter 7, Sections 7.26 to 7.27.3 of the notes. Runs the two-branch example in closed form and
trains a bias force on two interacting particles in the plane. Prints the weighted energies of the
two branch schedules, the interpolant and its energy, the reweighting arithmetic of the two-outcome
example, one Euler-Maruyama step, and the reward the trained force reaches.

    python lecture_7/branched.py              # 100 optimizer steps on the bias force
    python lecture_7/branched.py --steps 20   # a few seconds, for a first look
"""

# %% 01. Imports and configuration
from __future__ import annotations

import math

import torch
import torch.nn as nn

from cis6270.nets import TimeEmbedding, count_parameters
from cis6270.runner import (
    common_parser, pick_device, print_report, save_run, set_seed, train,
)

PARTICLES = 2  # n, the number of interacting particles.
DIMENSION = 2  # Each particle lives in the plane, so the projector below has something to remove.
MASS, FRICTION, TEMPERATURE = 1.0, 1.0, 1.0  # m_i, gamma and k_B tau, all in reduced units.
NOISE = math.sqrt(2.0 * FRICTION * TEMPERATURE / MASS)  # The noise scale of the reference.
TARGET_WIDTH = 0.4  # The width of the target density around the region B.
PATHS = 64  # Trajectories simulated per optimizer step.
SIM_STEPS = 30  # Euler-Maruyama steps per trajectory.


# %% 02. A population that splits
#
# Branch k is a controlled diffusion carrying the fraction w^k_t of the population, and the weights
# change at a learned growth rate with sum_k g^k_t = 0, so a branch grows only by taking mass from
# another. Section 7.26.


def branch_weights(growth, t: torch.Tensor) -> torch.Tensor:
    """w^1_t = integral of the growth rate from 0 to t, with branch 0 holding the remainder."""
    return growth(t)  # The weight of branch 1; branch 0 carries 1 - w^1_t.


def weighted_energy(drift: float, weight, points: int = 2001) -> float:
    """The time integral of (1/2) ||u||^2 times the weight of the branch, by the trapezoid rule."""
    t = torch.linspace(0.0, 1.0, points)
    density = 0.5 * drift ** 2 * weight(t)  # The control energy priced by the mass that spends it.
    return float(torch.trapezoid(density, t))


def branch_sample(weight_one: float, generator: torch.Generator) -> int:
    """Draw a branch with probability w^k_t, the way we sample the population."""
    return int(torch.rand(1, generator=generator) < weight_one)  # 1 for branch 1, 0 for branch 0.


# %% 03. The interpolant
#
# A network correction multiplied by t(1 - t) leaves both endpoints fixed for every value of the
# parameters, so the endpoint constraints hold during training with no penalty term. With no state
# cost the energy of a bent path exceeds the energy of the straight line, and the correction goes
# to zero in training. Section 7.26.2.


def interpolant(x0: float, x1: float, correction: float, t: torch.Tensor) -> torch.Tensor:
    """x_t = (1 - t) x_0 + t x_1 + t(1 - t) phi, with phi here a constant."""
    return (1.0 - t) * x0 + t * x1 + t * (1.0 - t) * correction


def interpolant_velocity(x0: float, x1: float, correction: float, t: torch.Tensor) -> torch.Tensor:
    """x_dot = (x_1 - x_0) + (1 - 2t) phi, the derivative of the interpolant in t."""
    return (x1 - x0) + (1.0 - 2.0 * t) * correction


def interpolant_energy(x0: float, x1: float, correction: float, points: int = 2001) -> float:
    """The integral of one half the squared velocity of the interpolant over the unit interval."""
    t = torch.linspace(0.0, 1.0, points)
    return float(torch.trapezoid(0.5 * interpolant_velocity(x0, x1, correction, t) ** 2, t))


# %% 04. The reference of Entangled Schrodinger Bridge Matching
#
# The reference is the dynamics the system follows, underdamped Langevin with the physical force,
# the friction and the thermal noise. Section 7.27.


def potential(positions: torch.Tensor) -> torch.Tensor:
    """U(R), a double well on each particle plus a spring that couples the two. [B, n, d] -> [B]."""
    radial = ((positions ** 2).sum(dim=-1) - 1.0) ** 2  # [B, n] a well at every unit circle.
    coupling = 0.5 * ((positions[:, 0] - positions[:, 1]) ** 2).sum(dim=-1)  # [B] the spring.
    return radial.sum(dim=-1) + coupling  # [B].


def physical_force(positions: torch.Tensor) -> torch.Tensor:
    """-grad_R U(R), differentiated by hand so the simulation runs without autograd. [B, n, d]."""
    radial = -4.0 * ((positions ** 2).sum(dim=-1, keepdim=True) - 1.0) * positions  # [B, n, d].
    separation = positions[:, 0] - positions[:, 1]  # [B, d] the displacement across the spring.
    spring = torch.stack([-separation, separation], dim=1)  # [B, n, d] equal and opposite.
    return radial + spring


def force_check(positions: torch.Tensor) -> float:
    """The largest gap between the force written by hand and the gradient autograd computes."""
    r = positions.detach().requires_grad_(True)
    gradient = torch.autograd.grad(potential(r).sum(), r)[0]
    return float((physical_force(positions) + gradient).abs().max())


TARGET = torch.tensor([1.0, 0.0])  # The centre of the region B, the same for both particles.


def log_target(positions: torch.Tensor) -> torch.Tensor:
    """r(X_1) = log pi_B(R_1), zero inside the region and negative outside. [B, n, d] -> [B]."""
    offset = positions - TARGET.to(positions)  # [B, n, d] displacement from the target.
    return -(offset ** 2).sum(dim=(1, 2)) / (2.0 * TARGET_WIDTH ** 2)  # [B].


def target_direction(positions: torch.Tensor) -> torch.Tensor:
    """s_hat_i, the unit vector along the gradient of log pi_B for particle i. [B, n, d]."""
    gradient = -(positions - TARGET.to(positions)) / TARGET_WIDTH ** 2  # [B, n, d].
    return gradient / gradient.norm(dim=-1, keepdim=True).clamp(min=1e-8)


# %% 05. The bias force and its parameterization
#
# We take a nonnegative scalar along the direction of steepest increase of the target density, plus
# a second output with its component along that direction removed, so the push toward the target is
# never negative and we read it off directly as alpha. Section 7.27.3.


def compose_bias(alpha: torch.Tensor, h: torch.Tensor, s_hat: torch.Tensor) -> torch.Tensor:
    """b = alpha s_hat + (I - s_hat s_hat^T) h, the parameterization of the force on each particle."""
    along = (h * s_hat).sum(dim=-1, keepdim=True) * s_hat  # [B, n, d] the removed component.
    return alpha.unsqueeze(-1) * s_hat + (h - along)  # [B, n, d].


class BiasForce(nn.Module):
    """A network from the positions, the velocities and the time to one force per particle."""

    def __init__(self, width: int = 64):
        super().__init__()
        self.time_embedding = TimeEmbedding(width, scale=1.0)
        self.input = nn.Linear(2 * PARTICLES * DIMENSION, width)  # Positions and velocities.
        self.body = nn.Sequential(nn.Linear(width, width), nn.SiLU(),
                                  nn.Linear(width, width), nn.SiLU())
        self.head = nn.Linear(width, PARTICLES * (1 + DIMENSION))  # alpha and h per particle.

    def forward(self, positions: torch.Tensor, velocities: torch.Tensor,
                t: torch.Tensor) -> torch.Tensor:
        """positions, velocities: [B, n, d], t: [B] -> the bias force [B, n, d]."""
        state = torch.cat([positions.flatten(1), velocities.flatten(1)], dim=1)  # [B, 2nd].
        hidden = self.body(self.input(state) + self.time_embedding(t))  # [B, width].
        output = self.head(hidden).view(-1, PARTICLES, 1 + DIMENSION)  # [B, n, 1 + d].
        alpha = nn.functional.softplus(output[..., 0])  # [B, n] nonnegative, so b never points away.
        return compose_bias(alpha, output[..., 1:], target_direction(positions))


# %% 06. Simulating the controlled dynamics
#
# The control enters the velocity equation divided by the mass, against the noise scale, so in units
# of the noise it is b / (m sigma) and Girsanov prices it at one half its squared size.
# Sections 7.27.1 and 7.27.2.


@torch.no_grad()
def roll_out(force: BiasForce | None, paths: int, steps: int, device: torch.device,
             generator: torch.Generator) -> dict:
    """Run the Langevin dynamics with the given bias force, storing the states and the increments."""
    dt = 1.0 / steps
    positions = torch.full((paths, PARTICLES, DIMENSION), -1.0, device=device)
    positions[:, :, 1] = 0.0  # Both particles start at the far well.
    velocities = torch.randn(paths, PARTICLES, DIMENSION, generator=generator).to(device) * \
        math.sqrt(TEMPERATURE / MASS)  # Maxwell-Boltzmann velocities at the reference temperature.
    stored_positions, stored_velocities, stored_u, increments = [], [], [], []
    energy = torch.zeros(paths, device=device)  # (1/2) int ||u||^2 dt along each path.
    cross = torch.zeros(paths, device=device)  # int u . dW along each path.
    for step in range(steps):
        t = torch.full((paths,), step * dt, device=device)
        u = (force(positions, velocities, t) / (MASS * NOISE)) if force is not None else \
            torch.zeros_like(positions)  # The control in units of the noise.
        noise = torch.randn(paths, PARTICLES, DIMENSION,
                            generator=generator).to(device) * math.sqrt(dt)
        stored_positions.append(positions.clone())
        stored_velocities.append(velocities.clone())
        stored_u.append(u.clone())
        increments.append(noise.clone())
        energy = energy + 0.5 * (u ** 2).sum(dim=(1, 2)) * dt
        cross = cross + (u * noise).sum(dim=(1, 2))
        acceleration = physical_force(positions) / MASS - FRICTION * velocities  # The reference.
        velocities = velocities + (acceleration + NOISE * u) * dt + NOISE * noise
        positions = positions + velocities * dt  # The position follows the velocity.
    return {
        "positions": torch.stack(stored_positions), "velocities": torch.stack(stored_velocities),
        "u_bar": torch.stack(stored_u), "noise": torch.stack(increments),
        "reward": log_target(positions), "energy": energy, "cross": cross, "dt": dt,
    }


def importance_weights(rollout: dict) -> torch.Tensor:
    """softmax over the batch of r(X_1) + log dR / dP^bbar, and the unknown Z cancels."""
    log_weight = rollout["reward"] - rollout["cross"] - rollout["energy"]
    return torch.softmax(log_weight, dim=0)  # [P] normalized weights, summing to one.


def cross_entropy_loss(force: BiasForce, rollout: dict, weights: torch.Tensor) -> torch.Tensor:
    """The weighted average of -log dP^{b_theta} / dR along the stored paths, the loss of 7.27.3."""
    steps, paths = rollout["u_bar"].shape[0], rollout["u_bar"].shape[1]
    dt = rollout["dt"]
    clock = torch.arange(steps, device=weights.device).to(torch.get_default_dtype()) * dt
    times = clock[:, None].expand(steps, paths)  # [T, P] the time of every stored state.
    u = force(rollout["positions"].flatten(0, 1), rollout["velocities"].flatten(0, 1),
              times.flatten()).view_as(rollout["u_bar"]) / (MASS * NOISE)  # [T, P, n, d].
    drift = (u * (rollout["u_bar"] * dt + rollout["noise"])).sum(dim=(0, 2, 3))  # [P].
    energy = 0.5 * (u ** 2).sum(dim=(0, 2, 3)) * dt  # [P] the control energy of each path.
    return (weights * (energy - drift)).sum()  # Weighted by the importance weights.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=100, lr=3e-3, width=64, samples=256)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    loud = not args.quiet
    generator = torch.Generator().manual_seed(args.seed)

    linear = lambda t: 0.3 * t  # The growth rate 0.3, so w^1_t rises at a constant rate.
    late = lambda t: 0.3 * t ** 2  # The growth rate 0.6 t, slow at first.
    near = {"linear": weighted_energy(2.0, lambda t: 1.0 - linear(t))
            + weighted_energy(-2.0, linear)}
    far = {"linear": weighted_energy(2.0, lambda t: 1.0 - linear(t)) + weighted_energy(-4.0, linear),
           "late": weighted_energy(2.0, lambda t: 1.0 - late(t)) + weighted_energy(-4.0, late)}
    half = float(branch_weights(linear, torch.tensor(0.5)))
    if loud:
        print("Two branches from the point 0, ending at +2 with weight 0.7 and at -2 with 0.3")
        print(f"  weights at t = 1/2: {1 - half:.2f} and {half:.2f}; at t = 1: "
              f"{1 - float(branch_weights(linear, torch.tensor(1.0))):.1f} and "
              f"{float(branch_weights(linear, torch.tensor(1.0))):.1f}")
        print(f"  weighted energies {weighted_energy(2.0, lambda t: 1.0 - linear(t)):.1f} and "
              f"{weighted_energy(-2.0, linear):.1f}, total {near['linear']:.1f}")
        print(f"  with the second target at -4: linear schedule {far['linear']:.1f}, "
              f"late schedule {far['late']:.1f}, saving {far['linear'] - far['late']:.1f}")

    bent = interpolant_energy(0.0, 2.0, 1.0)
    straight = interpolant_energy(0.0, 2.0, 0.0)
    grid = torch.tensor([0.0, 0.5, 1.0])
    if loud:
        print(f"\nThe interpolant with a correction of 1 passes through "
              f"{[round(float(z), 2) for z in interpolant(0.0, 2.0, 1.0, grid)]}")
        speed = [round(float(z), 0) for z in interpolant_velocity(0.0, 2.0, 1.0, grid)]
        print(f"  its velocity is {speed} and its energy {bent:.4f}, against "
              f"{straight:.1f} for the straight line")

    reference_mass, inside, outside = 0.1, 1.0, 0.01  # The two-outcome example of Section 7.27.1.
    normalizer = reference_mass * inside + (1.0 - reference_mass) * outside
    tilted_mass = reference_mass * inside / normalizer
    reward_average = tilted_mass * math.log(inside) + (1.0 - tilted_mass) * math.log(outside)
    tilt_entropy = reward_average - math.log(normalizer)
    if loud:
        print(f"\nA reference reaching B one time in ten, tilted by pi_B = 1 inside and 0.01 outside")
        print(f"  Z = {normalizer:.3f}, P*(B) = {tilted_mass:.3f}, "
              f"a factor of {tilted_mass / reference_mass:.2f}")
        print(f"  E[r] = {reward_average:.3f}, -log Z = {-math.log(normalizer):.3f}, "
              f"KL(P* || R) = {tilt_entropy:.3f}, J(b*) = {tilt_entropy - reward_average:.3f}")

    step_drift = (-0.3 + 0.5) * 0.01  # The physical force plus the bias, over one step.
    step_friction = -1.0 * 0.8 * 0.01  # Minus the friction times the current velocity.
    step_noise = NOISE * math.sqrt(0.01) * 0.4  # The noise scale times a normal draw of 0.4.
    unit_control = 0.5 / NOISE  # The bias force measured in units of the noise.
    plane = compose_bias(torch.tensor([[0.5]]), torch.tensor([[[0.3, -0.4]]]),
                         torch.tensor([[[1.0, 0.0]]]))[0, 0]  # The example of Section 7.27.3.
    log_weights = torch.tensor([-1.0, 0.0, 1.0])
    softmax = torch.softmax(log_weights, dim=0)
    if loud:
        print(f"\nOne Euler-Maruyama step at dt = 0.01, with a noise scale of {NOISE:.4f}")
        print(f"  dv = {step_drift:.4f} {step_friction:+.4f} {step_noise:+.4f} = "
              f"{step_drift + step_friction + step_noise:.4f}, dr = {0.8 * 0.01:.4f}")
        print(f"  the control in units of the noise is {unit_control:.4f}, contributing "
              f"{0.5 * unit_control ** 2 * 0.01:.6f} per step and "
              f"{0.5 * unit_control ** 2 * 0.01 * 100:.4f} over a hundred")
        print(f"  b = {[round(float(z), 1) for z in plane]}, with alpha "
              f"{float((plane * torch.tensor([1.0, 0.0])).sum()):.1f} along s_hat "
              f"and length {float(plane.norm()):.3f}")
        print(f"  three log-weights -1, 0, 1 normalize to {[round(float(z), 3) for z in softmax]}")

    force = BiasForce(args.width).to(device)
    reference_run = roll_out(None, args.samples, SIM_STEPS, device, generator)
    before = roll_out(force, args.samples, SIM_STEPS, device, generator)

    def loss_fn(step: int) -> torch.Tensor:
        rollout = roll_out(force, PATHS, SIM_STEPS, device, generator)  # Simulate under b_bar.
        return cross_entropy_loss(force, rollout, importance_weights(rollout))

    if loud:
        print(f"\ntraining the bias force, {count_parameters(force):,} parameters, "
              f"{PATHS} paths of {SIM_STEPS} steps per optimizer step")
    losses = train(force, loss_fn, steps=args.steps, lr=args.lr, quiet=not loud, log_every=50)
    after = roll_out(force, args.samples, SIM_STEPS, device, generator)
    effective = float(1.0 / (importance_weights(after) ** 2).sum() / args.samples)

    report = {
        "weight of branch 1 at t = 1/2": half,
        "weighted energy, branch 0": weighted_energy(2.0, lambda t: 1.0 - linear(t)),
        "weighted energy, branch 1": weighted_energy(-2.0, linear),
        "total, both ends at distance 2": near["linear"],
        "total with the far target, linear": far["linear"],
        "total with the far target, late": far["late"],
        "a branch drawn at t = 1/2": branch_sample(half, generator),
        "interpolant energy, correction 1": bent,
        "interpolant energy, straight": straight,
        "Z of the two outcomes": normalizer,
        "P*(B)": tilted_mass,
        "the factor the tilt applies": tilted_mass / reference_mass,
        "E_{P*}[r]": reward_average,
        "KL(P* || R)": tilt_entropy,
        "-log Z, the value of J(b*)": -math.log(normalizer),
        "one velocity step": step_drift + step_friction + step_noise,
        "control in units of the noise": unit_control,
        "control energy over a hundred steps": 0.5 * unit_control ** 2 * 0.01 * 100,
        "the composed force": [round(float(z), 2) for z in plane],
        "its length": float(plane.norm()),
        "softmax of -1, 0, 1": [round(float(z), 3) for z in softmax],
        "mean reward, reference": float(reference_run["reward"].mean()),
        "mean reward, before training": float(before["reward"].mean()),
        "mean reward, after training": float(after["reward"].mean()),
        "control energy after training": float(after["energy"].mean()),
        "effective sample size, fraction": effective,
        "force written by hand against autograd": force_check(after["positions"][-1]),
        "final loss": sum(losses[-10:]) / len(losses[-10:]),
    }
    print_report("Branched and Entangled bridges, Chapter 7 Sections 7.26 to 7.27.3", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()