OXE-AugE: A Large-Scale Robot Augmentation of OXE for Scaling Cross-Embodiment Policy Learning

In plain words

A robot learns a task by copying recorded demonstrations. Most of those recordings come from a few robot arms, so a learned policy often works only on the arm it saw. Collecting new demonstrations for every new arm and gripper is too expensive.

This paper reuses the recordings that already exist. For each video frame, it hides the original arm, fills in the background where the arm was, and draws a different arm in the same pose, using a physics simulator for the new arm. The result looks like the same task performed by another robot. The authors call this robot augmentation.

They applied the method to 16 widely used robot datasets from the Open X-Embodiment collection and added nine robot arm and gripper combinations. The output is a dataset of over 4.4 million trajectories. The authors then fine-tuned two large robot foundation models on it and tested them on arm and gripper pairs they had not seen. They report gains on those new pairs, in real robot tests.

TL;DR

• AugE-Toolkit extends the cross-painting pipeline. It masks the source arm with a fused SAM2 and simulation mask, inpaints the background with E2FGVI, and replays the trajectory on a target arm in simulation §4.2 fig. 2.
• OXE-AugE applies it to 16 OXE datasets, with 8 or 9 target robots each. It holds over 4.4M trajectories §6 app. A.2 tab. 1.
• Author-reported: in a simulation study, augmentation helps the source arm under visual perturbations and helps unseen arms §5.1.
• Author-reported: fine-tuning OpenVLA-OFT and pi0 on OXE-AugE raises their average scores by 24% and 45% over the no-augmentation baseline §6.2.
• Official code (MIT) is in the authors' repository. Its core pipeline lives in a submodule whose repository has no license file, so we do not use it README.md .gitmodules.

Problem

Generalist robot policies train on pooled data from many labs. The authors note that the Open X-Embodiment (OXE) dataset is highly imbalanced: four robot types supply over 85% of its real trajectories, so policies may overfit to those robot and scene combinations §1. Re-collecting demonstrations for each new robot is costly §1. Prior cross-painting methods either need a diffusion model trained per target robot, which lacks kinematic guarantees, or need known camera calibration, which large offline datasets often lack §4.1.

The paper asks three questions: whether augmentation improves the original robot under visual perturbations (robustness), whether more augmented robots improve performance on augmented robots (transfer), and whether policies trained on diverse augmentations generalize to robots never used in training (generalization) §1 §3.

Method

Setting. Imitation learning uses n source demonstrations, each a sequence of RGB observations, gripper 6D poses and actions §3. Augmentation produces a matching sequence for another embodiment (arm and gripper) that performs the same task. The method assumes the grippers are similar in shape and function (2- or 3-jaw), that Cartesian control is used, and that all frame transforms are known §3. The paper states no numbered equations; the formal setup is given in prose §3.

Cross-painting has three stages per frame: segment the source robot, inpaint the background, and replay and composite the augmented robot §4.1. Learning-based methods do this in pixels with diffusion models, and simulation methods render the robot from known camera parameters §4.1.

AugE-Toolkit keeps the three stages and adds three components §4.2.

1. Mask fusion §4.2 (1). SAM2 is fine-tuned on a small labelled set (20 trajectories from each of 16 OXE datasets) §4.2 (1). Its mask follows the image well but over- or under-segments near the gripper. A simulation mask is geometrically clean but can be shifted by unknown camera pose. The two are fused in three steps: shift the simulation mask on a small grid to maximise IoU with the learned mask; remove learned-mask pixels farther than a threshold tau from the aligned simulation boundary; take the union and apply morphological closing §4.2 (1).
2. Automatic base tuning §4.2 (2). Each target robot's base position is adjusted in simulation so that all end-effector poses of the source trajectory are reachable. Starting from an initial base, the method samples offsets of plus or minus delta along x, y and z, measures the maximum tracking error, keeps the best candidate and halves the step size. It stops when the maximum error is below 1 cm or an iteration limit is reached. Trajectories with no feasible base are discarded §4.2 (2).
3. Multi-robot deployment §4.2 (3). Replay uses URDF-style simulation rendering rather than a generative model, which avoids flicker and geometry artifacts §4.2 (3). Each target robot replays the same trajectory independently, so the stages run in parallel. A new robot is added by registering its MJCF model §4.2 (3).

The pipeline for one trajectory and one target robot, following §4.1 and §4.2 and the pipeline figure fig. 2:

# Cross-paint one trajectory onto one target robot  [§4.1] [fig. 2]
# Input:  frames (T, H, W, 3); poses (T, 7); target robot model
# Output: augmented frames (T, H, W, 3), or None if discarded

def augment_trajectory(frames, poses, target_robot):
    # (i) Segment the source arm in each frame  [§4.2 (1)]
    masks = []
    for frame, pose in zip(frames, poses):
        learned = sam2_mask(frame)
        simulated = render_source_mask(pose)
        masks.append(fuse_masks(learned, simulated))
    # (ii) Fill in the background behind the arm, across frames  [§4.1]
    background = video_inpaint(frames, masks)
    # Choose a base pose so the target can reach every pose  [§4.2 (2)]
    base = tune_base_position(poses, target_robot)
    if base is None:
        # No feasible base: the trajectory is discarded  [§4.2 (2)]
        return None
    # (iii) Replay the poses on the target and composite  [§4.1] [fig. 2]
    robot_frames = replay_in_simulation(target_robot, base, poses)
    return composite(background, robot_frames)

Fusing the masks for one frame, following §4.2 (1). The threshold tau and the grid size are not stated in the paper.

# Fuse a learned mask and a simulation mask  [§4.2 (1)]
# Input:  learned (H, W) bool; simulated (H, W) bool; tau in pixels
# Output: robot_mask (H, W) bool

def fuse_masks(learned, simulated, tau):
    # Step 1: shift the simulation mask to best match the learned mask
    shifted = best_grid_translation(simulated, learned)
    # Step 2: drop learned pixels farther than tau from the shifted edge
    pruned = prune_far_pixels(learned, shifted, tau)
    # Step 3: union of both masks, then morphological closing
    return morphological_close(shifted | pruned)

Base-position tuning, following §4.2 (2). The paper gives the 1 cm stop threshold but not delta or the iteration limit.

# Automatic base-position tuning  [§4.2 (2)]
# Input:  poses (T, 7) source end-effector poses; base (3,) in metres
# Output: feasible base (3,), or None

def tune_base_position(poses, base, delta, max_iters):
    # Start with a step size of delta  [§4.2 (2)]
    step = delta
    for _ in range(max_iters):
        # Candidates: the base moved by plus or minus step on each axis
        candidates = axis_offsets(base, step)
        # Replay the poses from each candidate and keep the best  [§4.2 (2)]
        errors = [max_tracking_error(poses, c) for c in candidates]
        best = argmin(errors)
        base, error = candidates[best], errors[best]
        # Stop once the maximum tracking error is below 1 cm  [§4.2 (2)]
        if error < 0.01:
            return base
        # Halve the step size for the next iteration  [§4.2 (2)]
        step = step / 2
    # No feasible base found within the iteration limit  [§4.2 (2)]
    return None

OXE-AugE dataset §6 app. A.2. The authors select 16 OXE datasets used in common robot foundation-model training mixtures. Their original robots are Franka, UR5, xArm, WidowX, Google Robot and Jaco. Each is augmented to up to nine robots: those six plus Sawyer, Kinova Gen3 and KUKA iiwa §6. Most datasets get all nine robots. RT-1 (Fractal) and Language Table get no WidowX augmentation, because most of their trajectories exceed the WidowX workspace app. A.2 tab. 1. Gripper assignments are listed in the appendix [app. A.2].

Simulation study §5. Five Robosuite tasks are used: Lift, Stack, Can, TwoPiece Assembly and Square §5. Each has 200 demonstrations from Robomimic or MimicGen, all on a Franka arm. Four additional arms (UR5e, Kinova Gen3, Sawyer, Jaco) are the augmentation targets §5. Policies are Diffusion Policy models §5 app. A.4. Four training regimes are compared: source only, one augmented robot alone, one augmented robot plus source, and all augmented robots plus source §5. Evaluation covers the source robot under lighting and occlusion perturbations, the augmented robots, and held-out robots §5.

Physical experiments §6.1 app. A.5. OpenVLA (with OpenVLA-OFT's LoRA recipe) and pi0 are fine-tuned on Bridge data, once with no augmentation and once with two or all augmented robots §6.1. The baselines are fine-tuned on Bridge alone, with the total number of steps matched §6.1. Four tasks are evaluated on a Franka FR3 with its default gripper ("Franka") and with a custom Robotiq gripper with coloured pads ("Robotiq++"). The second embodiment is unseen during training §6.1 fig. 5. Inputs are 256x256 third-person images with a language instruction §6.1.

Diffusion comparison §6.2 fig. 6. The authors generated 700K paired images with a MuJoCo-based pipeline and trained diffusion models to translate robot appearance, following RoVi-Aug. Fine-tuning on those images lowers final policy performance by 27 to 30% §6.2.

Hyperparameters

Values in the table come from the paper text, the appendix or the README. The paper states no equation numbers, so table entries cite sections, appendix items and tables.

Results

Author-reported. Setup: simulation with Diffusion Policy on five Robosuite tasks, 100 trials per task §5 §5.1. Physical: OpenVLA-OFT and pi0 on a Franka FR3, four tasks, 10 trials per task §6.1.

The per-task and per-robot numbers for the physical experiments are only in fig. 5; the text gives averages. The simulation curves in fig. 3 and fig. 4 carry standard errors, but the text does not repeat them.

Paper vs code

The paper's official code is the GuanhuaJi/oxe-auge repository, which the project page links. Its owner is the first author, and its MIT license names Guanhua Ji LICENSE. The repository states it is the official implementation of the paper, although the README calls the work "OXE++" README.md.

We searched the arXiv metadata (no code link), the paper text, the project page (which links the repository), and GitHub for "AugE-Toolkit". The GitHub search found BerkeleyAutomation/AugE-Toolkit, which the repository uses as a submodule. That repository has no license file, so we do not use its code .gitmodules.

Discrepancies, each with both citations:

• New robot calibration. The paper says adding a robot "only requires registering its model" and that "no retraining or calibration is needed" §4.2 (3). The README asks the user to measure the gripper with a finger-tracking script, record the actuator range and joint indices in core/gripper_utils.py, and update the gripper mapping [code: README.md#L128-L133].
• Camera. The paper says the method works on "uncalibrated or coarsely aligned" data §4.2 (1). The README requires interactive camera tuning for each dataset episode. It states this step needs macOS because the MuJoCo viewer uses native windowing README.md.
• Base position. The paper describes automatic base tuning §4.2 (2). The README's replay command takes a manual --translation vector, and suggests starting from zero. An optional translation_tuning script does iterative feedback README.md. That script is in the AugE-Toolkit submodule, which we do not have, so we cannot compare its algorithm with the paper's.
• Mask fusion. The paper describes three fusion steps and a closing step §4.2 (1). The README describes merge_mask.py, which shifts, prunes and unions, with --max_pix and --sam_radius parameters. It does not mention closing README.md. The script is in the sam2 submodule, so we cannot check it.
• SAM2 weights. The paper says SAM2 is fine-tuned on 20 trajectories per dataset §4.2 (1). The README gives a link to fine-tuned checkpoints and no training recipe README.md.
• Training code. The paper's policy training (Diffusion Policy in simulation, OpenVLA-OFT and pi0 fine-tuning) is not in the main repository. The README points to separate repositories for simulation experiments README.md, which we did not clone or license-check. The learning rates and step counts here come only from the paper app. A.4.
• Name. The paper's title is "OXE-AugE" §1. The README title is "OXE++" README.md.

Pitfalls

• "Success" in the physical tests is a partial score. Each trial gives 0.3, 0.5, 0.8 or 1.0, and the reported numbers average these app. A.5. The text says "success", so do not treat the numbers as binary success rates.
• The 24% and 45% gains have no stated unit or denominator. The text says performance improves "by 24%" and "by 45%" §6.2, without saying whether these are absolute points or relative change. The abstract says "24-45% on previously unseen robot-gripper combinations" [abstract], while the text gives averages without saying which embodiment they cover.
• Seen and unseen are mixed in the headline. The Franka embodiment matches one augmented robot in OXE-AugE §6.1, so the headline mixes transfer with generalization.
• Small sample. Ten trials per task is small, and the physical experiments have no stated error bars in the text §6.1.
• Baseline steps. The no-augmentation baseline is "lightly fine-tuned" with a matched total number of steps §6.1, but the exact count is not given.
• Gripper mismatch. Section 3 assumes similar grippers §3. In simulation, Jaco has a 3-jaw gripper §5 app. A.3. The Robotiq++ test gripper is custom and unseen during training §6.1.
• Unstated tuning values. tau, the translation grid, the closing kernel, delta and the iteration limit are not stated §4.2. A reimplementation needs its own calibration of them; the authors' code has them, but we could not read them.
• Replay errors are not all small. Only "over 95%" of trajectories reach under 0.25 cm, for most datasets §4.2 (2) fn. 1. The rest are kept without a stated error bound.
• Counts may not reconcile. The appendix says "over 550,000 demonstrations per robot" and 4.44M total app. A.2. Nine robots at 550,000 each is about 4.95M, which exceeds 4.44M. The text gives no per-dataset breakdown, so the card cannot resolve this.
• Cost. Dataset-generation hardware, total compute, and the number of GPU hours are not stated. The throughput figure assumes 32-way parallelism §4.2 (3).
• Limitations the authors state. Augmentation is 2D and ignores object-robot occlusion. It assumes all arms can do the task with the same strategy. The evaluation is mostly pick-and-place with single arms and fixed bases §7.
• Figure and section numbers were derived by counting the LaTeX source, not checked against a compiled PDF. The paper has no numbered equations and one table tab. 1. Figure 1 is the cover in the abstract fig. 1.

Lineage

• Builds on: RoVi-Aug (Chen et al., 2024), which uses diffusion models for cross-painting. This paper replaces the diffusion step with simulation replay §4.1 §4.2. It is not on this wiki yet.
• Builds on: Mirage (Chen et al., 2024), which does simulation-based cross-painting at test time §4.1. It is not on this wiki yet.
• Uses: E2FGVI for inpainting and SAM2 for segmentation §4.2.
• Fine-tuned models: OpenVLA with the OpenVLA-OFT recipe (OpenVLA-OFT) §6.1, and pi0 §6.1 (no card on this wiki yet).
• Compared against: the diffusion pipeline of RoVi-Aug, under the paper's setup §6.2.
• Superseded by: not known.