Teaching a Robot to Sort.

What began as a small imitation-learning exercise turned into a useful lesson about what a robot learns, what it skips, and how moving a few bins can reveal the difference.

Behavior Cloning PyTorch MuJoCo Multimodal Inputs Vision-Guided Control VLA Direction
Final closed-loop policy 06 / randomized success
After training with more varied examples, the robot can sort the cube even when the bins have moved.
How It Started

Motivation

I started this project because I wanted to learn more about modern robotic learning, like imitation learning, behavior cloning (BC), and VLAs. If you're reading this, I hope this project can help inspire and guide you through your own learnings, I certainly learned a lot that classrooms didn't teach.

This looked like a small sorting problem.

The idea was simple: for a 'waiter' robot holding a colored cube on a tray, train some policy to sort into the matching bin. Of course, this is a well-studied problem and is readily solved with a state machine or basic rule-based planning. That was the point. So I built the task in MuJoCo on top of my PhD work, and began writing the 'expert' or 'oracle' policy. I'll give a very abbreviated summary of Behavior Cloning below.

What is Behavior Cloning? Behavior Cloning (BC) is one of the simplest imitation learning techniques where a policy is trained to mimic an expert's actions. The expert can be a human or a scripted policy, and the learner tries to replicate the expert's behavior based on observations and actions. In this case, we use an expert that knows the environment and makes no mistakes (unless we make it).
Task Setup

Setting up the expert.

In each attempt, the camera sees a red or blue cube resting on the robot's tray, the robot itself, and the two trays. The camera is fixed in the environment; I chose an isometric view at first, but eventually decided on a more head-on view when we started domain randomization, which you'll see in the remaining videos.

Fixed camera view

Before any training, I had to create the scripted expert that already knew where everything was. In many posts and videos you'll see online today, people use the SO-101 arms and train basic tasks with hands-on teleoperation. I didn't want to be locked into a particular robot embodiment, and I wanted to test how a simulated expert would be able to teach a policy.

The scripted expert is just a state machine that does the following, using intrinsic information available in MuJoCo. It chooses the correct bin, carries the tray over, tips the cube out, and returns home. This version is less interesting as intelligence, but very useful as a patient teacher: it can repeat the same task over and over while I record examples for the learned policy.

Observation

Vision + robot state

24-Dimension state vector. RGB images from a fixed camera, plus joint state, force/torque, end-effector pose, and object pose.

Data Collection

Trials / episodes

Each episode lasts roughly 5 sec and the recording stride was set to 5 (observe every 5th frame) for efficient data collection.

Action

Δθ Joint deltas

The expert records absolute joint targets, but the learner predicts normalized joint deltas for stability.

Why this is not yet a full VLA

I began by studying lightweight open-source VLA models such as SmolVLA and LeRobot-style pipelines. For this first implementation, I scoped the problem to behavior cloning with multimodal inputs. The policy consumes images, robot state, and simplified task information, but it is not yet a general language-conditioned vision-language-action model.

The scripted oracle policy

The oracle knows where everything is. It reads the cube color from MuJoCo state, chooses the corresponding bin, carries the tray to a pre-defined pose, and tips the cube into the bin, then returns home. This version is less interesting as intelligence, but very useful as a patient teacher: it can repeat the same task over and over while I record examples for the learned policy.

01
Reset episode
Randomize cube color and reset the robot, tray, cube, and bins.
02
Read color
Oracle reads the environment's cube-color attribute directly (no vision).
03
Move to pre-drop
IK drives the tray to a hardcoded pose near the selected bin.
04
Tilt tray
The wrist tilts to dump the cube into the matching bin.
05
Record rollout
Images, state, instructions, actions, color, and success labels are saved.
01 / Expert basic sort
The scripted expert completes the basic sort using privileged simulator state. This is the behavior the learner is asked to imitate.

Rollout collection.

The first dataset contains 100 successful episodes from the expert. We don't record every single physics step. Instead, we sample at a lower rate to keep the dataset manageable without losing the shape of the movement.

Signal Shape / Example Purpose
Images Fixed 128x128 RGB camera Visual context for cube and bin geometry and poses.
State q, qdot, ft, ee_pos, obj_pos Robot proprioception and simulator-derived task state.
Action Expert joint targets → converted to Δθ Closed-loop command target for the learned policy.
Cube color red / blue Compact task label for ground truth and diagnostic variable for training.
Success Boolean episode outcome Tracks whether the matching bin received the cube.
Implementation notes

I initially recorded absolute joint position targets because that is what the expert command produced. During training, I converted these to joint deltas. This made normalization easier and reduced unstable closed-loop behavior during evaluation.

For speed, the recorder uses a stride over physics steps — for example, record_stride=20 yields roughly 50 Hz depending on the simulator step size. RGB rendering is also skipped on steps that are not being recorded.

Behavior Cloning

The learner doesn't get to cheat.

The learned policy doesn't get to read the cube's color and every pose straight from simulator state. It only has a camera image and its own proprioception, and has to recover the same decisions, such as which bin, which tray position, and how far to tip. It combines what it sees with information about the robot and the current task, and from that predicts a small change for each joint.

No memory across time These are plain behavior-cloning networks: each action is regressed straight from the current observation. There's no LSTM, attention, or other explicit temporal memory — nothing recurrent to carry state between steps.
First Attempts

The early runs were a little clumsy.

The first learned policy kept dropping the cube between or on the bin edges, regardless of color (hint: other forms of policy learning like diffusion policy resolve these issues).

What actually fixed it: missing regardless of color pointed at the image: the adaptive pool was collapsing the vision features down too aggressively, so the fused representation was effectively dominated by the other input features while the image contributed almost nothing. Widening the pooled output to a 2×2 grid gave the image roughly 200x more relative weight — much closer to what it should have been, and the policy started actually using what it saw.

02 / Behavior cloning first failure

The learner misses

The cube lands short, long, or on the lip of a bin — the policy cannot decide what to do correctly.

03 / Behavior cloning basic success

The baseline works

Same task, same layout, after fixing how much the policy actually looked at the image.

A promising first result Completing the fixed layout showed that the policy could imitate the demonstrations. I still wasn't sure whether it was looking for the bins or simply replaying a familiar motion. Notice how the video on the right has bins swapping places, but not in random positions.
A Small Experiment

So I moved and randomized the bins.

This was the simplest test I could think of. If the policy had learned to find the matching bin, a new layout should be manageable. If it had memorized "red means move this way," the same change would trip it up. It tripped up. So I collected a new set of demonstrations with the bins placed at different reachable positions around the robot, taking care to keep them separate and sensibly oriented. Check out the new expert below:

04 / Expertrandomized bins
The expert demonstrates the same task across changing layouts, making bin position part of the data rather than a constant.
Randomization details

The randomized-bin version samples each bin on a radius around the robot base and chooses an angle within the reachable workspace. Bins are rotated to be radially symmetric, the pre-drop pose is moved higher in z, and the tray tilts away from the robot along an axis tangent to the circle around the base.

What Happened

More varied practice made the difference — up to a point.

The original policy failed as soon as the bins left their usual spots. Re-collecting and re-training on demonstrations with randomized bin placement did not immediately fix that case: training loss was low, but closed-loop behavior on genuinely new layouts was not.

Basic BC's hidden assumption Basic BC worked with fixed bins because the task goal was effectively constant. Randomizing bin placement exposed a hidden assumption: the policy could imitate motion, but it had no reliable representation of where it was supposed to go. That's what moved me to goal-conditioned BC.
Low loss isn't the same as working The retrained policy matched the expert's actions closely during training. But BC doesn't guarantee good closed-loop behavior — small per-step errors compound once the policy is driving itself instead of being scored against a fixed demonstration.

The fix was to stop asking the policy to infer the goal from pixels alone and just tell it: I implemented goal-conditioned behavior cloning, feeding in the selected bin's pose — x, y, and yaw — alongside the image and robot state.

05 / Original BCrandomized failure

The shortcut breaks

Changing the bin geometry exposes what the fixed-layout policy did not learn.

06 / Goal-conditioned BCrandomized success

Conditioning on the goal works

Given the selected bin's pose directly, the policy completes the sort across continuous, random layouts.

Next Steps

There is plenty left to try.

This is still a fairly modest behavior-cloning project, which is part of why I like it. The next step is to give the robot richer visual and language inputs, then keep asking the same basic question: is it responding to the scene in front of it, or has it found another convenient pattern in the training data?

Perception

Wrist camera

Compare fixed-camera, wrist-camera, and multi-camera policies to study the cost and benefit of additional visual tokens.

Robustness

Domain randomization

Randomize bin placement, camera pose, lighting, object pose, and colors to reduce memorized shortcuts.

Learning

VLA fine-tuning

Move from simplified instruction/color conditioning toward a lightweight language-conditioned action model.

Open questions I am still testing How much bin randomization is enough before the policy must use vision? Does a wrist camera improve grounding or simply increase training cost? How should action chunking or flow-matching action heads be compared against a compact MLP baseline? These questions are the motivation for keeping the behavior cloning setup simple and diagnosable first.

Project Links

Code, notes, and demos.

This is a work in progress, and I expect the project—and this writeup—to keep changing as I try new cameras, layouts, and learning approaches.

GitHub ↗ Back to Projects