SegDAC: Visual Generalization in RL via Dynamic Object Tokens

Abstract

Visual reinforcement learning policies trained on pixel observations often struggle to generalize when visual conditions change at test time. Object-centric representations are a promising alternative, but most approaches use fixed-size slot representations, require image reconstruction, or need auxiliary losses to learn object decompositions. As a result, it remains unclear how to learn RL policies directly from object-level inputs without these constraints. We propose SegDAC, a Segmentation-Driven Actor-Critic that operates on a variable-length set of object token embeddings. At each timestep, text-grounded segmentation produces object masks from which spatially aware token embeddings are extracted. A transformer-based actor-critic processes these dynamic tokens, using segment positional encoding to preserve spatial information across objects. We ablate these design choices and show that both segment positional encoding and variable-length processing are individually necessary for strong performance. We evaluate SegDAC on 8 ManiSkill3 manipulation tasks under 12 visual perturbation types across 3 difficulty levels. SegDAC improves over prior visual generalization methods by 15% on easy, 66% on medium, and 88% on the hardest settings. SegDAC matches the sample efficiency of the state-of-the-art visual RL methods while achieving improved generalization under visual changes.

The Problem: Visual Generalization

Visual RL policies trained directly on pixels reach strong performance in simulation, but they are sensitive to visual changes that do not affect the task itself. A shift in background color, lighting, or texture at test time can break a policy even though the objects and the goal did not move.

Data augmentation on pixels helps close this gap, but it only covers the specific transformations chosen ahead of time. It does not remove the core issue, the policy is still reasoning over raw pixels, where object identity and background are entangled in the same representation.

DrQ-v2 evaluated on a held-out seed from the same training distribution (left) succeeds. The same policy on the same task, evaluated after a small background change (right), fails. Nothing about the task changed, only the visual appearance of the background.

SegDAC addresses this by changing what the policy reasons over. Instead of a full image, it reasons over a small set of object tokens obtained from segmentation, described next.

Method: Reasoning Over Objects, Not Pixels

SegDAC, short for Segmentation-Driven Actor-Critic, turns each frame into a small set of object tokens before it ever reaches the policy. The idea is simple: if the representation already separates objects from background, the policy does not have to learn that separation itself from reward alone.

Step 1 · Object-Centric Segmentation

Pixels to masks

At every timestep, SegDAC receives one RGB frame and a short list of text concepts, for example "robot", "cube", "background". An open-vocabulary detector (YOLO-World) proposes a bounding box for each concept found in the frame. A segmentation model (EfficientViT-SAM) turns each box into a pixel mask, using only the box coordinates, so a wrong text label cannot corrupt the mask itself. A light morphological cleanup (opening then closing) removes small holes and noise. Both models stay frozen, nothing here is trained. The number of segments N is not fixed, it changes every frame depending on what is visible.

Step 2 · Dynamic Object Tokens

Masks to contextual token embeddings

A vision transformer already encodes each patch of the image with some context from the rest of the scene, through self-attention. SegDAC reuses this: for each mask, it averages the patch embeddings that overlap the mask by at least 4 pixels into a single token. No new parameters are learned for this step. Because the encoder is frozen, these tokens can be computed once and cached in the replay buffer, instead of being recomputed at every gradient update.

Step 3 · Transformer Actor-Critic Over Dynamic Tokens

From object tokens to actions and Q-values

The object tokens, plus one proprioception token, go through a transformer decoder. A learned modality embedding tells the model whether a token is an object, proprioception, or a query. Each object token also gets a segment positional encoding built from its bounding box coordinates, so the model keeps some notion of where each object sits in the frame. The actor uses one learned query that cross-attends over the tokens to produce an action. The critic uses an action-conditioned query: it combines a learned token with the actual action before cross-attending, so it evaluates the Q-value for that specific action.

Segment positional encoding helps sample efficiency

Transformer decoders & learned query fusion help training stability

Since N changes every step, transitions cannot be batched the usual way. SegDAC packs all tokens in a batch into one sequence, with an attention mask so each token only attends within its own timestep. This sequence-packing trick is common in large language model training and lets SegDAC process a variable number of tokens without padding or a hard cap.

SegDAC Is Simple To Train

Plain SAC loss
No image reconstruction
No auxiliary losses
No hand-crafted data augmentation
No ground-truth masks

Robust to a Changing Number of Objects

The number of detected segments is not stable. It moves with occlusion, detector confidence, and how the scene happens to look at that instant. Within a single episode, we observe the segment count moving roughly between 5 and 21.

Segment count and identity changing across a rollout

Segments change in count and identity across a rollout, on PushCube and UnitreeG1PlaceAppleInBowl. SegDAC is not trained assuming any fixed set of objects.

The segment count within a single episode is far from constant. SegDAC has to work with as few as a handful of segments or as many as twenty.

The detector and segmentation model are not perfect either. Segments can flicker, split, merge, or disappear from one frame to the next, on top of the count just changing as objects move in and out of view. This is the kind of noisy, unstable input SegDAC has to be robust to, not just a clean variable count, and it matters for real-world use, where perception is never perfect.

All segments SegDAC consumes at each timestep, stitched into one image for preview, on UnitreeG1PlaceAppleInBowl

Every segment SegDAC consumes at each timestep on UnitreeG1PlaceAppleInBowl, stitched into one image so you can see the raw detector and segmentation noise at once. This stitching is only for preview, SegDAC always receives these as a sequence of separate tokens, never as one image.

Watch how much the segments move around even though the scene itself is fairly static: a small shift in lighting or in the detector's own confidence is enough to make a segment split in two, disappear for a frame, or merge with its neighbor. UnitreeG1PlaceAppleInBowl is also the task where the fixed-length ablation below shows the largest gap, which lines up with what this looks like: a fixed slot count has no way to absorb this much frame-to-frame noise, while SegDAC keeps training and acting normally through it.

We test whether this dynamic, variable-length design is actually needed, by removing each piece in isolation.

Ablation removing segment positional encoding

Removing segment positional encoding does not stop learning, but it consistently hurts sample efficiency and increases variance across tasks.

Ablation comparing fixed and dynamic segment count

Clamping the token count to a fixed number (padding or truncating to 5) is competitive on simple tasks but degrades on tasks needing finer object localization, most visibly on UnitreeG1PlaceAppleInBowl.

Ablation comparing object-centric tokens with a single global token

Replacing the object tokens with one globally pooled vector collapses performance on 7 of 8 tasks. The gain comes from keeping objects separate, not just from the strength of the underlying frozen encoder.

Both segment positional encoding and variable-length processing are each individually necessary for strong performance.

What the Critic Pays Attention To

With no attention supervision at all, we look at where the critic's cross-attention lands on the object tokens while it predicts a Q-value.

Critic attention on visual segments and proprioception during Q-value prediction, without any frame stacking or history.

On LiftPegUpright, attention starts mostly on the peg, roughly 75% early in the trajectory, then shifts partly toward proprioception once the peg is upright, likely useful for stabilizing it.

On PickCube, the goal marker is fully hidden behind the arm early in the trajectory. The critic keeps its attention on the cube, the correct sub-goal at that point, and the policy keeps moving toward the goal region correctly even while it stays occluded.

Without any supervision on what to attend to, the critic learns to weight task-relevant objects and shifts that weight as the sub-goal changes: object first, goal after grasping.

A Visual Generalization Benchmark on ManiSkill3

To measure visual generalization directly, we built a benchmark on top of ManiSkill3 with 8 manipulation tasks across two robot embodiments, a Franka Panda arm and a Unitree G1 humanoid: LiftPegUpright, PickCube, PokeCube, PullCube, PullCubeTool, PushCube, UnitreeG1PlaceAppleInBowl, and UnitreeG1TransportBox.

Scene Taxonomy

Every scene is broken into 4 entity types, so a perturbation can target exactly one of them: the Manipulation Object (what gets grasped or moved), the Receiver Object (an object needed for the task but not directly manipulated, for example a bowl), the table (the main interaction surface), and background elements (walls, floor, sky).

A hard table-color perturbation on PushCube, with entities labeled according to the scene taxonomy.

Perturbation Types

We apply 12 perturbation types across 4 categories: camera (pose, field of view), lighting (direction, color), color (of the manipulation object, receiver object, table, and background), and texture (of the same 4 entities). A task without a receiver object gets fewer than 12.

Difficulty Levels

Difficulty is defined by how much the scene changes visually, and whether that change creates a semantic conflict between entities.

Easy

Small visual change, no semantic conflict, close to the default scene.

Medium

A larger visual shift, still no semantic conflict between entities.

Hard

A large visual change together with a deliberate semantic conflict, for example texturing the manipulation object with a pattern close to the goal marker, or coloring the table the same as the cube.

All 8 Tasks, All Perturbations

Each grid below shows the full set of tested perturbations for one task, one row per difficulty level.

Lift Peg Upright

Pick Cube

Place Apple in Bowl

Poke Cube

Pull Cube

Pull Cube Tool

Push Cube

Transport Box

Results

+15%

Easy

+66%

Medium

+88%

Hard

Improvement in IQM return over the strongest prior visual generalization baselines, at each difficulty level.

At hard difficulty, most baselines lose more than half of their unperturbed performance: DrQ-v2 loses over 90%, SADA and MaDi lose 59%. SAM-G is the only baseline that stays competitive at hard, and SegDAC still beats it by 88% higher IQM.

Visual generalization results by perturbation category and difficulty

IQM return by perturbation category and difficulty level. Gray bars show unperturbed performance, colored bars show performance under perturbation with 95% confidence intervals.

Sample Efficiency

SegDAC matches DrQ-v2 sample efficiency across all 8 tasks and is ahead on 3 of 8, without using any data augmentation. It also does not collapse on any task, unlike some of the generalization baselines, for example MaDi does well on PullCube but fails on LiftPegUpright.

Aggregated sample efficiency, IQM normalized return over training, across 5 seeds.

Baselines are grouped into three families: pixel-based (SAC-AE, DrQ-v2, MaDi, SADA), self-supervised (SMG), and segmentation-based (SAM-G).

Results in Action

Rollouts of DrQ-v2, MaDi, and SegDAC, side by side, one video per difficulty level.

Easy

Small visual change, no semantic conflict.

Medium

Larger visual shift, still no semantic conflict.

Hard

Large visual change with a deliberate semantic conflict.

Takeaways

Strong visual generalization is possible without hand-crafted data augmentation.

Reasoning over objects instead of pixels is a favorable inductive bias for both robust visual generalization and efficient learning.

Training on a dynamic set of object tokens, changing in count and identity at every step, stays stable under plain SAC and promotes robustness to state variation.

Simple to train

Only needs the RL reward
Plain SAC loss
Drop-in for any model-free algorithm

Limitations

SegDAC is evaluated on short-horizon manipulation tasks in simulation. The policy is memoryless, it does not reason over object state across timesteps, which may limit performance on tasks requiring long-horizon planning. The 2D positional encoding may not generalize to extreme viewpoint shifts. Text input is used only for open-vocabulary detection, the policy does not condition on language instructions.

BibTeX

Please use the following BibTeX entry to cite this work:

@article{brown2026segdac,
    title={SegDAC: Visual Generalization in Reinforcement Learning via Dynamic Object Tokens},
    author={Alexandre Brown and Glen Berseth},
    journal={Reinforcement Learning Journal},
    volume={7},
    pages={},
    year={2026}
}

SegDAC: Visual Generalization in Reinforcement Learning via Dynamic Object Tokens

SegDAC reasons over a dynamic set of object tokens instead of pixels, giving both efficient learning and strong visual generalization.

Abstract

The Problem: Visual Generalization

Method: Reasoning Over Objects, Not Pixels

Step 1 · Object-Centric Segmentation

Step 2 · Dynamic Object Tokens

Step 3 · Transformer Actor-Critic Over Dynamic Tokens

Robust to a Changing Number of Objects

What the Critic Pays Attention To

A Visual Generalization Benchmark on ManiSkill3

Scene Taxonomy

Perturbation Types

Difficulty Levels

All 8 Tasks, All Perturbations

Results

Sample Efficiency

Results in Action

Takeaways

Limitations

BibTeX