Hierarchical Entity-centric Reinforcement Learning with Factored Subgoal Diffusion

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

• The paper proposes a hierarchical framework that combines a low-level value-based RL agent with a high-level conditional diffusion subgoal generator.
• It leverages unsupervised object-centric representations to factor the state space, mitigating combinatorial complexity and improving sample efficiency.
• Empirical results demonstrate over 150% success rate improvements in challenging multi-entity tasks and robust zero-shot generalization.

Hierarchical Entity-centric Reinforcement Learning with Factored Subgoal Diffusion

---

Introduction and Motivation

The paper "Hierarchical Entity-centric Reinforcement Learning with Factored Subgoal Diffusion" (2602.02722) delivers a modular, scalable approach for offline Goal-Conditioned Reinforcement Learning (GCRL) in domains characterized by multiple interacting entities and long-horizon sparse-reward tasks. Standard GCRL paradigms exhibit sample inefficiency and struggle to generalize in combinatorially complex state spaces, particularly under high-dimensional visual observations. This work targets two fundamental bottlenecks: compounding value approximation errors over extended temporal horizons and the explosion of combinatorial complexity with increasing entities in the environment.

The proposed Hierarchical Entity-Centric RL (HECRL) framework employs a two-level hierarchy consisting of a low-level value-based GCRL agent and a high-level conditional diffusion model that generates entity-factored subgoals. Crucially, the subgoal generator and RL agent are trained independently and composed post hoc via a value-based subgoal selection mechanism, ensuring flexibility and compatibility with a wide range of underlying GCRL algorithms.

---

Technical Framework

Goal-Conditioned RL and Hierarchy

The baseline formulation assumes goals are states in the environment ($\mathcal{G} = \mathcal{S}$) and utilizes a sparse reward shaping. Hierarchical RL, instantiated in prior models such as HIQL, typically suffers from subgoal quality issues when the high-level policy is trained with AWR-based regression, resulting in weighted average subgoals that do not correspond to valid or easily achievable states.

HECRL instead leverages unsupervised object-centric representations (e.g., Deep Latent Particles, DLP) to encode both states and goals in an entity-centric format, facilitating the explicit incorporation of the underlying factored structure. The high-level subgoal generator is implemented as a conditional diffusion model, trained to produce subgoals a fixed number of steps ($K$) into the future, conditioned on the current state and goal. At test-time, a set of $N$ candidate subgoals are sampled from the diffusion model and filtered using the learned value function to ensure reachability and select the subgoal maximizing proximity to the final goal. This decoupling of subgoal generation and selection enables flexible control over subgoal difficulty and policy competence radius ($R$).

Figure 1: Left illustrates the value-based subgoal sampling and selection procedure for constrained reachability; right depicts factored subgoal structure, favoring modifications to minimal subsets of entities to simplify the low-level task.

Entity-centric Diffusion for Factored Subgoals

Entity-centric architectures, especially Transformer-based diffusion models, are shown to naturally induce sparsity in generated subgoals—frequently requiring modifications to only a few entities as opposed to the full state vector. This sparsity emerges from the model architecture and suitability of the underlying representations, rather than explicit regularization. The multi-modal sampling capability of diffusion enables generation of distinct subgoal modes, avoiding the degeneracy observed in deterministic AWR-trained regression policies.

The low-level agent executes policies conditioned on these subgoals for a fixed temporal interval, after which the process is repeated iteratively, yielding a receding horizon control policy analogous to Model Predictive Control (MPC) in state-space.

---

Empirical Evaluation

Environment and Task Suite

Experiments were conducted on a diverse suite of multi-entity manipulation tasks incorporating high-dimensional image observations: PPP-Cube (pick&place and push), Stack-Cube (stacking), Scene (complex object interaction), and Push-Tetris (block position/orientation control). All tasks are extended from known benchmarks (e.g., OGBench) and are characterized by combinatorial state spaces, challenging perception, and sparse rewards.

Figure 2: Environment suite encompasses PPP-Cube, Stack-Cube, Scene, and Push-Tetris, capturing a spectrum of manipulation and entity interaction complexities.

Numerical and Qualitative Results

HECRL (EC-SGIQL) consistently surpasses all baselines, often achieving more than 150% improvement in success rates over standard entity-centric RL and unstructured GCRL variants on the most challenging tasks, particularly PPP-Cube with image input. Numerical results are strong across both state- and image-based settings, with substantial gains in long-horizon sparse-reward environments.

Figure 3: Scene environment—factored subgoal generation and achievement. Top: latent reconstructions of subgoals with highlighted entity modifications; bottom: original environment images at subgoal time steps.

Figure 4: PPP-Cube and Push-Tetris—visualization of image-based subgoal achievement. Top: DLP reconstructions of latent subgoals; circles/arrows denote sparse entity modifications and corresponding policy effects.

Ablations demonstrate that value-based subgoal filtering and selection markedly dominate random or regression-based alternatives in both sample efficiency and success metric. The diffusion model’s entity factorization leads to an average modification of only 1–1.3 entities per subgoal, compared to nearly 3 for regression-based approaches.

Hyperparameter sensitivity analyses confirm robustness to variations in subgoal generation parameters ($K$, $N$, $T_{sg}$); performance only degrades when the subgoal interval is far out of phase with the underlying policy competence.

Figure 5: Success rates on PPP-Cube (Image) indicate stable hyperparameter sensitivity across subgoal generation training interval $K$, rollout duration $T_{sg}$, and candidate sample number $N$.

Subgoal Quality and Generalization

Qualitative analysis using DLP reconstructions demonstrates that subgoals generated by the diffusion model are typically sparse, compositional, and map closely to achievable states. AWR-based policies, in contrast, yield ambiguous or duplicated entity representations. HECRL additionally exhibits nontrivial zero-shot compositional generalization to increasing entity counts and horizon lengths—a key advantage in domains requiring scaling and transfer.

---

Theoretical and Practical Implications

HECRL’s modular architecture enables scalable, hierarchical RL in environments where entity factorization and object-centric representations are accessible. The use of conditional diffusion for subgoal generation addresses the multi-modality and combinatorial branching inherent in complex tasks, while decoupling subgoal and goal training targets avoidance of compounding approximation errors pervasive in temporal-difference learning for long-horizon GCRL.

Practically, the method is adaptable to contemporary RL algorithms exploiting new value functions or policy architectures, and is applicable to complex robotic and multi-agent scenarios with pixel-level observations, given adequate unsupervised object-centric representation learning. The approach is also robust to imperfect subgoal specification and supports flexible test-time tuning of policy competence radius.

---

Future Directions

Several lines of inquiry are promising for future research:

• Scaling Entity-Centric RL: Extending HECRL to domains with higher entity counts through improved factored state estimators, curriculum learning, or offline-to-online fine-tuning.
• Joint Representation and Policy Learning: End-to-end co-training of object-centric representations and RL components—mitigating reliance on representation quality bottlenecks.
• Subgoal Horizon Adaptation: Dynamic inference of the subgoal training interval $K$ from data or model feedback.
• Real-world Deployment: Bridging advances in unsupervised object-centric representations (e.g., DLP, slot attention, etc.) for robust transfer to physical robotics and vision-based agents.

---

Conclusion

HECRL introduces a hierarchical, modular, and entity-centric framework for offline GCRL in combinatorially complex, multi-entity environments. Through conditional diffusion-based subgoal generation and value-based selection, the approach achieves strong empirical performance and compositional generalization. Theoretical contributions include mitigation of compounding value errors and exploitation of structured state representations for scalable RL. The method’s flexibility and extensibility make it a viable candidate for tackling increasingly complex control and planning tasks in both simulated and real-world domains.