Hierarchical Entity-centric Reinforcement Learning with Factored Subgoal Diffusion

AI summary

Overview

• Research area: Offline goal-conditioned reinforcement learning (GCRL), hierarchical RL, diffusion models for decision-making, and entity-centric (object-centric) representation learning, evaluated on multi-object robotic manipulation from pixels.
• Technical level: Advanced. The paper assumes familiarity with temporal-difference value learning, goal relabeling, subgoal hierarchies, denoising diffusion models, and unsupervised object-centric representations.
• Scope (one sentence): The paper introduces HECRL, a modular two-level framework that pairs a value-based entity-centric GCRL agent with an entity-factored conditional diffusion subgoal generator, and shows it improves long-horizon, sparse-reward, multi-entity manipulation from image observations.

What This Paper Is About

Offline GCRL agents struggle when environments contain many entities, because the state space grows combinatorially and sparse rewards must propagate back through long horizons of temporal-difference updates. The authors combine two existing ideas that are usually treated separately: subgoal hierarchies (treating intermediate trajectory states as goals) and factored entity-centric structure (representing the state as a collection of objects). Their goal is an agent that reaches long-horizon goals in multi-entity, image-based domains from a fixed dataset of suboptimal trajectories, without supervised object labels.

Key Contributions

1. A modular hierarchical framework (HECRL) where the low level is a value-based entity-centric GCRL agent and the high level is a conditional diffusion model that generates subgoals. The two levels are trained independently and composed only at test time, so the method can be attached to potentially any value-based GCRL algorithm without modifying it.
2. A test-time subgoal generation and selection procedure that samples N subgoal candidates from the diffuser, filters them for reachability using the value function against a threshold R-hat (the "policy competence radius"), and selects the candidate closest to the goal; if the goal is closer than the chosen subgoal, the goal is used directly.
3. Entity-factored subgoal diffusion, adapting the entity-centric diffuser of Qi et al. (2025) so that each state, goal, and subgoal entity is a separate input token to a Transformer denoiser. This inductive bias encourages subgoals that modify only a small subset of entities, and enables small diffusion models with as few as 10 denoising steps.
4. New benchmark variations of OGBench Cube/Scene and Push-T (PPP-Cube, Stack-Cube, Scene, Push-Tetris) with increased image resolution, multi-view perception, removed visual simplifications, and no goal-based termination at test time, plus datasets, checkpoints, and code released for reproduction.

Main Findings

• Hierarchy boosts the same underlying agent. EC-SGIQL (the proposed method) uses the same low-level policy and value function as EC-IQL yet consistently improves it. On PPP-Cube (Image), EC-SGIQL reaches 64.3 ± 4.9 success versus 25.0 ± 5.7 for EC-IQL — described in the abstract and introduction as more than a 150% increase in success rate on the hardest task in the suite.
• Entity-centric structure matters. EC-IQL is reported as the second most performant method across the suite, which the authors cite as evidence for the value of factored structure in these domains.
• Full comparison table (Table 1). PPP-Cube (State): EC-SGIQL 82.5 ± 3.1, EC-IQL 51.5 ± 4.4, EC-Diffuser 44.8 ± 6.7, HIQL 48.3 ± 7.3, IQL 34.3 ± 4.9. Stack-Cube (State): 43.5 ± 1.9, 29.0 ± 2.9, 43.8 ± 9.2, 0.0 ± 0.0, 19.3 ± 3.0. PPP-Cube (Image): 64.3 ± 4.9, 25.0 ± 5.7, 0.3 ± 0.5, 0.0 ± 0.0, 0.0 ± 0.0. Scene (Image): 61.5 ± 5.9, 53.0 ± 5.5, 3.3 ± 2.5, 8.3 ± 1.3, 17.5 ± 2.7. Push-Tetris (Image, reported as pixel coverage): 61.4 ± 3.3, 31.6 ± 1.3, 7.9 ± 0.5, 5.2 ± 0.8, 3.4 ± 0.8.
• Both value-based selection and filtering help. Removing the filter ("Max Value") still performs on par on some tasks, indicating the policy competence radius may be larger than the radius of generated subgoals there. Randomly sampling a subgoal ("Random Sample") and replacing the diffuser with a deterministic Transformer trained with advantage-weighted regression ("AWR") both degrade performance; AWR is the lowest overall, attributed to weighted averages of observations that are not valid states. Random sampling and AWR also need more timesteps when success rates are comparable.
• Subgoals are measurably sparser. In the state-based Cube environments, the number of modified entities (excluding the agent) in generated subgoals is 1.36 ± 0.01 for the entity-centric diffuser versus 2.96 ± 0.01 for EC-AWR and 2.98 ± 0.00 for AWR on PPP-Cube; and 1.04 ± 0.01 versus 2.82 ± 0.02 and 2.98 ± 0.01 on Stack-Cube, out of 3 cubes plus the arm, averaged over 400 randomly sampled initial states and goals.
• Qualitative subgoal behaviour. Diffuser subgoals often compose the input state and goal images with sparse changes. Occasional errors duplicate an entity, usually as an exact copy from one of the two inputs, which the authors read as support for the Transformer selectively copying input tokens; they attribute this to the DLP representation and the diffusion objective. AWR subgoals often contain multiple duplicates of entities representing different futures inside a single subgoal, giving the low-level policy ambiguous goals.
• Zero-shot compositional generalization (Table 4). Trained on 3 cubes, EC-SGIQL reaches 69.8 ± 7.0 on 2-Cubes-Stack, 43.5 ± 1.9 on 3-Cubes-Stack, and 28.8 ± 3.9 on 4-Cubes-Stack-2; on PPP with state observations it degrades from 97.0 ± 1.8 at 1 cube to 25.7 ± 1.5 at 6 cubes, and with images from 94.5 ± 2.9 at 1 cube to 9.5 ± 1.3 at 6 cubes. In one case (2-Cubes-Stack) EC-Diffuser scores higher (87.8 ± 5.4 versus EC-SGIQL's 69.8 ± 7.0).
• Comparison of the three columns is not uniformly in favour of the proposed method. The paper states it outperforms all baselines with two exceptions where it performs on par.
• Entity-centric methods reach peak performance with fewer gradient updates than unstructured ones when the latter achieve non-trivial performance.

Methodology in Plain English

The authors start from a standard goal-conditioned setup where goals are states and the reward is 0 at the goal and -1 otherwise, learning offline from a fixed dataset. They build a two-level system.

The low level is any value-based goal-conditioned agent that has been given entity-centric structure — a policy and a value function that take a set of objects rather than a single flat vector. In their experiments this is an ECRL-style agent trained with IQL, and image inputs are converted into object slots using a pretrained unsupervised object-centric model (DLPv2), which they note requires no supervision.

The high level is a diffusion model trained separately on the same offline dataset. Training pairs are built mechanically: sample a state from a trajectory, sample a later state in the same trajectory as the goal, and use the state K steps after the first one (or the goal itself, whichever comes first) as the target subgoal. The model learns the distribution of such near-future states given the current state and the goal. Because the dataset contains no goal-directed behaviour, this distribution is multi-modal, which is the stated motivation for using diffusion rather than regression.

At test time, the system does not blindly trust the diffuser. It samples N subgoal candidates, keeps only those the value function says are reachable from the current state (value above a threshold), and among those picks the one the value function ranks as closest to the final goal. The low-level policy then executes towards that subgoal for a fixed number of timesteps before the procedure repeats — the authors describe this as constrained sample-based planning with receding horizon control, where the "model" is the diffuser-value pair and optimization happens in state space rather than action space. If the goal itself turns out to be closer than the chosen subgoal, the subgoal is replaced by the goal.

The entity-factored part comes from feeding each entity of the state, the goal, and the noisy subgoal into a Transformer denoiser as a separate labeled token. The authors argue this lets attention selectively copy input entities to the output, which is what produces subgoals that change only a few entities.

Why This Matters

• Impact on research: The paper gives a concrete recipe for combining offline GCRL, diffusion-based subgoal generation, and unsupervised object-centric representations, and shows that the benefits are largely modular — the high level is trained separately and bolted on without touching the RL agent. It also argues, with measurements, that the standard weighted-regression subgoal objective produces averaged, invalid subgoals while diffusion produces valid, mode-specific ones. It additionally contributes harder benchmark variants for multi-entity domains, since the authors note non-trivial performance on these environments with more than 2 objects had not been attained prior to this work within the data regime they consider.
• Real-world applications (as named or implied by the paper):
  • Robotic object manipulation, including pick, place, push, and stacking tasks with a UR5e arm (the paper's own experimental domain).
  • Multi-robot path planning.
  • Autonomous driving.
  • Video games.
• Industry relevance: The method targets learning long-horizon behaviour from fixed, suboptimal datasets with sparse rewards and realistic images — the setting typical of industrial robot data collection, where online trial and error is costly. The combination of small diffusion models (as few as 10 denoising steps) and modular test-time composition is relevant to practitioners who want to improve an existing GCRL stack rather than replace it, since the paper claims compatibility with potentially any value-based GCRL algorithm and other test-time constraints on states.

Future Directions

• Remove the fixed horizon hyperparameter K. The Subgoal Diffuser is trained with subgoals up to a fixed K steps ahead; the authors state the method is more robust to this than one might expect because of test-time filtering, but suggest automatically inferring K from data as future work.
• Loosen the reliance on the value function. The approach assumes the value function gives the low-level policy a non-negligible competence radius and can also guide the subgoal generator. This held in the tested domains but may limit applicability elsewhere.
• Scale to more entities. The paper reports generalization degrading with increasing numbers of objects, and suggests curriculum learning or offline-to-online finetuning as routes to scaling, explicitly leaving this for future work.
• Improve the factored state estimator. Image-based performance depends on a good unsupervised object-centric representation; the authors tie progress in in-the-wild scenarios to advances in that area and note that the observed entity-duplication artifacts might be remedied with larger diffusion Transformers.

Target Audience

Researchers and graduate students working on offline reinforcement learning, goal-conditioned RL, hierarchical RL, and diffusion-based decision-making, especially those interested in long-horizon, sparse-reward problems. It is also relevant to roboticists who need manipulation policies learned from fixed datasets of pixels, and to anyone working on object-centric or entity-centric representation learning who wants to see those representations used as a structural prior inside a hierarchy. Practitioners seeking a drop-in test-time improvement on top of an existing value-based GCRL agent are the most direct applied audience; the paper's code and checkpoints are released at https://github.com/DanHrmti/HECRL and rollout videos at https://sites.google.com/view/hecrl.

Authors’ abstract

We propose a hierarchical entity-centric framework for offline Goal-Conditioned Reinforcement Learning (GCRL) that combines subgoal decomposition with factored structure to solve long-horizon tasks in domains with multiple entities. Achieving long-horizon goals in complex environments remains a core challenge in Reinforcement Learning (RL). Domains with multiple entities are particularly difficult due to their combinatorial complexity. GCRL facilitates generalization across goals and the use of subgoal structure, but struggles with high-dimensional observations and combinatorial state-spaces, especially under sparse reward. We employ a two-level hierarchy composed of a value-based GCRL agent and a factored subgoal-generating conditional diffusion model. The RL agent and subgoal generator are trained independently and composed post hoc through selective subgoal generation based on the value function, making the approach modular and compatible with existing GCRL algorithms. We introduce new variations to benchmark tasks that highlight the challenges of multi-entity domains, and show that our method consistently boosts performance of the underlying RL agent on image-based long-horizon tasks with sparse rewards, achieving over 150% higher success rates on the hardest task in our suite and generalizing to increasing horizons and numbers of entities. Rollout videos are provided at: https://sites.google.com/view/hecrl