Laplacian Representations for Decision-Time Planning

AI summary

Overview

Research area: Model-based reinforcement learning, specifically representation learning and decision-time planning for offline goal-conditioned RL (GCRL).

Technical level: Advanced. The paper assumes familiarity with Markov decision processes, spectral graph theory, model-predictive control, and hierarchical RL.

Scope: The paper proposes a hierarchical decision-time planning algorithm, ALPS, built on the Laplacian representation of an environment, and evaluates it against model-free baselines on tasks from the OGBench benchmark.

What This Paper Is About

Planning with a learned model is hard because small prediction errors accumulate over long horizons, so imagined trajectories drift away from reality. The authors argue that the choice of state representation is the key to fixing this: if the latent space preserves genuine temporal distances between states, a planner can break a long-horizon task into shorter, more reliable subproblems. Their goal is to show that the Laplacian representation (built from eigenvectors of the graph Laplacian induced by the environment's dynamics) provides exactly this kind of latent space, and to use it to build a planner that outperforms model-free methods.

Key Contributions

1. A demonstration that the scaled Laplacian representation is a suitable latent space for planning. The authors show that Euclidean distance in the scaled Laplacian space, ψ-space, approximates the commute time distance (CTD) of the underlying graph, and that the representation naturally partitions the environment at its bottlenecks, making it useful both as a low-level cost metric and as a basis for high-level subgoal clustering.

2. ALPS (Augmented Laplacian Planning with Subgoals), a hierarchical decision-time planning algorithm. ALPS learns the Laplacian representation with the ALLO objective, learns a one-step forward model in the original state space via a multi-step auto-regressive objective, learns a behavior prior, clusters the dataset with k-means in ψ-space to build a cluster graph, and then plans at decision time with Dijkstra over cluster centers plus CEM for low-level action selection.

3. A behavior prior that accelerates planner convergence. Instead of sampling action sequences from an unconditional Gaussian, ALPS uses a learned goal-conditioned deterministic policy to propose a mean action sequence, around which temporally-correlated Gaussian noise is added, with a cost function that also penalizes large actions.

4. Empirical results on OGBench and Maze2D—PointMass. ALPS outperforms model-free GCRL baselines on a suite of OGBench tasks, a benchmark the paper describes as previously dominated by model-free methods, and also outperforms PcLast on Maze2D—PointMass.

Main Findings

• The scaled Laplacian space captures environment geometry at multiple time scales. In a visualization of the pointmaze-large environment, the CTD approximation produces distances that are temporally consistent, and clustering in ψ-space produces tight clusters that respect the environment dynamics. Early eigenvectors encode global structure such as rooms or regions, while later ones capture increasingly local distinctions.

• ALPS beats PcLast on Maze2D—PointMass. In Table 1, using 1 cluster (low-level planner only) and 16 clusters (both planners), ALPS† scores 94±3 and 97±2 on Hallway, 92±3 and 96±2 on Rooms, and 91±4 and 94±2 on Spiral, compared to PcLast's 51±4 and 62±4, 30±3 and 57±10, and 35±4 and 60±6, respectively. These results are reported as mean and standard deviation over 10 seeds.

• Both planners together help. Both algorithms performed better across all three Maze2D domains when the high- and low-level planners were used jointly. PcLast showed a substantial performance reduction when the high-level planner was removed, while ALPS† performed similarly in the low- and high-level settings on the Spiral domain.

• ALPS outperforms model-free baselines on OGBench. According to a two-sided paired Wilcoxon signed-rank test applied to per-domain mean performance with Holm-Bonferroni correction, ALPS is significantly better with p < 0.001. The paper notes that all GCRL baselines considered are model-free.

• Large margins on the hardest "giant" and "explore" tasks. Examples from Table 2 include pointmaze-giant-stitch-v0, where ALPS reaches 98±1 while GCBC, GCIVL, GCIQL, CRL, and HIQL all report 0±0; antmaze-giant-stitch-v0, where ALPS reaches 92±3 versus HIQL's 2±2; antmaze-medium-explore-v0, where ALPS reaches 100±0 versus HIQL's 37±10; and visual-antmaze-giant-stitch-v0, where ALPS reaches 55±6 while every baseline reports 0±0.

• ALPS is more robust to maze size. The authors state that ALPS is more robust to the size of the maze the agent navigates, though the excerpt of the paper provided cuts off mid-sentence at this point.

• ALPS does not win everywhere. It is beaten on several tasks, among them cube-double-play-v0 (ALPS 2±1 versus GCIQL 40±5), scene-play-v0 (ALPS 26±3 versus GCIQL 51±4), humanoidmaze-medium-stitch-v0 (ALPS 68±5 versus HIQL 88±2), pointmaze-teleport-stitch-v0 (ALPS 13±4 versus GCIVL 44±2), pointmaze-teleport-navigate-v0 (ALPS 40±6 versus GCIVL 45±3), antmaze-teleport-navigate-v0 (ALPS 45±3 versus CRL 53±2), antmaze-teleport-stitch-v0 (ALPS 35±11 versus GCIVL 39±3), and visual-antmaze-teleport-stitch-v0 (ALPS 21±4 versus HIQL 37±2). On cube-single-play-v0 it ties GCIQL at 68±6.

Methodology in Plain English

The agent is given a fixed offline dataset of trajectories, with no interaction allowed during learning. The approach has a pre-training phase and a planning phase.

During pre-training, four things are learned from the dataset. First, the Laplacian representation is learned using the ALLO objective, which is optimized with stochastic gradient descent on sampled transition pairs, avoiding the cubic cost of an explicit eigendecomposition. Second, a one-step forward model is trained in the original state space with a multi-step auto-regressive loss over a horizon, so gradients flow back through several predicted steps and the model is discouraged from drifting. Third, a behavior prior is trained by behavior cloning to predict an action that moves from a current state toward a state some number of steps in the future, expressed in terms of the scaled Laplacian representation. Fourth, k-means clustering is applied in the scaled Laplacian space to partition the dataset into regions, and a graph over cluster centers is built from observed inter-cluster transitions, with infrequent transitions pruned by nucleus sampling.

At decision time, the start and goal states are mapped into ψ-space and assigned to clusters. Dijkstra's algorithm finds a shortest path over the cluster graph, giving a sequence of subgoals. At each step, the low-level planner uses the Cross-Entropy Method: the behavior prior proposes a mean action sequence, candidate sequences are sampled around it with temporally-correlated Gaussian noise, each candidate is rolled out through the forward model, and candidates are scored by a cost that combines distance to the current subgoal in ψ-space with a penalty on large actions. The best-scoring sequences update the sampling distribution over several iterations, and the first action of the best sequence is executed. If the agent leaves the planned path, the high-level plan is recomputed.

Experiments use the Maze2D—PointMass tasks, where PcLast was originally evaluated, and the locomotion and manipulation tasks from OGBench. Maze2D—PointMass uses a 100×100 single-channel image observation, actions within [-0.2, 0.2] per coordinate, requires reaching within 0.03 units of the target within 30 actions, and is trained on 500K transitions from a uniform random policy. OGBench locomotion covers pointmaze, antmaze, and humanoidmaze, controlling a 2-DoF ball, 8-DoF ant, and 21-DoF humanoid. Manipulation covers Cube and Scene tasks with a 6-DoF UR5e arm and a Robotiq 2F-85 gripper in a 5-D action space. Pixel-based variants use 64×64×3 third-person images. Performance is averaged over five pre-defined state-goal pairs, with 8 seeds for state-based and 4 seeds for pixel-based tasks, and the OGBench protocol averages over 750 rollouts (3 evaluation epochs × 5 test-time goals × 50 rollouts). Baselines are GCBC, GCIVL, GCIQL, QRL, CRL, and HIQL.

Why This Matters

The paper's central claim is that representation choice, not just model accuracy, is what limits decision-time planning with learned models. If a latent space preserves a meaningful notion of temporal distance, long-horizon problems can be decomposed into subgoals automatically, which mitigates the compounding-error problem that has kept model-based methods behind model-free ones on benchmarks like OGBench. The paper also shows that a learned Laplacian representation can be scaled to large state and action spaces, supporting a shift away from the assumption that model-free methods are the default for offline GCRL.

Real-world applications (plausible, based on the task domains studied):

• Robot navigation in mazes and buildings. The pointmaze, antmaze, and humanoidmaze tasks involve navigating large, walled environments where Euclidean proximity is misleading, a setting that maps onto indoor or warehouse navigation.
• Robot manipulation. The Cube and Scene tasks cover pick-and-place, sequencing operations, opening drawers, and toggling lock states, which resemble tabletop assembly and household manipulation.
• Learning from large pre-recorded demonstration datasets. The offline setting means a policy can be trained from logged data without risky or expensive online interaction.
• Vision-based control. The pixel-based variants use only third-person images, relevant to robots that rely on cameras rather than privileged state information.

Industry relevance: the approach targets domains where collecting online experience is costly or unsafe but large logged datasets exist, such as manufacturing robotics, warehouse automation, and simulation-to-reality pipelines. The paper provides code at https://github.com/machado-research/ALPS and a project page at https://dikshuy.github.io/ALPS/.

Future Directions

• Address the tasks where ALPS underperforms. ALPS falls behind baselines on cube-double-play-v0, scene-play-v0, humanoidmaze-medium-stitch-v0, and several teleport variants; understanding why, and whether the subgoal decomposition or the forward model is responsible, is an open question.
• Improve robustness of the forward model. The paper's focus on compounding errors suggests further work on stabilizing multi-step predictions or replacing the one-step model with a better calibrated alternative.
• Investigate the scaling behavior more fully. The paper claims robustness to maze size but the provided text is truncated mid-statement, so the full extent of that claim and its analysis are not visible here.
• Extend the Laplacian representation to settings with changing dynamics or reward structure. The representation is learned from a fixed offline dataset, and how it adapts if the environment or task distribution changes is not reported.

Target Audience

This paper is best suited to reinforcement learning researchers and graduate students working on model-based RL, representation learning, or hierarchical planning, as well as practitioners applying offline RL to robotics who are familiar with benchmarks such as OGBench and with spectral methods. Readers without a background in MDPs, graph Laplacians, or model-predictive control will find the formalism dense, although the high-level idea of using environment geometry to pick better subgoals is accessible.

Authors’ abstract

Planning with a learned model remains a key challenge in model-based reinforcement learning (RL). In decision-time planning, state representations are critical as they must support local cost computation while preserving long-horizon structure. In this paper, we show that the Laplacian representation provides an effective latent space for planning by capturing state-space distances at multiple time scales. This representation preserves meaningful distances and naturally decomposes long-horizon problems into subgoals, also mitigating the compounding errors that arise over long prediction horizons. Building on these properties, we introduce ALPS, a hierarchical planning algorithm, and demonstrate that it outperforms commonly used baselines on a selection of offline goal-conditioned RL tasks from OGBench, a benchmark previously dominated by model-free methods.