"""
=============================================================================
第 16 讲 · 01：为什么机器人的动作需要生成模型 —— 回归、离散化与流匹配
=============================================================================

对应讲义：lectures/lecture16_robotics_agents.md 第 3 节（RT-2 的动作 token；π0 的流匹配动作专家）

-----------------------------------------------------------------------------
【一个"绕开障碍"的决策】
-----------------------------------------------------------------------------
机器人正前方有一个障碍。人类示范里，一半人从左边绕（横向速度约 −1），一半从右边绕（约 +1）。
这就是**多峰**（multimodal）的动作分布：两种做法都对，但它们的"平均"（直走，速度约 0）会撞上障碍。
撞障碍的判定：|横向速度| < 0.5。

三种学习"动作"的方法：
    · 回归（MSE）：学一个函数直接输出动作，最小化均方误差 —— 最优解是示范的**平均**。
    · 离散化 + 采样（RT-2 的做法）：把动作范围切成 256 个格子，每个格子当作一个 token，
      像语言模型一样学"下一个 token 的分布"，然后从中采样。toy 版本直接用示范的直方图当作学到的分布。
    · 流匹配（flow matching，π0 的动作专家）：从噪声 x₀ ~ N(0, 1) 出发，沿着一个学到的"速度场"
      v(x, t) 走到 t = 1，终点就是一个动作样本。训练目标：对配对的 (x₀, x₁)（x₁ 是示范动作），
      在中间点 x_t = (1 − t) x₀ + t x₁ 上让 v 去拟合 x₁ − x₀。
      toy 版本不训练神经网络，而是用核回归直接估计最优的速度场：
          v(x, t) = 加权平均(x₁ − x₀)，权重 = exp(−(x − x_t)² / (2h²))
      （最优的速度场恰好是"给定 x_t = x 时 x₁ − x₀ 的条件期望"，核回归就是在估计这个期望。）
      生成时用欧拉法积分 K 步：x ← x + v(x, k/K) / K。

运行时间：约 1 秒（只用标准库）。
"""

import math
import random

N_DEMOS = 400               # 示范的数量（也用作流匹配的训练配对数）
N_SAMPLES = 300             # 每种方法生成多少个动作来统计
BANDWIDTH = 0.15            # 核回归的带宽 h
N_BINS, LOW, HIGH = 256, -2.0, 2.0


def demonstrations(rng):
    """示范的横向速度：一半 N(+1, 0.1²)，一半 N(−1, 0.1²)。"""
    return [rng.gauss(1.0 if rng.random() < 0.5 else -1.0, 0.1) for _ in range(N_DEMOS)]


def mse_policy(demos):
    """
    回归（MSE）：最小化 Σ(a − aᵢ)² 的 a 就是平均值。

    【思路】这里只有一个状态（障碍前），所以"学一个函数"退化成"学一个数"。
        有很多状态时也一样：在每个状态上，MSE 回归的最优输出是"这个状态下示范动作的条件平均"。
    """
    return math.fsum(demos) / len(demos)


def histogram_policy(demos):
    """
    离散化：把 [LOW, HIGH] 切成 N_BINS 格，统计每格的示范数（加一点平滑），得到每个 token 的概率。

    【返回】长度 N_BINS 的概率列表
    【边界】超出范围的示范被截断到边界的格子。
    """
    counts = [0.1] * N_BINS                                   # 加 0.1 的平滑，避免概率为 0
    for a in demos:
        i = int((a - LOW) / (HIGH - LOW) * N_BINS)
        counts[min(N_BINS - 1, max(0, i))] += 1
    total = math.fsum(counts)
    return [c / total for c in counts]


def sample_histogram(probs, rng):
    """从格子的分布里采一个 token，返回格子中心对应的动作值。"""
    u, acc = rng.random(), 0.0
    for i, p in enumerate(probs):
        acc += p
        if u < acc:
            return LOW + (i + 0.5) * (HIGH - LOW) / N_BINS
    return LOW + (N_BINS - 0.5) * (HIGH - LOW) / N_BINS


def flow_pairs(demos, rng):
    """流匹配的训练配对：每个示范动作 x₁ 配一个随机噪声 x₀ ~ N(0, 1)。"""
    return [(rng.gauss(0, 1), x1) for x1 in demos]


def velocity(x, t, pairs):
    """
    核回归估计的速度场 v(x, t) ≈ E[x₁ − x₀ | x_t = x]。

    【思路】每个配对在时刻 t 的位置是 x_t = (1 − t) x₀ + t x₁；离 x 越近的配对权重越大。
    【边界】如果所有权重都极小（x 离所有配对都很远），退回到所有配对的平均速度，避免除以 0。
    【自我改进联系】真实的 π0 用一个约 3 亿参数的"动作专家"网络来拟合这个速度场，输入还包括图像和语言指令。
    """
    num = den = 0.0
    for x0, x1 in pairs:
        xt = (1 - t) * x0 + t * x1
        w = math.exp(-((x - xt) ** 2) / (2 * BANDWIDTH ** 2))
        num += w * (x1 - x0)
        den += w
    if den < 1e-300:
        return math.fsum(x1 - x0 for x0, x1 in pairs) / len(pairs)
    return num / den


def sample_flow(pairs, steps, rng):
    """从噪声出发，用欧拉法积分 steps 步，得到一个动作。"""
    x = rng.gauss(0, 1)
    for k in range(steps):
        x += velocity(x, k / steps, pairs) / steps
    return x


def summarize(actions):
    """撞障碍的比例（|a| < 0.5），以及离最近的示范模式（±1）的平均距离。"""
    crash = sum(1 for a in actions if abs(a) < 0.5) / len(actions)
    off = math.fsum(abs(abs(a) - 1.0) for a in actions) / len(actions)
    left = sum(1 for a in actions if a < 0) / len(actions)
    return crash, off, left


def text_histogram(actions, width=40):
    """在终端里画一个简单的直方图：把 [−2, 2] 分成 16 格。"""
    bins = [0] * 16
    for a in actions:
        i = int((a + 2) / 4 * 16)
        bins[min(15, max(0, i))] += 1
    peak = max(bins)
    lines = []
    for i, b in enumerate(bins):
        lo = -2 + i * 0.25
        lines.append(f"    [{lo:+.2f}, {lo + 0.25:+.2f})  " + "█" * int(round(b / peak * width)))
    return "\n".join(lines)


def main():
    rng = random.Random(1600)
    demos = demonstrations(rng)
    pairs = flow_pairs(demos, rng)
    probs = histogram_policy(demos)

    print("=" * 78)
    print(f"绕开障碍：{N_DEMOS} 个示范（一半向左 −1，一半向右 +1），每种方法生成 {N_SAMPLES} 个动作")
    print("=" * 78)
    print(f"{'':36}{'撞障碍':>8}{'离示范模式的距离':>18}{'向左的比例':>12}")
    rows = [("回归（MSE）", [mse_policy(demos)] * N_SAMPLES),
            ("离散化（256 个 token）+ 采样", [sample_histogram(probs, rng) for _ in range(N_SAMPLES)])]
    flow_samples = {}
    for steps in (1, 2, 5, 10, 20):
        flow_samples[steps] = [sample_flow(pairs, steps, rng) for _ in range(N_SAMPLES)]
        rows.append((f"流匹配（欧拉法 {steps} 步）", flow_samples[steps]))
    for name, actions in rows:
        crash, off, left = summarize(actions)
        width = sum(2 if ord(ch) > 127 else 1 for ch in name)
        print(name + " " * max(1, 36 - width) + f"{crash:8.3f}{off:18.3f}{left:12.3f}")

    print("\n流匹配（欧拉法 1 步）生成的动作分布：")
    print(text_histogram(flow_samples[1]))
    print("\n流匹配（欧拉法 10 步）生成的动作分布：")
    print(text_histogram(flow_samples[10]))

    print("\n观察：")
    print("  · 回归（MSE）输出示范的平均 —— 一个几乎为 0 的横向速度，每次都直直地撞上障碍。")
    print("    （它的\"向左的比例\"是 1.000，只是因为这个平均值碰巧略小于 0。）")
    print("    两种做法都对，它们的平均却是错的：这就是机器人动作需要**生成模型**的原因。")
    print("  · 离散化 + 采样（RT-2 把每个动作维度切成 256 个 token）能表示多峰分布，几乎不撞；")
    print("    代价是精度受格子宽度限制，动作维度多、控制频率高时 token 数很多（π0-FAST 用压缩的动作 token 缓解）。")
    print("  · 流匹配（π0 的动作专家）同样能生成两个峰，而且动作是连续的。")
    print("  · 但积分步数太少时，流匹配也会退化成\"平均\"：只走 1 步时，x₀ 和 x₁ 还没有任何关联，")
    print("    速度场只能指向示范的平均，于是 97% 的样本撞上障碍；走 5 步以上才能把两个峰分开。")
    print("    生成模型的\"采样步数\"是精度和推理速度之间的权衡（机器人要每秒输出几十个动作）。\n")


if __name__ == "__main__":
    main()