Skip to content

AI Engineering

I Taught a Robot to Pick Up a Cube — Without Owning a Robot

From zero-percent RL success to reliable grasping via imitation learning — all in a physics sim on a MacBook, orchestrated by an LLM.

 ·  12 MIN READ


Alexandre Agius

Alexandre Agius

AWS SOLUTIONS ARCHITECT

SHARE

I spent a week teaching a robot arm to pick up a cube. The robot doesn’t exist. It lives in a physics simulator on my MacBook, controlled by an LLM that has never touched anything physical. The arm is a SO-100 — a popular open-source manipulator — simulated joint-by-joint in MuJoCo with full contact physics. The brain giving it instructions? Claude Sonnet 4.5 running on Amazon Bedrock, wrapped in the Strands Agents SDK. No hardware, no lab, no budget. Just Python, a physics engine, and the question: can a language model actually control a robot?

The answer turned out to be more nuanced — and more interesting — than I expected. This post is the full journey: nine progressive scripts, three different learning approaches, one spectacular failure, and the architectural insight that ties it all together.

The Setup: Why Simulated Robotics, Why Now

Two things converged to make this project possible in 2026 that weren’t practical even eighteen months ago.

First, MuJoCo went open-source and got fast. DeepMind’s physics engine now runs real-time on a laptop CPU with sub-millisecond step times. You can simulate a 6-DOF arm with contact physics at thousands of steps per second. No GPU needed for simulation — just for training.

Second, agent frameworks matured enough to treat physical actuators as tools. The Strands Agents SDK introduced the concept of an AgentTool — a callable that an LLM can invoke with structured arguments, inspect results from, and reason about iteratively. The strands-robots package extends this to robotics: it wraps a MuJoCo simulation (or a real robot) as a single tool that an agent can call.

This is the Robot-as-AgentTool pattern, and it’s the architectural core of the project:

┌─────────────────────────────────────────────────┐
│  Strands Agent (Claude Sonnet 4.5 on Bedrock)   │
│                                                 │
│  "Move joint 3 to 1.2 radians, then close      │
│   the gripper"                                  │
└──────────────────────┬──────────────────────────┘
                       │ tool_call(action=[...])

┌─────────────────────────────────────────────────┐
│  strands_robots.Robot() — AgentTool             │
│  • Accepts joint targets or Cartesian goals     │
│  • Steps MuJoCo simulation                      │
│  • Returns observation dict (positions, forces) │
└──────────────────────┬──────────────────────────┘
                       │ physics step

┌─────────────────────────────────────────────────┐
│  MuJoCo Engine — SO-100 Arm + Scene             │
│  • Contact dynamics, gravity, friction          │
│  • Renders to viewer (optional)                 │
└─────────────────────────────────────────────────┘

The LLM doesn’t know anything about motor control, inverse kinematics, or PID loops. It just sees a tool that accepts actions and returns observations. The intelligence is in deciding what to do next given what it sees.

From Hello World to Agent Control

The project is structured as nine scripts that build on each other progressively. The first two establish the pattern.

Script 01 is a pure simulation sanity check — load the SO-100 URDF, step the sim, render a frame. No agent, no learning. Just “does MuJoCo work on my machine?”

Script 02 is where it gets interesting. This is the minimal Robot-as-AgentTool:

from strands import Agent
from strands.models.bedrock import BedrockModel
from strands_robots import Robot

# Create the simulated robot as a Strands AgentTool
sim = Robot(
    robot_id="so100",
    embodiment="mujoco",
    task="pick_cube"
)

# Wire it to an LLM agent
model = BedrockModel(
    model_id="us.anthropic.claude-sonnet-4-5-20250514-v1:0",
    region_name="us-west-2"
)

agent = Agent(
    model=model,
    tools=[sim],
    system_prompt=(
        "You control a SO-100 robotic arm in simulation. "
        "Use the robot tool to move joints and observe results. "
        "Your goal: pick up the red cube on the table."
    )
)

# Natural language → robot actions
agent("Pick up the red cube and lift it 10cm above the table")

That’s it. The LLM receives the tool’s observation schema, reasons about joint positions and gripper state, and issues sequential tool calls. It’s clumsy — the arm waves around, overshoots, drops things — but it works. The agent closes the loop between perception and action through pure language reasoning.

The limitation is obvious: an LLM making 200ms API calls per action step is comically slow for real-time control. But as an orchestrator deciding which policy to execute? That’s a different story. More on this later.

The RL Journey: Where Observation Design Beats Compute

Scripts 03 through 05 explore reinforcement learning. The task: reach the cube, grasp it, lift it. Standard manipulation benchmark.

I started with FastSAC (a sample-efficient variant of Soft Actor-Critic) and the naive observation space: joint positions, joint velocities, gripper state, cube position. Reward: negative distance from end-effector to cube, plus a bonus for lifting.

Result after 400,000 steps: 0% grasp success. The arm could reach the cube reliably but never closed the gripper at the right moment. It would hover, oscillate, or slam into the table.

The fix wasn’t more compute or a different algorithm. It was observation design.

I added a single 3D vector to the observation: the displacement from the gripper fingertips to the cube center. This is information the agent could theoretically compute from joint positions + cube position + forward kinematics, but making it explicit removed a computational burden the policy network was failing to learn implicitly.

# The observation that changed everything
obs = {
    "joint_positions": qpos[:6],        # 6 joint angles
    "joint_velocities": qvel[:6],       # 6 joint velocities
    "gripper_state": gripper_pos,       # scalar: 0=open, 1=closed
    "cube_position": cube_xyz,          # 3D world coords
    "grip_to_cube": cube_xyz - fingertip_xyz,  # THIS. This one vector.
}

Result after 400,000 steps with the new observation: 80% grasp success. Same algorithm, same hyperparameters, same compute budget. The only change was giving the policy network a pre-computed relative vector.

This is the first big lesson: for RL in manipulation, observation design matters more than compute. A well-engineered observation space can be worth orders of magnitude more than additional training steps.

In fact, when I doubled the training to 800,000 steps, success dropped to around 65%. The policy overfit to specific cube placements in the training distribution and lost generalization. More compute actively hurt.

ConfigurationStepsGrasp SuccessNotes
Naive observations400k0%Reaches cube, never grasps
+ grip→cube vector400k80%One vector = breakthrough
+ grip→cube vector800k~65%Overfitting, lost generalization

The Hard Truth: RL Can’t Grasp (Well)

Even at 80% reach-and-grasp, the RL policy had a fundamental problem: it couldn’t lift reliably. Getting the gripper around the cube is one thing. Closing fingers with the right force, maintaining grip during the lift acceleration, compensating for the shifted center of mass — these are contact-rich dynamics that RL struggles with.

The issue is structural. RL optimizes a scalar reward signal through exploration. But contact events are discontinuous — the reward landscape has cliffs. You’re either gripping or you’re not. There’s no smooth gradient from “fingers 1mm too wide” to “solid grasp.” The policy needs to discover a narrow region of action space (correct gripper closure timing, force, and approach angle simultaneously) through random exploration. For a 6-DOF arm with continuous actions, that’s finding a needle in a haystack.

This matches the broader robotics literature: RL works well for reaching, pushing, and non-contact locomotion. For contact-rich manipulation — grasping, insertion, reorientation — imitation learning dominates.

The Expert + Imitation Path

Scripts 06 through 08 take a completely different approach.

The IK Expert Controller

First, I built a perfect controller using inverse kinematics. No learning at all — just geometry:

  1. Compute the target end-effector pose above the cube
  2. Solve IK for the joint angles
  3. Move there with a simple trajectory
  4. Lower to grasp height
  5. Close gripper
  6. Lift
def expert_pick_and_lift(sim, cube_pos, lift_height=0.10):
    """Deterministic IK-based pick and lift. 100% success when tuned."""
    # Pre-grasp: approach from above
    pre_grasp = cube_pos + np.array([0, 0, 0.05])
    joints_pre = inverse_kinematics(sim, target_pos=pre_grasp)
    move_to(sim, joints_pre, steps=100)

    # Lower to grasp
    joints_grasp = inverse_kinematics(sim, target_pos=cube_pos)
    move_to(sim, joints_grasp, steps=80)

    # Close gripper
    close_gripper(sim, force=0.8, steps=50)

    # Lift
    lift_target = cube_pos + np.array([0, 0, lift_height])
    joints_lift = inverse_kinematics(sim, target_pos=lift_target)
    move_to(sim, joints_lift, steps=100)

    return check_cube_lifted(sim, lift_height)

Result: 12/12 successful pick-and-lifts. 100% success rate. No randomness, no exploration, no reward shaping. Just math.

But this expert is brittle — it assumes perfect knowledge of the cube position, a clean table, no obstacles, and no perturbations. It doesn’t generalize. That’s fine, because its job isn’t to be the deployed policy. Its job is to generate demonstrations.

From Expert Demos to Learned Policy

Script 07 records 50 expert demonstrations — full trajectories of observations and actions during successful pick-and-lifts with randomized cube positions. Script 08 trains an ACT (Action Chunking with Transformers) policy using the LeRobot framework from Hugging Face.

The key differences from RL:

AspectReinforcement LearningImitation Learning (ACT)
SignalSparse reward (lift bonus)Dense supervision (full trajectory)
ExplorationRandom — must discover graspingNone — learns from success only
Contact handlingStruggles with discontinuitiesLearns continuous trajectory through contact
Data neededMillions of sim steps50 demonstrations
Resulting policyReach: good, Grasp: fragileReach + Grasp + Lift: robust
Failure modeNever discovers graspSlight distribution shift on unseen configs
Training timeHours (CPU)30 min (GPU)

The imitation policy doesn’t discover novel strategies — it reproduces the expert’s behavior with learned generalization across cube positions. For manipulation, that’s exactly what you want. The creativity belongs to the high-level planner (the LLM orchestrator), not the motor policy.

Scaling Training to the Cloud

Script 09 wraps the training pipeline in AWS infrastructure for reproducibility and GPU access. The stack:

  • AWS CDK defines the infrastructure as code
  • ECR hosts the training container (PyTorch + LeRobot + MuJoCo)
  • SageMaker runs training on a g5.xlarge (single A10G GPU — more than enough for ACT)
  • Step Functions orchestrates the pipeline: build container → push to ECR → launch training job → evaluate → store artifacts

The whole pipeline is defined in ~200 lines of CDK TypeScript. Training a 50-demo ACT policy takes about 30 minutes on the g5.xlarge. The resulting policy checkpoint is 47MB — small enough to deploy anywhere, including on-robot.

I deliberately kept the cloud piece minimal. The point of this project isn’t infrastructure complexity — it’s the learning journey. But having a reproducible GPU pipeline means anyone can retrain with their own demonstrations or a different arm without wrestling with local CUDA.

What I Learned

After nine scripts, three learning approaches, and more failed experiments than I’m willing to admit, here are the distilled lessons:

LessonDetail
Observation design > computeOne pre-computed vector (grip→cube) was worth more than doubling training steps
More training can hurt800k steps overfit where 400k generalized — always validate on held-out configs
RL struggles with contactDiscontinuous contact dynamics break smooth reward optimization
Imitation wins for manipulation50 expert demos → more reliable grasping than 400k RL steps
The expert doesn’t need to be smartA rigid IK controller is fine as a demo source — robustness comes from the learned policy
LLM as orchestrator, not controllerThe right architecture is LLM decides what to do, specialized policies execute how
Simulation is good enoughMuJoCo on a MacBook is sufficient for the full research loop: train, eval, iterate

The Architecture That Emerges

The final mental model that fell out of this project:

┌────────────────────────────────────────────┐
│  LLM Orchestrator (Strands Agent)          │
│  "Pick up the red cube, then place it      │
│   on the blue platform"                    │
│                                            │
│  Decomposes into: pick(red_cube) →         │
│                   place(blue_platform)      │
└─────────────────┬──────────────────────────┘
                  │ selects policy + params

┌────────────────────────────────────────────┐
│  Policy Library                            │
│  • pick_policy.pt (ACT, trained on demos)  │
│  • place_policy.pt (ACT, trained on demos) │
│  • reach_policy (RL, works fine here)      │
│  • push_policy (RL)                        │
└─────────────────┬──────────────────────────┘
                  │ 200Hz control loop

┌────────────────────────────────────────────┐
│  Robot (sim or real)                       │
└────────────────────────────────────────────┘

The LLM is too slow for real-time control (200ms per token vs. 5ms control cycles). But it’s perfect for task decomposition, error recovery, and multi-step sequencing. The specialized policies handle the millisecond-level motor control. Each layer does what it’s good at.

This is the Robot-as-AgentTool pattern evolved: the tool isn’t raw joint control anymore — it’s “execute this trained policy with these parameters.” The LLM stays in the loop for high-level decisions without bottlenecking low-level execution.

Where This Goes Next

Three directions I’m exploring:

Sim-to-real transfer. The SO-100 is a real, purchasable arm (~$200). MuJoCo’s contact model is good enough that policies trained in sim transfer with minor domain randomization. The next step is deploying the ACT policy on physical hardware.

Multi-task sequencing. One pick policy is a building block. Stacking, sorting, assembly — these require the LLM orchestrator to chain policies, handle failures, and replan. The Strands Agent framework already supports this through tool-use loops.

Shared policy libraries. If 50 demos produce a reliable pick policy, a community of sim users could collectively build a library of manipulation primitives. LeRobot already provides a model hub. The missing piece is a standardized observation/action interface — which is exactly what strands-robots defines.


The entire project — all nine scripts, the CDK pipeline, the trained checkpoints — is open source at github.com/agiusalexandre/strands_robot.

You can reproduce everything in this post on a MacBook with no hardware. A uv sync installs all dependencies (Python 3.12, MuJoCo, PyTorch, Strands SDK). The RL scripts train on CPU in under an hour. The only thing that needs AWS is Bedrock (for the LLM agent) and optionally SageMaker (for GPU training of the ACT policy — though a local GPU works fine too).

The age of robots-that-need-a-lab is ending. The age of robots-that-live-in-a-terminal is here. If you can write Python and call an API, you can build a robot that grasps.

ABOUT THE AUTHOR

Alexandre Agius

Alexandre Agius

AWS Solutions Architect

Passionate about AI & Security. Building scalable cloud solutions and helping organizations leverage AWS services to innovate faster. Specialized in Generative AI, serverless architectures, and security best practices.

ONE LETTER A MONTH · NO TRACKER · UNSUBSCRIBE ANYTIME

CONTINUE READING

Related dispatches

Comments

Sign in to leave a comment