In this third installment of our series, we dive into the results of our first training run on the Koch SO-ARM101 using a JEPA (Joint-Embedding Predictive Architecture). With a limited dataset of only 12 episodes, we weren't expecting a miracle, but rather a clear diagnostic.

Is the model learning the laws of physics or just memorizing pixels? Does our latent space show signs of "collapse," or is it ready to support a high-level planner like a GFlowNet? By analyzing Latent MSE, PCA projections, t-SNE clustering, and Linear Probing, we will map the boundaries of our model's "internal world" and define the roadmap for our next scaling phase.

The training was performed on RunPod using the following template:
https://console.runpod.io/deploy?template=f83357qr5r&ref=7x06vrca
(Note: This is not a commercial collaboration or a paid partnership).

Data analysis

Fig 1: Predictive Drift — Stability or Inertia?

This graph measures the error between what the model "imagines" (prediction) and what it "sees" (real encoding) as the prediction horizon extends further into the future.

• X-Axis (step_idx): Represents the number of "steps" into the future. At Step 0, the model starts from a ground-truth image. By Step 100, it has generated 100 successive states autoregressively, relying solely on its own previous latent predictions and the provided action vectors.

• Theory vs. Reality: In a standard predictive architecture, the MSE (Mean Squared Error) should increase linearly or even exponentially. This is due to the inevitable accumulation of temporal errors—a phenomenon to which JEPA models are particularly sensitive.

Analysis of the Two Phases:

1. The Stagnation Phase (0 to 400 steps): The error remains abnormally low and flat. This is the sign of a "lazy" model. With a limited dataset of only 12 episodes, the JEPA has learned that the static background and noise are the easiest signals to predict. By predicting almost no change, it minimizes its Loss without actually learning the arm's physics. We are witnessing background overfitting.

2. The Rupture Phase (After 400 steps): We see a sudden explosion in MSE. This occurs when the model exceeds the average sequence length of its training data. Having no more known structure to cling to, its "imaginary world" completely collapses.

The Diagnostic: The apparent stability in the first phase is a statistical illusion. The predictor and encoder are "consistent" only because they have both agreed to ignore the dynamics of the robot in favor of a static world.

Fig 2: Consistency Analysis — Latent Neighbors Topology

Visualizing a PCA (Principal Component Analysis) on a model using the SIGReg regularizer may seem paradoxical. Where PCA seeks to isolate the axes of maximum variance, SIGReg aims to uniformize it into a Gaussian distribution (isotropy).

However, it is precisely this "conflict" that makes this visualization valuable: it allows us to observe the struggle between the model's dynamics and its statistical constraints.

The Visual Grammar of Latent Space

To interpret this graph, we use three "health" indicators:

• The Perfect Circle: The regularization has won. The model has "killed" the physical signal in favor of Gaussian statistics.

• Filaments or Loops: The physics is "resisting." Episode trajectories are forcing the space to structure itself despite the constraint.

• Clustering (Point or Line): A sign of total collapse. The encoder is no longer producing useful information.

Observations from Run 1

On this graph, we observe:

1. Local Consistency: The neighbors (blue points) are tightly grouped around the anchors (red crosses). This indicates that the model respects a certain mathematical continuity.

2. "Spider Web" Structure: The distribution is messy but filamentous. This demonstrates that a physical structure is attempting to emerge from the constraints imposed by SIGReg, even if it hasn't yet formed clear semantic clusters (e.g., "closed gripper" vs. "open gripper").

Limitations of the Current Analysis

Technical honesty is required here: my current evaluation script does not yet visually isolate the frames selected by the Nearest Neighbor (NN). While mathematical proximity is evident, I cannot yet state with absolute certainty that two neighbors share identical visual characteristics. This is a key improvement planned for the next diagnostic pipeline.

Fig 3: Global Latent Mapping (PCA Projection)

In this visualization, we use the same PCA projection but color the points according to their source episode. This allows us to track how the model separates individual demonstrations.

1. Sequential Victory: Filaments over Clouds

We observe that each episode forms a semi-independent "filament" rather than an exploded cloud of points. This is a significant victory for the encoder: it demonstrates a clear understanding that state t+1 is closely related to state t. The model has successfully captured the sequential structure of the video.

2. Chaotic Entanglement = Absence of Semantics

However, this is where the "Post-Mortem" gets gritty. In a model that truly "understands" the task (a generalist model), we shouldn't see isolated filaments by episode. Instead, we should see these filaments break apart to form thematic nodes:

• All "rest positions" (starting points) should cluster in the same area.

• All "grasping phases" should form a compact cluster, regardless of which episode they come from.

Currently, we have chaos because:

• The model identifies each episode by its unique noise (slight lighting differences, imperceptibly different starting positions).

• It fails to extract the invariant (the act of grasping) to merge these trajectories into a shared latent meaning.

3. SIGReg and the "Tangled Cables" Effect

Without SIGReg, these filaments would likely be drifting far apart in an immense, unorganized space. SIGReg forces them all into the same "small box" (the Gaussian distribution).

The result? They all end up in the same mathematical coordinate space, but because they lack semantic ties, they cross and intertwine without any physical logic. It is the latent equivalent of a drawer full of tangled cables.

Fig 4: Linear Probing R2 Scores – The Latent-to-Physics Bridge

If the PCA was our "X-ray," the Linear Probing is our "lie detector." It measures exactly how much physical information is actually accessible within the latent space for a controller or a planner.

We perform a simple linear regression from the latent state to predict real-world values. A high score indicates high linear availability; a low score means the information is either missing or "twisted" in a way that a simple system cannot use.

1. Component Breakdown:

• Action: This is the highest score. The model captures "intent"—the correlation between the visual context and the command sent. It confirms that the link between images and actions is the loudest signal, yet it remains significantly underdeveloped.

  $$R^2 \approx 0.35$$

• Proprioception: This is the weak link. The model struggles to locate the arm’s joints in space. A score of 0.30 means it is effectively "blind" to 70% of the actual joint positions. For a robot, this is the equivalent of driving with your eyes closed, trying to guess the steering wheel's position by touch alone.

$$R^2 \approx 0.30$$

• State: The "State" (a combination of image and proprioception) stagnates. This confirms the encoder’s failure to extract a stable world representation from a dataset of only 12 repetitive episodes.

  $$R^2 \approx 0.30$$

2. The "Glass Ceiling" Effect

Seeing all scores clustered around the 0.3 mark reveals a fundamental semantic resolution problem:

1. Non-Linearity: Since the probing is linear, a low score suggests the information might be present but is stored non-linearly. It is "trapped" in a format the model cannot easily translate into action.

2. The Predictor Paradox: As noted in our Epoch 10 analysis, the predictor often outperforms the encoder. This proves the architecture is trying to learn the physics, but the encoder isn't receiving enough visual diversity to "calibrate" its vision.

The Verdict: A global score of 0.31 is an admission of powerlessness for any high-level planner (like a GFlowNet). A system cannot make rational decisions if its perception of reality is 70% wrong. This graph mathematically validates our intuition: the model has learned the "noise" of the trajectory, but it remains blind to the "precision" of the gesture.

Fig 6: Non-Linear Clustering Analysis (t-SNE)

While PCA looks for global variance, t-SNE focuses on local structures and manifold geometry. This visualization is where the model's struggle with generalization becomes most apparent.

What is t-SNE?

Before diving into the visualization, let’s define the tool. t-Distributed Stochastic Neighbor Embedding (t-SNE) is a non-linear dimensionality reduction technique specifically designed to visualize high-dimensional data.

Unlike PCA, which is linear and focuses on preserving the global spread (variance) of the data, t-SNE is "obsessed" with local structure. It works by converting the distances between points in the high-dimensional latent space into probabilities of being neighbors. It then tries to recreate those same neighborhood relationships in a 2D map.

In short: if two points are close on a t-SNE plot, they are very likely to be "neighbors" in the model's mind. In robotics, we use it to see if the model naturally groups similar physical states (like "closed grip") into distinct clusters.

1. The "Spaghetti" Syndrome: Episode Overfitting

The most striking feature of this graph is the presence of distinct, colorful "curves" or filaments. Each curve represents a single episode from the training set.

• The Diagnostic: Instead of grouping similar physical states together (e.g., all "gripper closing" moments in one cluster), the model has isolated each episode into its own trajectory. This is a clear sign of overfitting: the JEPA has memorized the specific "flavor" of each recording rather than extracting the universal laws of the robot's physics.

2. Latent Stability: Encoder vs. Predictor

There is, however, a silver lining. Notice how closely the Encoder points (blue) and Predictor points (orange) overlap within each filament.

• The Good News: The internal "dialogue" between the encoder and the predictor is highly stable. The predictor has successfully learned to mimic the encoder's latent transitions. Within its own "hallucinated" world, the model is consistent; it’s just that this world is currently a collection of 12 separate stories rather than one unified physical reality.

3. Implications for GFlowNet Planning

For a GFlowNet to function as a System 2 (reasoning) layer, it needs an environment where it can jump between states based on logic, not just follow a pre-recorded path. The current "blurred" boundaries and dispersed clusters confirm that the latent space is not yet "navigable."

Conclusion: The model is an excellent "historian" of its 12 episodes, but a poor "physicist." To break these filaments and force the model to merge these trajectories into meaningful semantic clusters, we need to flood the system with the Massive Diversity planned for Run 3.

Nota Bene: The Multi-View Experiment

A second training run was conducted by doubling the dataset size through a multi-view approach (using the same episodes but captured from two different camera angles).

After a thorough analysis, I decided not to dedicate a full article to this second run. The results were strictly identical, exhibiting the same patterns and limitations as the first run. Adding another camera angle on the same 12 episodes did not provide the "physical diversity" the model craved.

Key Takeaway: The 10-Epoch Convergence Interestingly, I compared the metrics at Epoch 10 versus Epoch 40 during this second run. The results were slightly better at Epoch 10. This confirms the findings of the original LeWorldModel (LeWM) paper: this architecture converges extremely fast, typically within 10 epochs.

The Pivot for Run 3 Moving forward, we will stop "wasting" compute on long training cycles. For all future experiments, we will limit training to 10 epochs. Instead of depth, we will focus on breadth: drastically increasing the number of episodes by building a massive, merged dataset from various heterogeneous sources.

For my POC on the Koch arm (SO-ARM101), I’m using the LeRobot ecosystem. It’s a goldmine of data, but its default storage format isn't built for the intensive training cycles required by World Models. Here is why I had to build a technical "bridge" to the HDF5 format.

🔗 Bridge repository

1. The Format War: Storage vs. Training

The Problem: Training a World Model (JEPA) requires sampling random time windows across thousands of episodes. Attempting a seek in an MP4 file for every single frame in a batch of 256 is performance suicide. HDF5 allows us to treat the dataset as one massive tensor living on the disk.

2. The "Little Mac" Challenge: Optimize or Die

Fetching datasets via the lerobot library is a breeze. The real challenge began during conversion. My Mac, with its limited resources, suffered several Kernel Panics before I got it right.

To succeed without saturating the RAM, I had to implement a "Lean & Mean" pipeline:

The "Low-Memory" Conversion Strategy

• Linear Pipeline: I abandoned aggressive parallelism. We process one episode at a time, one camera at a time. It’s slower, but it’s predictable.

• Micro-Batching (64 frames): Instead of loading an entire episode, we decode and write to the HDF5 in small chunks.

• Video Streaming: Using iterators (PyAV/OpenCV) to ensure a full video never materializes in the RAM.

• LZF Compression: The perfect compromise. It's ultra-fast, CPU-light, and significantly reduces the final file weight.

Safety and Reliability (Diagnostic Mode)

Because a 2-hour conversion crashing at 99% is unacceptable, I integrated several safeguards:

• Pre-validation: We check the integrity of metadata (episodes, frames, flags) before touching the video files.

• Watchdog & Heartbeat: If the script shows no progress for 120 seconds, it "fails-fast" rather than wasting heat.

• RAM Estimation: The script calculates and displays the estimated memory footprint of the batch before starting.

3. Alignment for LeWM (World Model)

The LeWM model is demanding. Conversion isn't enough; we need adaptation:

• Resize 224x224: The standard for modern vision backbones. Resizing is done on-the-fly during conversion.

• Key Normalization: LeRobot names columns one way, while LeWM expects another (e.g., pixels, action, state, done). My bridge handles the translation automatically.

4. The Golden Rule: No "Dirty" Data In low-cost robotics, datasets are often imperfect—truncated episodes or missing done flags are common.

My Policy: dirty_episode_policy=fail

On small datasets, the model is extremely sensitive to overfitting. Introducing inconsistent trajectories or ill-defined episode endings condemns the model to learn nonsense. I would rather have a converter that refuses to work than one that produces toxic data.

Tanguy's Advice

If you attempt this: watch your HDF5 chunks. A mismatch between your chunk size and your conversion micro-batches can turn your hard drive into a massive bottleneck. Yes, I learned this the hard way—remember, I only have a little Mac!

Next Step: We launch the training on an RTX 4090 runpod instance and see if our latent space survives physical reality.

1. The Current Deadlock: The Limits of VLA (Vision-Language-Action)

Today, state-of-the-art robotics relies heavily on VLA models. The concept is straightforward: give the robot a text instruction and an image, and it generates an action. However, this approach hits two major walls:

• The "Weight" Problem: Integrating a Large Language Model (LLM) into the control loop makes the embedding extremely heavy. We end up with models boasting billions of parameters just to decide whether to squeeze a 2cm gripper.

• Physical Hallucination: Language is discrete and symbolic; physics is continuous and unforgiving. By relying on LLM-style architectures, VLAs are prone to hallucinations: the robot "thinks" it has successfully grabbed an object because the statistical probability of the next word is high, even as the object slips away.

"Even when not 'complete,' VLAs rely on massive language model backbones. They inherit the bloat and probabilistic nature of LLMs, whereas JEPA offers a compact architecture dedicated solely to world dynamics."

2. What is the JEPA Architecture?

The JEPA (Joint-Embedding Predictive Architecture), championed by Yann LeCun, proposes a shift away from the "reconstruction" paradigm and toward "understanding."

The Fundamental Intuition

Instead of predicting pixels—like a generative model wasting energy trying to reconstruct the exact reflection of a light bulb on a table—JEPA predicts abstract representations.

The model isn't trying to draw the future; it’s trying to comprehend it..

Architecture inspired by "LeWorldModel: Stable End-to-End JEPA from Pixels" (Maes et al., 2024).

This schema is licensed under CC-BY 4.0. You are free to share and adapt it for any purpose, even commercially, as long as you keep the watermark or credit the original article.

1. Encoding Images into Latent Vectors

The process starts with a shared Encoder (e.g., a ViT-Tiny). Its job is to condense a high-dimensional 224 x224 image into a compact [B, 128] latent vector. We encode both the initial observation at Step t and the "target" observation at Step t+k. This compressed representation is what we call the Latent Space.

2. Preventing Model Collapse via Statistical Regularization

Older JEPA architectures often suffered from "model collapse," where the encoder would output the same constant value for every image to artificially minimize the loss. To prevent this, LeWorldModel (LeWM) introduces SIGReg (Statistical Isotropic Gradient Regularization). This forces the Latent Space to follow an Isotropic Gaussian distribution:

• Mean = 0

• Variance = 1 in every dimension/direction.

3. Predicting the Future

We feed the regularized Latent Space z_t and the Action Tensor (shape [B, k, n]) into a Transformer-based Predictor. The goal is to predict the future state in the latent space:

$$\hat{z}_{t+k}$$

4. The Optimization Challenge

The training involves minimizing two competing losses:

1. Consistency Loss (PredLoss): Ensures the predicted latent matches the real encoded future.

2. SIGReg Loss: Prevents collapse by enforcing the Gaussian distribution.

The core challenge lies in this tension: To achieve good predictions, the encoder wants to "group" similar physical states together. However, SIGReg constantly tries to "spread" the points out to satisfy the statistical constraint. This "struggle" is exactly what defines the geometry of our model's internal world.

3. Critical Analysis

Advantages

• Efficiency: With only 15M parameters, LeWM outperforms giant models in planning speed (up to 48x faster).

• Intuitive Physics: The model naturally learns concepts like gravity and object permanence.

Disadvantages

• Black Box: It is impossible to "see" the prediction without adding a third-party decoder.

• Short Horizon: Error accumulation in autoregressive mode limits long-term planning.

4. The GFN > JEPA Intuition: Marrying System 1 & 2

The major hurdle for JEPA in robotics is reaching distant goals. My intuition is to use GFlowNets (GFN) to decompose the task:

• JEPA (System 1 - Instinct): Handles local physics and immediate execution within the latent space.

• GFN (System 2 - Reason): Explores the diversity of possible trajectories and defines intermediate sub-goals that JEPA can then easily reach.

My Roadmap: From Signal to Sense

To test this intuition, my work is divided into two major experimental phases:

Phase 1: World Model Stress-Test

I will begin by training a JEPA (based on the LeWM architecture) on specific Low-Cost robotics datasets (LeRobot ecosystem, Koch SO-ARM101 arm). The goal is to study the "scaling laws" of the representation:

• Data Density: What is the impact of increasing the number of episodes?

• Multi-View: Does the model become more stable by synchronizing multiple angles (e.g., laptop camera + mobile camera) for the same task?

• Generalization: How does the model react to non-robotic videos or different robotic arms?

Phase 2: GFN-JEPA Alignment

If the latent space proves robust, I will tackle the "bridge" to deliberate planning. We will seek to define the limits of GFlowNets and align their workspace with JEPA’s. The end goal: a user provides a textual command, and the GFN translates it into a sequence of latent sub-goals for the JEPA to execute.

Follow the Progress

This project is an Open-Science exploration. I will publish results, failures, and breakthroughs in this serie.

A Note from the Author

This log is as much about the "how" as the "why." If you're interested in the intersection of latent dynamics and physical embodiment, stay tuned.