xBerry Robotic Arm RL Week 3: Camera Training and Asymmetric A2C Architecture

After weeks of training on perfect numerical data, we removed the cheat sheet. From now on, the robot has to learn from what it sees.

TL;DR

Week 3 marks the transition from Phase A (proprioception-only) to Phase B: training with raw camera images as input. We used the Critic’s weights from Phase A as a warm start, built an asymmetric A2C architecture where the Actor sees only pixels and joints while the Critic retains access to privileged data, and hit 96% success rate (48/50 trials) on ball transfer in simulation.

From Numbers to Pixels

Phase A gave us a solid foundation: a policy trained on perfect numerical state information – exact joint positions, ball coordinates, bowl position – running across 4,000+ parallel simulation environments. That policy reached near-perfect performance in simulation. But there was a catch: on a real robot, most of that privileged data simply does not exist. The actual SO-101 arm does not receive exact ball coordinates during inference. It sees what its cameras see.

Phase B closes that gap. The goal is to train a policy that can operate from camera frames alone – the same information stream available during real-world deployment.

The Warm Start Trick

Switching from numbers to pixels sounds like starting over. It’s not and that distinction matters enormously for training time.

We transferred the Critic’s weights directly from Phase A into Phase B. The Critic – the network component responsible for evaluating how good a given action is – already knows what good robot behavior looks like. It spent all of Phase A learning to score actions based on rich state observations. By warm-starting Phase B with this knowledge, the Actor (the actual motion controller) gets reliable feedback signals from the very first training step, rather than learning from a Critic that itself is still fumbling in the dark.

Why this matters: Training vision-based RL policies is slow and expensive. Without a warm start, the Critic would need thousands of steps before its action evaluations became meaningful. The warm start compresses that calibration cost to near zero.

Asymmetric A2C: One Policy, Two Different Realities

The central architectural challenge of Phase B is that the Actor and the Critic need to receive fundamentally different types of information.

A2C (Advantage Actor-Critic) – a reinforcement learning algorithm where two networks collaborate: the Actor decides what action to take, and the Critic evaluates how good that action was.

During inference on a real robot, the Actor is the only component that runs. The Critic trains alongside it but never deploys. This asymmetry creates an opportunity: the Critic can be fed privileged information – exact ball position and bowl position expressed as quaternions (x, y, z, w) – that would be unavailable in the real world, while the Actor learns exclusively from what the real robot will actually have access to.

The Actor receives: camera embeddings (compressed visual representations from a CNN encoder) and joint positions from the robot’s servo encoders.

The Critic receives everything the Actor sees, plus privileged observations about the exact 3D positions of both the ball and the bowl.

How We Taught the Robot to See: The CNN Encoder

Raw camera frames cannot flow directly into a policy network. A single 480×640 RGB image contains hundreds of thousands of pixel values – feeding that unprocessed into a neural network is computationally impractical and statistically noisy.

CNN (Convolutional Neural Network) – a type of neural network designed to process spatial data like images, detecting features through learned filters applied across the image.

Our PixelEncoder consists of 4 convolutional layers followed by SpatialSoftmax — a layer that computes the expected 2D position of each learned feature across the image, effectively turning a feature map into a set of spatial coordinates. A final linear head converts those coordinates into a compact embedding vector.

Embeddings – compact numerical vectors representing an image in latent space. Similar images produce similar embedding vectors, helping the policy generalise across different visual conditions.

class SpatialSoftmax(nn.Module):
    def __init__(self):
        super().__init__()
        self._grid_cache: dict[tuple[int, int, torch.device], tuple] = {}

    def _grids(self, h: int, w: int, device: torch.device):
        key = (h, w, device)
        cached = self._grid_cache.get(key)
        if cached is None:
            xs = torch.linspace(-1, 1, w, device=device)
            ys = torch.linspace(-1, 1, h, device=device)
            y_grid, x_grid = torch.meshgrid(ys, xs, indexing="ij")
            cached = (x_grid.reshape(-1), y_grid.reshape(-1))
            self._grid_cache[key] = cached
        return cached

    def forward(self, feat):
        B, C, H, W = feat.shape
        x_grid, y_grid = self._grids(H, W, feat.device)
        w = torch.softmax(feat.reshape(B, C, H * W), dim=-1)
        x = (w * x_grid).sum(-1)
        y = (w * y_grid).sum(-1)
        return torch.cat([x, y], dim=-1)


class PixelEncoder(nn.Module):
    def __init__(self, in_channels=3, feature_dim=128):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, stride=1),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, stride=1),
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, stride=1),
            nn.ReLU(),
        )
        self.spatial_softmax = SpatialSoftmax()
        self.head = nn.Sequential(
            nn.Linear(2 * 32, feature_dim),
            nn.LayerNorm(feature_dim),
            nn.Tanh(),
        )

    def forward(self, x):
        return self.head(self.spatial_softmax(self.conv(x)))

The result: during inference, raw camera frames are transformed into a fixed-size vector the Actor’s neural network can reason about. The encoding adds minimal latency and enables the entire policy to run from pixels alone.

The Memory Ceiling

The transition to camera-based training surfaced a hard physical constraint: GPU memory.

Phase A ran comfortably at 4,096 parallel simulation environments simultaneously. Phase B camera observations require storing image tensors for each environment at every training step. At 4,096 environments that exhausts GPU memory almost instantly. The solution was trial and error: systematically reduce the environment count until training runs without OOM (out-of-memory) errors. We landed at approximately 64 parallel environments – a 64× reduction from Phase A.

This significantly slows per-step throughput, but the warm-started Critic compensates by making every gradient step count more from the start.

Results: 48 out of 50

After training Phase B with the camera-based asymmetric A2C setup, we ran a 50-trial evaluation. The policy succeeded on 48 of those trials – 96% success rate on the ball transfer task.

Notably, this matches Phase A performance despite the policy now receiving only camera frames and joint positions during evaluation. The warm-started Critic, armed with privileged state data during training, gave the Actor enough signal to develop a robust visual policy without ever needing perfect state information at inference time.

What’s Next

Week 3 ends with a simulation-validated checkpoint at 96% accuracy. Week 4 is deployment: the model ships to the physical SO-101 arm. The real question is no longer whether the policy works in simulation – it’s whether the visual features it learned generalise to real lighting, real cameras, and real physics.

FAQ

What is Phase B in robot reinforcement learning training?

Phase B is a training stage where the robot’s policy learns to operate from raw camera images instead of privileged numerical state data, directly preparing the policy for real-world deployment where exact object positions are unavailable.

What is an asymmetric A2C architecture?

Asymmetric A2C (Advantage Actor-Critic) is a neural network setup where the Actor receives only the observations available during real deployment (camera embeddings, joint positions), while the Critic additionally receives privileged information (exact object coordinates) to provide more accurate training feedback.

What is SpatialSoftmax and why is it used in robot learning?

SpatialSoftmax converts a convolutional feature map into 2D coordinate expectations for each feature channel, effectively locating key visual features spatially in the image. In robot learning, it helps the encoder detect and track task-relevant objects – such as a ball or bowl – across consecutive camera frames.

What is a warm start in RL training?

A warm start means initialising a new training run with weights from a previously trained network rather than random initialisation. In Phase B, the Critic’s weights from Phase A are reused so the Actor receives meaningful action evaluations immediately, substantially reducing overall training time.

Why does camera-based training reduce the number of parallel environments?

Each camera observation is an image tensor containing hundreds of thousands of values per environment. At 4,096 parallel environments the combined memory for all image tensors exceeds GPU capacity. Reducing to approximately 64 environments keeps memory usage within bounds while preserving training stability.

What result did the robot achieve in Phase B simulation evaluation?

The robot achieved a 96% success rate – 48 successful ball transfer trials out of 50 – using only camera frames and joint positions as input during evaluation, matching the performance of the proprioception-trained Phase A policy.