patchpol

This repository reproduces a robotics research paper called Patch Policy, which trains an AI system to control a robot arm pushing a T-shaped block toward a target position. The task, known as Push-T, is a standard benchmark in embodied AI research. The paper proposes that robot control policies work better when they use dense visual information from every part of an image rather than a single compressed summary. The reproduction reports a final coverage score of 0.772 across 100 rollouts, falling about 0.06 short of the paper's target of 0.83.

The system starts with 224x224 pixel camera frames and passes them through a frozen vision model called DINOv2, which converts each frame into 256 patch tokens of 384 dimensions each. Two frames of these tokens then flow through a transformer with a specific masking pattern the README calls block-causal. Tokens within a single frame can see each other freely, while tokens from later frames cannot leak information backward to earlier ones. A diffusion model then takes the transformer's output and produces five actions at a time, values between -1 and 1 that drive the robot.

Training runs for 50,000 steps and takes about three and a half hours on an RTX 4090 GPU. The README includes detailed notes on memory requirements, explaining that a custom attention mask forces a fallback code path that uses more memory than the optimized path. A flag that enables mixed-precision math is required on a 24 GB card, cutting peak memory from roughly 24 GB down to 15.8 GB while keeping the same loss curve. Evaluation runs 100 rollouts using a physics simulation and reports coverage metrics plus optional video output.

The gap between the reproduction and the paper appears to be a terminal-holding problem. The policy drives the T-block to the goal but sometimes drifts away before the 300-step limit, losing coverage it had already achieved. The README identifies the most likely cause: the conditioning setup reads out only the last frame instead of producing a readout at every frame as the paper specifies. It also lists several other deviations from the paper, including the trunk size, diffusion schedule details, and training precision.