Beyond the basics

In this write-up, I share what took my policy from roughly 25% success to over 95%, with steerability and strong generalization. The task is sorting laundry with a ReBot arm: socks go to a basket, shirts go to a bin. The base model is MolmoAct2, a 5B VLA, fine-tuned on one GPU.

It was an iterative process with many changes, but three key factors made the biggest difference.

Data diversity. A robust policy that generalizes and can be steered needs diverse training data. It is tempting to think that post-training only on one’s own dataset is best, but 5B-parameter VLAs post-trained with the current techniques on a non-diverse dataset seem to suffer from memorization and collapse the representations learned during pre-training (Schwager, 2026). My observation was that after assembling a 404-episode, 120-task, 27-environment, 5-robot dataset and co-training it with my own data, I got the strongest results on my own task. I deliberately selected only a small number of episodes per task, aiming for a high task-to-episode ratio. The goal was to expose the policy to many different tasks, objects, and environments within a compact training set. The dataset is available here.

Steerability via subtasks. One of the appeals of steerability is that it raises the level of abstraction: for a new task, instead of having to teleoperate, data can be collected by instructing the robot with subtasks. Steerability is also good on its own, since it gives more control over how the model behaves. This is one of the key capabilities discussed in the π0.7 paper (Physical Intelligence et al., 2026). My observation was that steerability came from dataset diversity combined with VLM subtask annotations, some of which were validated by a human. Without the diverse data, the policy would ignore subtask commands or even tasks commands.

The bf16 trap: weight stagnation. Training in bfloat16 (bf16) is appealing because it allows larger models to be trained faster. The cost is reduced precision: when a weight update is smaller than the smallest change bf16 can represent, it is rounded away and the weight does not move at all. At a given learning rate, this can leave most of the model effectively frozen. This was silently affecting my post-training runs far more than I realized. This issue is well documented in the LLM literature (Ozkara et al., 2025; Rae et al., 2021; Zamirai et al., 2020). At a typical learning rate of 5e-5, only 23% of the trainable weights changed over 56 steps of AdamW. With stochastic rounding, 73% did.

The rest of this write-up goes through each of these in detail.

Data diversity

With only my local data, training loss kept decreasing while validation loss started rising after a few hundred steps. That is memorization, and with 46 episodes and a 5B model, it is unsurprising.

To address this, I incorporated the diverse dataset directly into post-training by filling half of each batch with diverse data and the other half with my local demonstrations. This prevented the same memorization pattern. Validation loss reached a lower minimum than the local-only run (0.063 against roughly 0.071), and then stayed there: it remained between 0.063 and 0.070 for the rest of training, up to 2,000 steps. Training loss settled higher than in the local-only run, which is expected: half of each batch now comes from a different distribution, so the model cannot fit the local demonstrations as closely.

The mixture also helped a run that had already memorized. I resumed from the local-only checkpoint at step 1,200 and switched the diverse half on. Within about 200 steps, validation loss dropped from above 0.08 to 0.064, nearly matching the best value of the run that was mixed from the start (0.063). It drifted up afterwards, but 1,000 steps later it was still around 0.071, no worse than the best the local-only run ever reached. I did not expect memorization to be this reversible.”

The diverse dataset

The dataset has 404 episodes and about 4.7 hours of recordings from five public sources. Every source has a wrist RGB camera, because the policy is conditioned on one. FMB also supplies wrist depth from the same D405 camera I have.

Learning from suboptimal data. Following π0.7 (Physical Intelligence et al., 2026), I added quality, speed, and mistake metadata to the diverse dataset and passed it to the policy as text. This gives the model the context needed to distinguish different levels of performance during training, while allowing us to request high-quality, fast behavior without mistakes at inference.

A VLM reviewed the data first, and it flagged episodes for me to review. I dropped two kinds of training samples: those where the robot does not move during the upcoming action window, since some datasets contain idle stretches, and chunks where I did not think any useful behavior would be learned.

Three details mattered when making that data usable:

• Set the mixture explicitly. Half of each batch comes from local data and half from the diverse dataset, so that training still emphasizes my own data. Sampling from everything pooled together, without fixing the ratio, might be an experiment worth running.
• Normalize per action layout. ReBot records joints in degrees; every other source uses radians. A single pooled normalization would have compressed the entire diverse dataset into about 1% of the range. Whether a source records commanded targets or measured states also matters, so DROID and FMB are normalized separately even though both use a Franka arm.
• Keep useful failures. I kept mistakes and recoveries rather than filtering them out. Each is labelled, and the label is passed to the policy as text. A failed attempt can still contain useful behavior, particularly in the attempt that follows it. The clip below shows what this looks like on the robot.

The sampling and normalization details are in the appendix.

Generalization

The clearest effect of the diverse data was on objects and instructions outside my own recordings. I asked the policy to “grasp the cup from the rim.” Neither a paper cup nor that phrase appears in my local training data. It took a few attempts, then the policy grasped the rim and carried the cup to the basket.

This shows transfer beyond my local examples. It does not establish that the pretrained model had never seen a cup.

Steerability via subtasks

Steerability in this context means that the policy follows both a task and subtasks that specify how to execute it. Each subtask is atomic: one action, written as verb → object, with a destination when needed. Examples are “grasp the black sock”, “move the black sock to the basket”, and “place the object on the fixture”. Grasping, moving, and releasing are separate atoms. These atoms can be combined into tasks never seen during training, a capability I saw glimpses of. This raises the level of abstraction.

For a new task, instead of teleoperating, one can collect data by instructing the robot step by step, and the robot does the low-level work. Steerability is also worth having on its own: it is direct control over how the model behaves. π0.7 makes the same point (Physical Intelligence et al., 2026). Its authors describe the prompt’s conditioning information as what “makes it possible to steer the model precisely to perform many tasks with different strategies”, and they pair that with training on very diverse data.

My observation was that steerability came from dataset diversity combined with subtask annotations. The diverse dataset was annotated with one-action subtask clauses generated with a VLM and validated by hand for a subset. The two clips below show examples of the policy following different requested subtask sequences, including both the requested object order and destination.

In the next two clips, I moved the basket to the left of the table, a position absent from the local training data. In the first, I asked for the black sock to go to the basket, and the policy placed it there. In the second, I asked for it to go to the bin behind the basket, and the arm carried it past the basket and into the bin. The sock-to-bin pairing was also absent from my demonstrations.

These runs change both where a container is and which container an object should go to. They are more informative than another successful sort with the usual layout.

The bf16 trap

Bfloat16 (bf16) is appealing because it allows larger models to train faster while using half the memory per value of fp32. It retains fp32’s exponent range, but at the cost of precision. The spacing between representable values increases with their magnitude:

Near 0.01, the spacing is already comparable to a typical fine-tuning learning rate of 5e-5; near 0.1, it is about ten times larger. When weights are updated directly in bf16, an update smaller than half the spacing rounds away, leaving the weight unchanged. For example, an update of 5e-5 cannot move a weight near 0.1. The actual update depends on the gradients and optimizer state, not just the learning rate. This can fail silently: loss may decrease even while many weights remain unchanged.

This issue is known in the LLM literature (Ozkara et al., 2025; Rae et al., 2021; Zamirai et al., 2020). I had not realized how much it was affecting my runs and am sharing the measurements in case they help others.

How bad it was in my runs

Each run in the sweep starts from the foundation checkpoint and trains for 64 steps at a fixed learning rate. The measurement is the percentage of trainable bf16 weights whose stored values differ between steps 8 and 64, comparing ordinary AdamW with stochastic-rounding AdamW.

At 5e-5, 23% of the weights differed with ordinary AdamW, compared with 73% with stochastic rounding. At 5e-6, the final learning rate of the cosine schedule, those figures were 5% and 41%.

Stochastic rounding

Stochastic rounding chooses between the two neighboring representable values with probabilities that preserve the unrounded value on average (Ozkara et al., 2025; Zamirai et al., 2020). Small updates therefore have a chance to change the stored weight instead of being consistently rounded away.

I ended up using a fixed learning rate of 5e-5 with stochastic rounding and no weight decay. After training, 99.9% of the trainable weights differed from the original checkpoint. This comparison covers the full run; the sweep above measures changes only between steps 8 and 64.

On the robot, the stochastic-rounding checkpoint followed subtasks better but grasped less reliably than its AdamW counterpart. The runs also differed in learning-rate schedule and weight decay, so the behavior difference cannot be attributed to rounding alone.

The clips in the steerability and generalization sections use a linear blend of 60% stochastic-rounding checkpoint and 40% local-only checkpoint. This is also the policy used for the success rate reported in the introduction.

Appendix: setup and data

The arm is a ReBot B601, with seven joints counting the gripper. The policy receives an external RGB view, wrist RGB and depth, joint state, and a text prompt. The model is MolmoAct2, a roughly 5.4-billion-parameter vision-language-action model, fine-tuned in full in bf16 on one workstation GPU.

The local-only run uses 46 episodes, about 181,000 frames. The later local corpus has 61 episodes and about 230,000 frames at 30 Hz, or 2.1 hours. It includes 15 recorded rollout episodes as well as the demonstrations. Episodes carry subtask annotations, quality grades from 1 to 5, a mistake flag, and speed labels where available.

Diverse dataset sampling

Within the diverse half of a batch, I sample a source, then an episode uniformly, then an anchor within the episode. Source weights are proportional to the square root of their episode counts. The resulting shares are about 35% RoboChallenge, 25% FMB, 18% each DROID split, and 5% UR7e. This keeps long episodes and large sources from dominating simply because they contain more frames.

Each sample predicts a one-second chunk of 30 actions at 30 Hz, expressed relative to the joint state at the start of the chunk. Lower-rate action sequences are interpolated onto that grid. State and action dimensions are padded with masks, and cameras are mapped to consistent roles. Missing cameras are masked out.

Normalization is separate for each action layout. RoboChallenge has no commanded action stream in this corpus, so its measured joint state is used as the target and kept under a separate layout. The diverse sources are used for training only; evaluation is on the local task.

Validation loss

The loss plots report flow-matching loss on 128 held-out frames, evaluated every 40 training steps. Validation episodes are excluded from the corresponding training run. Training curves use a 200-step smoothing window.

The local-only run uses four held-out demonstration episodes. The later validation set contains three of those episodes and a policy rollout. On identical weights, that change accounts for about a 0.005 shift at the branch point in the first loss plot. The later training corpus also adds rollouts and revises annotations. The within-run trends are therefore easier to interpret than a precise difference between the curves.

Weight updates

The short optimizer sweep starts each run from the foundation weights with fresh optimizer state and a fixed learning rate. The reported percentage counts trainable bf16 elements whose stored values differ between steps 8 and 64. Frozen parameters are excluded. The 5e-3 runs are missing and are marked as such in the figure.

The comparison is PyTorch AdamW versus this project’s stochastic-rounding AdamW. The latter also computes the final update in fp32 and combines decay with the adaptive update before rounding. It is an implementation comparison, rather than an isolated change to one rounding operation.

For the longer training runs, both methods use the same mixed dataset, batch size of 128, and peak learning rate of 5e-5. Plain AdamW uses a cosine schedule down to 5e-6 and weight decay of 0.1. The stochastic-rounding run keeps a flat learning rate and turns weight decay off.