Download lecture_7/dynamic_bridge.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 13.1 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/dynamic_bridge.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_7/dynamic_bridge.py
-
curl -L -o dynamic_bridge.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_7/dynamic_bridge.py
13.1 kB
| """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() | |