robomotic/lewm-breakout-plays

--- title: LeWM Breakout Plays Dataset emoji: 🏓 colorFrom: blue colorTo: purple sdk: static --- # LeWM Breakout Plays Dataset This repository contains Atari Breakout trajectories generated by various heuristic policies for the LeWorldModel (LeWM) project. The datasets are organized by the heuristic policy that generated them, and each contains the corresponding Parquet files and statistics. ## Available Heuristics - **center**: A policy that constantly tries to keep the paddle in the exact center of the screen. This generates data with minimal paddle variance, useful for causal inference on object presence. - **heuristic**: An expert-level algorithmic policy that tracks the ball's X position dynamically to successfully hit it. This yields high-reward, long-episode sequences. - **passive**: A completely inactive policy that performs no moving actions, generating sequences that show how the game evolves purely passively (without player interaction). - **random_active**: A policy that selects actions randomly but ensures the game starts (by occasionally pressing FIRE). Generates diverse, short-lived chaotic trajectories. - **wall_hugger**: A policy that drives the paddle to one of the walls and stays there. This creates a data distribution where the paddle's spatial location is extremely biased. --- ## 🔬 Representation Quality Evaluation These datasets are used to train and evaluate the **LeWorldModel (LeWM)**, a JEPA-based world model for Atari Breakout. The quality of the learned representations is measured via **MLP probing** on frozen latent embeddings from the trained encoder and its causal predictor. ### Probe Procedure After training, lightweight MLP probes (3-layer: 256→128→output) and linear probes are trained on top of **frozen** encoder/predictor representations. No gradient flows back through the JEPA model. Each probe addresses a specific question about what physics the encoder has absorbed: | Task | Input | Target | Why | | :--- | :--- | :--- | :--- | | `ball_pos` | Encoder CLS at $t$ | Ball $(x,y)$ at $t$ | Does the encoder localise the ball? | | `ball_vel` | Encoder CLS at $t$ | Ball $(v_x,v_y)$ at $t$ | Does the encoder track motion? | | `ball_pos_pred` | ARPredictor output at $t$ | Ball $(x,y)$ at $t+1$ | Does the world model predict next-frame position? | | `ball_vel_pred` | ARPredictor output at $t$ | Ball $(v_x,v_y)$ at $t+1$ | Does the predictor encode implicit momentum? | | `paddle_pos` | Encoder CLS at $t$ | Paddle $x$ at $t$ | Does the encoder locate the controlled paddle? | | `paddle_pos_act` | CLS + action at $t$ | Paddle $x$ at $t$ | How much does action conditioning improve paddle prediction? | Baselines include: (1) a **static persistence** predictor that copies the current label as the prediction; (2) a **randomly-initialised** ViT encoder (sanity check — should score near R²≈0); (3) a **patch-mean** alternative to the CLS token; and (4) **linear probes** to measure linear separability. ### Probe Results — 100-Epoch Reference Run **Evaluation split**: held-out `my_datasets_eval/` (≈13,225 sequences from heuristic + random_active policies, never seen during training). #### MLP Probe — R² (test set) | Task | ep 1 | ep 10 | ep 25 | ep 50 | ep 100 | | :--- | :---: | :---: | :---: | :---: | :---: | | `ball_pos` | +0.178 | **+0.306** | +0.212 | +0.234 | +0.272 | | `ball_vel` | -0.222 | -0.314 | -0.269 | -0.296 | -0.271 | | `ball_pos_pred` | +0.131 | +0.243 | +0.230 | +0.226 | +0.253 | | `ball_vel_pred` | -0.144 | -0.299 | -0.260 | -0.249 | -0.270 | | `paddle_pos` | -0.253 | -0.548 | -0.370 | -0.562 | **+0.319** | | `paddle_pos_act` | -0.169 | -0.170 | -0.187 | -0.536 | **+0.456** | #### Linear Probe — R² (test set) | Task | ep 1 | ep 10 | ep 25 | ep 50 | ep 100 | | :--- | :---: | :---: | :---: | :---: | :---: | | `ball_pos` | -0.143 | +0.244 | +0.249 | +0.327 | +0.301 | | `ball_vel` | -0.022 | -0.062 | -0.059 | -0.094 | -0.078 | | `ball_pos_pred` | -0.194 | +0.198 | +0.282 | +0.278 | +0.276 | | `ball_vel_pred` | -0.017 | -0.031 | -0.040 | -0.058 | -0.068 | | `paddle_pos` | -0.384 | -0.173 | -0.110 | -0.009 | **+0.478** | | `paddle_pos_act` | -0.076 | +0.077 | +0.143 | +0.183 | **+0.607** | #### Sanity-Check Baselines | Baseline | `ball_pos` R² | `ball_vel` R² | `paddle_pos` R² | | :--- | :---: | :---: | :---: | | Static copy (predict $t$ as $t+1$) | +0.955 | +0.821 | +0.778 | | Random encoder | -0.351 | -0.018 | -0.034 | ### Key Findings - **Ball position is learned early**: R² reaches +0.31 by epoch 10. The linear probe improves monotonically, confirming representations become more linearly decodable over the full 100 epochs. - **Ball velocity is not encoded**: Negative R² across all epochs and probe types — JEPA training does not provide an explicit signal to retain first-order dynamics. - **Predictor tracks ball position well**: At epoch 100, the ARPredictor R² (+0.25) is within 2% of the encoder R² (+0.27), demonstrating genuine next-frame anticipation. - **Paddle position emerges late**: Val R² grows steadily (reaching +0.92 at ep100), but test R² is negative until epoch 100 when it jumps to +0.32/+0.48 — a distribution-shift effect from the held-out policy mix. - **Action conditioning is effective**: Adding the action embedding boosts paddle prediction by +0.13 (linear, ep100), confirming the ARPredictor correctly absorbs action-conditioned dynamics. - **Random encoder is uninformative**: R² ≈ −0.35 for ball_pos under a random ViT confirms all probe gains come from training, not pixel statistics. For full methodology, all results, and reproduction instructions, see the companion model repository at [robomotic/lewm-breakout](https://huggingface.co/robomotic/lewm-breakout).--- title: LeWM打砖块游戏轨迹数据集 emoji: 🏓 colorFrom: blue colorTo: purple sdk: static --- # LeWM打砖块游戏轨迹数据集 本仓库收录了面向LeWorldModel（LeWM）项目生成的多款启发式策略所产生的雅达利打砖块（Atari Breakout）游戏轨迹数据。数据集按照生成其数据的启发式策略进行分类，每个分类下均包含对应的Parquet文件与统计信息。 ## 可用启发式策略 - **居中策略（center）**：该策略始终尝试将挡板保持在屏幕正中央，生成的轨迹中挡板位移方差极小，适用于针对物体存在性的因果推断任务。 - **专家启发策略（heuristic）**：该策略为专家级算法方案，可动态追踪球体的X轴位置以精准击球，能够生成高奖励、长时长的游戏序列。 - **被动策略（passive）**：该策略完全不执行任何移动操作，生成的序列可直观展示无玩家交互时游戏的纯被动演化过程。 - **随机主动策略（random_active）**：该策略随机选择操作，但会偶尔按下开火键以确保游戏正常启动，可生成多样化、短时长的混乱轨迹数据。 - **贴墙策略（wall_hugger）**：该策略将挡板移动至某一侧墙壁并保持静止，生成的数据集会呈现出挡板空间位置高度偏向的分布特征。 ## 🔬 表征质量评估 这些数据集被用于训练与评估**LeWorldModel（LeWM）**——一款面向雅达利打砖块游戏、基于JEPA的世界模型。所习得表征的质量，将通过对训练好的编码器及其因果预测器输出的冻结隐嵌入向量执行**多层感知器探测（MLP probing）**进行评估。 ### 探测流程 训练完成后，将基于**冻结**的编码器/预测器表征，训练轻量级多层感知器探测模型（3层结构：256→128→输出维度）与线性探测模型。训练过程中，梯度不会反向传播至JEPA模型。每个探测模型均针对编码器所习得的物理规则，回答一个特定问题： | 任务 | 输入 | 目标 | 探究方向 | | :--- | :--- | :--- | :--- | | `ball_pos` | 编码器在时刻$t$的CLS令牌 | 时刻$t$的球体$(x,y)$坐标 | 编码器是否能够定位球体？ | | `ball_vel` | 编码器在时刻$t$的CLS令牌 | 时刻$t$的球体速度$(v_x,v_y)$ | 编码器是否能够追踪运动？ | | `ball_pos_pred` | ARPredictor在时刻$t$的输出 | 时刻$t+1$的球体$(x,y)$坐标 | 世界模型是否能够预测下一帧的球体位置？ | | `ball_vel_pred` | ARPredictor在时刻$t$的输出 | 时刻$t+1$的球体速度$(v_x,v_y)$ | 预测器是否编码了隐式动量信息？ | | `paddle_pos` | 编码器在时刻$t$的CLS令牌 | 时刻$t$的挡板$x$坐标 | 编码器是否能够定位受控挡板？ | | `paddle_pos_act` | 时刻$t$的CLS令牌+动作嵌入 | 时刻$t$的挡板$x$坐标 | 动作条件对挡板预测的提升效果如何？ | 基准测试方案包括：(1) **静态持久化预测器**：直接将当前标签作为预测结果；(2) **随机初始化的视觉Transformer（ViT）编码器**：作为合理性校验基准，其$R^2$得分应接近0；(3) **补丁均值替代方案**：用图像补丁均值替代CLS令牌作为输入；(4) **线性探测模型**：用于衡量表征的线性可分性。 ### 探测结果——100轮训练参考运行结果 **评估划分**：采用预留的`my_datasets_eval/`数据集（包含来自专家启发策略与随机主动策略的约13225条游戏序列，训练过程中从未见过这些数据）。 #### 多层感知器探测结果——测试集$R^2$得分 | 任务 | ep 1 | ep 10 | ep 25 | ep 50 | ep 100 | | :--- | :---: | :---: | :---: | :---: | :---: | | `ball_pos` | +0.178 | **+0.306** | +0.212 | +0.234 | +0.272 | | `ball_vel` | -0.222 | -0.314 | -0.269 | -0.296 | -0.271 | | `ball_pos_pred` | +0.131 | +0.243 | +0.230 | +0.226 | +0.253 | | `ball_vel_pred` | -0.144 | -0.299 | -0.260 | -0.249 | -0.270 | | `paddle_pos` | -0.253 | -0.548 | -0.370 | -0.562 | **+0.319** | | `paddle_pos_act` | -0.169 | -0.170 | -0.187 | -0.536 | **+0.456** | #### 线性探测结果——测试集$R^2$得分 | 任务 | ep 1 | ep 10 | ep 25 | ep 50 | ep 100 | | :--- | :---: | :---: | :---: | :---: | :---: | | `ball_pos` | -0.143 | +0.244 | +0.249 | +0.327 | +0.301 | | `ball_vel` | -0.022 | -0.062 | -0.059 | -0.094 | -0.078 | | `ball_pos_pred` | -0.194 | +0.198 | +0.282 | +0.278 | +0.276 | | `ball_vel_pred` | -0.017 | -0.031 | -0.040 | -0.058 | -0.068 | | `paddle_pos` | -0.384 | -0.173 | -0.110 | -0.009 | **+0.478** | | `paddle_pos_act` | -0.076 | +0.077 | +0.143 | +0.183 | **+0.607** | #### 合理性校验基准结果 | 基准方案 | `ball_pos` $R^2$得分 | `ball_vel` $R^2$得分 | `paddle_pos` $R^2$得分 | | :--- | :---: | :---: | :---: | | 静态复制预测器（将时刻$t$的标签作为$t+1$的预测结果） | +0.955 | +0.821 | +0.778 | | 随机初始化编码器 | -0.351 | -0.018 | -0.034 | ### 核心发现 - **球体位置可被快速习得**：训练至第10轮时，$R^2$得分已达到+0.31。线性探测的$R^2$得分随训练轮次单调提升，证实表征的线性可解性在完整100轮训练中持续优化。 - **未习得球体速度表征**：所有训练轮次与探测类型下的$R^2$得分均为负值，说明JEPA训练未提供保留一阶动力学信息的显式信号。 - **预测器可有效追踪球体位置**：训练至第100轮时，AR预测器的$R^2$得分（+0.25）与编码器的$R^2$得分（+0.27）相差不到2%，证实模型具备真实的下一帧预测能力。 - **挡板位置表征出现较晚**：验证集$R^2$得分稳步提升（第100轮时达到+0.92），但测试集$R^2$得分在第100轮前始终为负值，之后跃升至+0.32/+0.48，这是预留数据集的策略混合带来的分布偏移效应。 - **动作条件信息有效**：加入动作嵌入后，挡板预测的$R^2$得分提升了+0.13（线性探测，第100轮），证实AR预测器正确学习到了动作条件下的动力学规律。 - **随机初始化编码器无有效表征**：随机初始化的ViT编码器在`ball_pos`任务上的$R^2$得分约为-0.35，证实探测模型的性能提升均来自训练过程，而非像素本身的统计特征。 如需获取完整方法论、全部实验结果与复现指南，请访问配套模型仓库：[robomotic/lewm-breakout](https://huggingface.co/robomotic/lewm-breakout).