Reinforcement Learning for Block Pickup with Visual State Estimation
Xiaoxiang (Tony) Dong & Shuaiwei Huang
Embedded Files
Link to Full Project Report： https://drive.google.com/file/d/1ncxcAw7XX9vs5hMZO8ytne2Zudzr4AHf/view?usp=sharing
Project Github Repo: https://github.com/Tony050519/16831_project.git
Robotic pick-and-place remains a foundational benchmark for sensorimotor learning, but training a policy end-to-end is bottlenecked by long-horizon credit assignment, the precision required for stable grasp synthesis, and the gap between the privileged state available in simulation and the partial, noisy perception available on real hardware. In this work, we explore reinforcement learning for robotic pick-and-place under two input regimes: privileged simulation state and a single fixed RGB camera. Using an Agile-X Piper 6-DoF arm in IsaacLab, we design a shaped reward and compare three state-based policies: (1) PPO with a two-stage reach-then-lift curriculum, (2) SAC, and (3) the model-based TD-MPC2. We then train a vision-based CNN-MLP student to perform the same task without access to the block pose, distilling a privileged PPO teacher via three progressive stages of DAgger. PPO and TD-MPC2 each reach a 100% lift success rate, SAC reaches 70%, and the vision-based student reaches 80%. The results suggest that, in this contact-rich setting, a learned world model recovers PPO's asymptotic performance at an order of magnitude fewer environment interactions, and that teacher-student distillation provides a practical path from privileged-state policies to perception-driven policies suitable for real hardware.
RL Methodologies
Proximal Policy Optimization (PPO) PPO training employs a two-stage procedure. In the first stage, the robot is incentivized to position the gripper at a desired pose above the object block, preparing for subsequent grasping and lifting. In the second stage, additional incentives are introduced to encourage successful grasping once the gripper reaches a suitable pose.
Soft Actor Critique (SAC) SAC training directly adopts the same reward formulation as the PPO finetuning stage and is trained in a single stage with finetuning reward directly.
Model-Based Reinforcement Learning (MBRL) MBRL uses TD-MPC2 which jointly trains a world model with six models embedded: encoder, latent dynamics, policy prior, reward model, Q-function, and termination model. Instead of learning trajectory directly on the observation space, it encodes the observation to a latent space via encoder, then samples an action from policy prior, and predicts the next latent state via latent dynamics. In the trajectory planning phase, TD-MPC2 trains a reward model to predict reward given a pair of latent state and action, then estimate q-value via Q-function. Termination model is a network to learn termination condition for piper robot arm environment. TD-MPC2 trains in a single stage with finetuning reward directly.
Agents Reward vs. Step Curve
SAC demonstrates good sample efficiency, achieving strong performance with significantly fewer environment interactions. PPO, particularly with the proposed two-stage training strategy, achieves substantially higher final performance. MBRL pushes sample efficiency further, significantly exceeding the asymptotic performance of SAC and PPO which requires an order of magnitude more data to reach reward of 3,000. Notably, MBRL not only peaks the earliest of all methods but also achieves the highest sustained return, suggesting that for this manipulation task, a well-trained world model provides a more informative learning signal than the reward gradients available to model-free methods.
Agents Performance Visualization
PPO_Agent_Performance.mp4PPO Agent
SAC_Agent_Performance.mp4SAC Agent
MBRL_Performance.mp4MBRL Agent
robot_random_walk.mp4Random Agent
Teacher-Student Policy Distillation
Rather than providing the policy with the ground-truth block position directly, we attach a static RGB camera to the environment that observes the robot arm and the block simultaneously. The student policy receives two inputs: (i) the camera image, and (ii) the robot's proprioceptive state, instead of the block's position, orientation, and its relative displacement to the gripper. We then adopt a teacher–student distillation framework to transfer manipulation behavior from a privileged state-based teacher (PPO policy) to a vision-based student.
Distillation Loss Curve
Distilled Policy Performance Visualization
distilled_policy_performance.mp4Student Policy Visual Input during Rollout
Page updated
Report abuse