Blame for lecture_7/dynamic_bridge.py

"""A diffusion reference: Girsanov, the control problem, Hopf-Cole, and the Brownian bridge.

We change the drift of a diffusion and keep its noise, the only change under which the two path
laws stay comparable, and over one short step the two increments differ only in their mean. We sum
those steps to get Girsanov's formula, and averaging it under the controlled law turns relative
entropy on paths into an integral of a squared drift. The Hopf-Cole change of variables then turns
the resulting Hamilton-Jacobi-Bellman equation into the linear backward equation of the reference.

Chapter 7, Sections 7.16 to 7.18 of the notes. Runs on a driftless diffusion in one dimension.
Prints the two computations of the relative entropy of a constant control, the residual of the
backward equation and of the Hamilton-Jacobi-Bellman equation, and the Brownian bridge between two
fixed points.

    python lecture_7/dynamic_bridge.py              # 200 Euler-Maruyama steps
    python lecture_7/dynamic_bridge.py --steps 20   # a coarse time grid, for a first look
"""

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

import math

import torch

from cis6270.runner import common_parser, print_report, save_run, set_seed

DTYPE = torch.float64  # Double precision for the closed forms and the Monte Carlo check.

EPSILON = 1.0  # The reference noise variance per unit time.
START, END = 0.0, 2.0  # The two points the bridge is pinned at.
PATHS = 8192  # Trajectories in the Monte Carlo checks.


# %% 02. One short step
#
# Over a step of length dt both increments are Gaussian with the covariance epsilon dt, and they
# differ only in the mean, by u dt. The logarithm of the ratio of the two densities is the inner
# product of w, the increment beyond the reference drift, with u over epsilon, minus the squared
# norm of u times dt over 2 epsilon. Section 7.16.1.


def log_density_ratio_step(w: torch.Tensor, u: torch.Tensor, dt: float,
                           epsilon: float = EPSILON) -> torch.Tensor:
    """log p^u / p^R over one step, with w = Delta X - b dt the part beyond the reference drift."""
    return (w * u).sum(dim=-1) / epsilon - (u ** 2).sum(dim=-1) * dt / (2.0 * epsilon)


# %% 03. Summing over the steps, and the average under the controlled law
#
# We add the per-step contributions and let dt go to zero. That gives Girsanov's formula, a
# martingale term built from the noise and a drift term built from the size of the control. When we
# average under P^u the two energy terms are left with coefficients 1/epsilon and -1/(2 epsilon).
# Sections 7.16.2 and 7.16.3.


def simulate_controlled(u: float, steps: int, paths: int, epsilon: float = EPSILON,
                        generator: torch.Generator | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    """Euler-Maruyama for dX = u dt + sqrt(epsilon) dB, returning the endpoints and the log ratio.

    We accumulate the log ratio step by step from the same Brownian increments the path uses, the
    way Section 7.27.3 accumulates it during training.
    """
    dt = 1.0 / steps
    x = torch.zeros(paths, 1, dtype=DTYPE)  # [P, 1] every path starts at the origin.
    log_ratio = torch.zeros(paths, dtype=DTYPE)  # [P] the Girsanov integral along each path.
    control = torch.full((paths, 1), u, dtype=DTYPE)  # [P, 1] a control constant in space and time.
    for _ in range(steps):
        noise = torch.randn(paths, 1, dtype=DTYPE, generator=generator) * math.sqrt(dt)  # dB.
        w = control * dt + math.sqrt(epsilon) * noise  # Delta X with no reference drift.
        log_ratio = log_ratio + log_density_ratio_step(w, control, dt, epsilon)
        x = x + w  # The controlled state after this step.
    return x.squeeze(1), log_ratio


def control_energy_entropy(u: float, epsilon: float = EPSILON) -> float:
    """KL(P^u || R) = (1 / 2 epsilon) E int_0^1 ||u||^2 dt, for a control constant in time."""
    return u ** 2 / (2.0 * epsilon)  # The integral of a constant over the unit interval.


def gaussian_entropy_gap(shift: float, variance: float) -> float:
    """KL between two Gaussians with a common variance, the squared difference of the means over 2v."""
    return shift ** 2 / (2.0 * variance)


# %% 04. The value function and the optimal control
#
# Bellman's principle and Ito's formula give the Hamilton-Jacobi-Bellman equation. Only two of its
# terms contain the control. We complete the square in those two and isolate u* = -grad V. We
# then substitute it back and are left with a closed equation for V with a squared gradient in it.
# Sections 7.17 to 7.17.2.


def optimal_control(value, x: float, t: float) -> float:
    """u*(x, t) = -grad V_t(x), the control pointing downhill on the cost-to-go."""
    xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
    v = value(xs, torch.tensor(t, dtype=DTYPE))
    return float(-torch.autograd.grad(v, xs)[0])


def hjb_residual(value, x: float, t: float, epsilon: float = EPSILON) -> float:
    """d_t V + (epsilon / 2) grad^2 V - (1/2) ||grad V||^2, with no reference drift."""
    xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
    ts = torch.tensor(t, dtype=DTYPE, requires_grad=True)
    v = value(xs, ts)
    dv_dt = torch.autograd.grad(v, ts, create_graph=True)[0]
    dv_dx = torch.autograd.grad(v, xs, create_graph=True)[0]
    d2v_dx2 = torch.autograd.grad(dv_dx, xs, create_graph=True)[0]
    return float((dv_dt + 0.5 * epsilon * d2v_dx2 - 0.5 * dv_dx ** 2).detach())


# %% 05. Removing the squared gradient
#
# We write the value function as minus epsilon times a logarithm. That cancels the squared
# gradient against the extra term the Laplacian produces and leaves the Kolmogorov backward
# equation of the reference, a linear equation. In the new variable the control is epsilon times
# the gradient of the logarithm of phi. Sections 7.17.3 and 7.17.4.


def hopf_cole(potential, epsilon: float = EPSILON):
    """V_t = -epsilon log phi_t, the change of variables that linearizes the equation."""
    return lambda x, t: -epsilon * torch.log(potential(x, t))


def backward_residual(potential, x: float, t: float, epsilon: float = EPSILON) -> float:
    """d_t phi + (epsilon / 2) grad^2 phi, zero for a reference transition density."""
    xs = torch.tensor(x, dtype=DTYPE, requires_grad=True)
    ts = torch.tensor(t, dtype=DTYPE, requires_grad=True)
    phi = potential(xs, ts)
    dphi_dt = torch.autograd.grad(phi, ts, create_graph=True)[0]
    dphi_dx = torch.autograd.grad(phi, xs, create_graph=True)[0]
    d2phi_dx2 = torch.autograd.grad(dphi_dx, xs, create_graph=True)[0]
    return float((dphi_dt + 0.5 * epsilon * d2phi_dx2).detach())


# %% 06. The Brownian bridge
#
# We condition on a single final point y. The potential becomes the reference density of
# arriving at y from the current state. Epsilon times the gradient of its logarithm is then the
# remaining displacement over the remaining time, and the noise level cancels from the drift and
# survives in the variance. Section 7.18.


def bridge_potential(z: torch.Tensor, t: torch.Tensor, y: float = END,
                     epsilon: float = EPSILON) -> torch.Tensor:
    """h(z, t) = R_{1|t}(y | z), the full Gaussian density with variance epsilon (1 - t)."""
    remaining = 1.0 - t  # The time left, or the variance divided by epsilon.
    return torch.exp(-((y - z) ** 2) / (2.0 * epsilon * remaining)) / torch.sqrt(
        2.0 * math.pi * epsilon * remaining)


def bridge_drift(z: torch.Tensor, t: torch.Tensor, y: float = END) -> torch.Tensor:
    """(y - z) / (1 - t), the velocity that reaches y on time if nothing further perturbs the state."""
    return (y - z) / (1.0 - t)


def bridge_law(t: float, x: float = START, y: float = END,
               epsilon: float = EPSILON) -> tuple[float, float]:
    """The mean (1 - t) x + t y and the variance epsilon t (1 - t) of the bridge at time t."""
    return (1.0 - t) * x + t * y, epsilon * t * (1.0 - t)


def simulate_bridge(steps: int, paths: int, epsilon: float = EPSILON,
                    generator: torch.Generator | None = None) -> torch.Tensor:
    """Euler-Maruyama for the bridge SDE, returning the state at every time. [steps + 1, P]."""
    dt = 1.0 / steps
    x = torch.full((paths,), START, dtype=DTYPE)
    trajectory = [x.clone()]
    for step in range(steps):
        t = torch.tensor(step * dt, dtype=DTYPE)
        noise = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(dt)
        x = x + bridge_drift(x, t) * dt + math.sqrt(epsilon) * noise  # One Euler-Maruyama step.
        trajectory.append(x.clone())
    return torch.stack(trajectory)  # [steps + 1, P].


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=200)  # --steps is the number of Euler-Maruyama steps.
    args = parser.parse_args()
    set_seed(args.seed)
    loud = not args.quiet
    steps = max(args.steps, 2)
    generator = torch.Generator().manual_seed(args.seed)

    endpoints, log_ratio = simulate_controlled(2.0, steps, PATHS, EPSILON, generator)
    girsanov = float(log_ratio.mean())  # The average of the log density under P^u.
    closed_form = control_energy_entropy(2.0)
    endpoint_form = gaussian_entropy_gap(2.0, EPSILON)
    if loud:
        print(f"A constant control u = 2 at epsilon = {EPSILON:g}, over {steps} steps")
        print(f"  (1 / 2 eps) int ||u||^2 dt = {closed_form:.3f}")
        print(f"  KL(N(x + 2, 1) || N(x, 1)) = {endpoint_form:.3f}")
        print(f"  averaged Girsanov log ratio over {PATHS} paths = {girsanov:.3f}")
        print(f"  endpoint mean {float(endpoints.mean()):.3f}, "
              f"variance {float(endpoints.var()):.3f}")

    value = hopf_cole(bridge_potential)  # V_t = -epsilon log h, with h the pinned potential.
    test_points = [(-0.5, 0.2), (0.5, 0.5), (1.5, 0.5), (2.5, 0.8)]
    backward = max(abs(backward_residual(bridge_potential, x, t)) for x, t in test_points)
    hjb = max(abs(hjb_residual(value, x, t)) for x, t in test_points)
    control_gap = max(
        abs(optimal_control(value, x, t)
            - float(bridge_drift(torch.tensor(x, dtype=DTYPE), torch.tensor(t, dtype=DTYPE))))
        for x, t in test_points)
    if loud:
        print(f"\nThe pinned potential h(z, t) and its two equations")
        print(f"  largest backward residual d_t h + (eps/2) h'' = {backward:.2e}")
        print(f"  largest Hamilton-Jacobi-Bellman residual = {hjb:.2e}")
        print(f"  largest gap between -grad V and (y - z) / (1 - t) = {control_gap:.2e}")

    mean_half, variance_half = bridge_law(0.5)
    sample = simulate_bridge(steps, PATHS, EPSILON, generator)
    half = sample[steps // 2]  # The state at t = 1/2 on the grid.
    state = 0.6732  # The particular state Section 7.18 evaluates the drift at.
    drift_here = float(bridge_drift(torch.tensor(state, dtype=DTYPE),
                                    torch.tensor(0.25, dtype=DTYPE)))
    mean_quarter, variance_quarter = bridge_law(0.25)
    spreads = {e: math.sqrt(e * 0.25) for e in (0.25, 1.0, 4.0)}  # Standard deviations at t = 1/2.
    if loud:
        print(f"\nThe bridge from {START:g} to {END:g} at epsilon = {EPSILON:g}")
        print(f"  at t = 0.25: mean {mean_quarter:.4f}, variance {variance_quarter:.4f}, "
              f"standard deviation {math.sqrt(variance_quarter):.4f}")
        print(f"  at X_0.25 = {state}: drift = ({END:g} - {state}) / 0.75 = {drift_here:.4f}, "
              f"against the straight-line velocity {END - START:.0f}")
        print(f"  at t = 1/2: mean {mean_half:.4f}, variance {variance_half:.4f}")
        print(f"  simulated mean {float(half.mean()):.4f}, variance {float(half.var()):.4f}")
        print(f"  standard deviation at t = 1/2 for eps = 0.25, 1, 4: "
              f"{[round(s, 2) for s in spreads.values()]}")

    report = {
        "KL from the control energy": closed_form,
        "KL from the two endpoint laws": endpoint_form,
        "KL from the simulated paths": girsanov,
        "simulated endpoint mean": float(endpoints.mean()),
        "backward equation residual": backward,
        "Hamilton-Jacobi-Bellman residual": hjb,
        "-grad V against the bridge drift": control_gap,
        "bridge mean at t = 0.25": mean_quarter,
        "bridge variance at t = 0.25": variance_quarter,
        "bridge drift at X = 0.6732": drift_here,
        "straight-line velocity": float(END - START),
        "bridge mean at t = 1/2": mean_half,
        "bridge variance at t = 1/2": variance_half,
        "simulated mean at t = 1/2": float(half.mean()),
        "simulated variance at t = 1/2": float(half.var()),
        "simulated endpoint, mean": float(sample[-1].mean()),
        "spread at t = 1/2, eps = 0.25": spreads[0.25],
        "spread at t = 1/2, eps = 1": spreads[1.0],
        "spread at t = 1/2, eps = 4": spreads[4.0],
    }
    print_report("A diffusion reference, Chapter 7 Sections 7.16 to 7.18", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()