Spaces:
Running
Running
| """Clean-room BC with action quantization from "Understanding Behavior Cloning with Action Quantization" | |
| (arXiv 2603.20538). numpy, CPU. | |
| MDP: x_{t+1} = alpha x_t + u_t + noise (P-IISS: alpha<1 contracting). Expert u* = -alpha x (RTVC: Lipschitz). | |
| Quantizer: binning (smooth) maps u to nearest bin center; epsilon_q = bin width. | |
| BC: fit a policy to quantized expert actions via least-squares. Regret = sum_t (u_t - u*_t)^2. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| ALPHA = 0.9 # dynamics contraction (P-IISS) | |
| def expert_action(x): | |
| return -ALPHA * x | |
| def binning_quantizer(u, n_bins, u_range=(-2, 2)): | |
| """Smooth binning quantizer: maps u to nearest bin center. epsilon_q = bin width.""" | |
| edges = np.linspace(u_range[0], u_range[1], n_bins + 1) | |
| centers = (edges[:-1] + edges[1:]) / 2 | |
| idx = np.clip(np.digitize(u, edges) - 1, 0, n_bins - 1) | |
| return centers[idx], (u_range[1] - u_range[0]) / n_bins # quantized action, epsilon_q | |
| def nonsmooth_quantizer(u, n_bins, rng, u_range=(-2, 2)): | |
| """Non-smooth quantizer: random assignment within bins (violates smoothness).""" | |
| edges = np.linspace(u_range[0], u_range[1], n_bins + 1) | |
| idx = np.clip(np.digitize(u, edges) - 1, 0, n_bins - 1) | |
| # assign to a RANDOM center (not the nearest) -> non-smooth | |
| centers = rng.uniform(u_range[0], u_range[1], n_bins) | |
| return centers[idx], (u_range[1] - u_range[0]) / n_bins | |
| def run_bc(n_samples, H, n_bins, smooth=True, seed=0): | |
| """Collect n expert (x,u) pairs, quantize, fit BC, evaluate regret over horizon H.""" | |
| rng = np.random.default_rng(seed) | |
| # collect expert data | |
| X_data = rng.standard_normal(n_samples) * 2 | |
| U_data = np.array([expert_action(x) for x in X_data]) | |
| # quantize | |
| if smooth: | |
| U_q, eps_q = binning_quantizer(U_data, n_bins) | |
| else: | |
| U_q, eps_q = nonsmooth_quantizer(U_data, n_bins, rng) | |
| # BC: fit u_hat = theta * x (linear) to quantized actions | |
| theta = np.sum(X_data * U_q) / max(np.sum(X_data ** 2), 1e-9) | |
| # evaluate regret over horizon H | |
| x = rng.standard_normal() * 2; regret = 0.0 | |
| for _ in range(H): | |
| u_bc = theta * x; u_expert = expert_action(x) | |
| regret += (u_bc - u_expert) ** 2 | |
| x = ALPHA * x + u_bc + rng.standard_normal() * 0.1 | |
| return regret, eps_q, theta | |
| def run_bc_augmented(n_samples, H, n_bins, seed=0): | |
| """Model-based augmentation: augment data with model predictions (Theorem 7).""" | |
| rng = np.random.default_rng(seed) | |
| X_orig = rng.standard_normal(n_samples) * 2 | |
| U_orig = np.array([expert_action(x) for x in X_orig]) | |
| U_q, eps_q = binning_quantizer(U_orig, n_bins) | |
| # augmentation: generate additional data via the model x' = alpha*x + u | |
| X_aug = np.concatenate([X_orig, ALPHA * X_orig + U_q + rng.standard_normal(n_samples) * 0.1]) | |
| U_aug = np.concatenate([U_q, np.array([expert_action(x) for x in ALPHA * X_orig + U_q])]) | |
| U_aug_q, _ = binning_quantizer(U_aug, n_bins) | |
| theta = np.sum(X_aug * U_aug_q) / max(np.sum(X_aug ** 2), 1e-9) | |
| x = rng.standard_normal() * 2; regret = 0.0 | |
| for _ in range(H): | |
| u_bc = theta * x; regret += (u_bc - expert_action(x)) ** 2 | |
| x = ALPHA * x + u_bc + rng.standard_normal() * 0.1 | |
| return regret, eps_q | |