---
title: 'COMPASS: RL with Latent Space Search'
url: https://www.emergentmind.com/topics/compass-combinatorial-optimization-with-policy-adaptation-using-latent-space-search
type: topic
---

# COMPASS: RL with Latent Space Search

COMPASS (COMbinatorial Optimization with Policy Adaptation using Latent Space Search) is a reinforcement learning (RL) framework for combinatorial optimization that parameterizes a distribution over specialized policies conditioned on a continuous latent space. The approach is designed to address the limitations of standard RL heuristics for solving NP-hard combinatorial optimization tasks, particularly their restricted search diversity at inference and limited generalization under distribution shift. COMPASS introduces a pre-training procedure that anticipates instance-level search, enabling efficient adaptation to new or out-of-distribution problem instances at test time.

## 1. Problem Setting and Motivations

COMPASS targets classical NP-hard combinatorial optimization (CO) problems, including:

- **Travelling Salesman Problem (TSP):** Find a minimum-length tour visiting each of \(n\) cities exactly once.
- **Capacitated Vehicle Routing Problem (CVRP):** Route a fleet of vehicles with limited capacity to serve all demands at minimum cost.
- **Job-Shop Scheduling Problem (JSSP):** Sequence operations on machines to minimize overall makespan.

The cardinality of the feasible set increases exponentially with instance size \(n\), precluding exact enumeration for large-scale tasks. While RL has been investigated as a framework for learning construction heuristics, prior work predominantly focuses on training a single policy, limiting the diversity and adaptability of solutions at inference. Post-training search strategies either rely on stochastically sampling many trajectories from the same policy (yielding minor variations) or directly fine-tuning the policy on novel instances—an approach that incurs significant computational overhead. COMPASS addresses these deficiencies by pre-training a family of policies conditioned on a continuous latent space and employing principled search (CMA-ES) to adapt solutions at inference.

## 2. Reinforcement Learning Formulation and Latent Space Conditioning

In COMPASS, each combinatorial optimization problem is modeled as an episodic Markov Decision Process \((S, A, T, R, H)\), where the agent constructs a solution step-by-step. The policy gradient objective is given by

\[
J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ R(\tau) \right], \qquad R(\tau) = \sum_{t=0}^{H} \gamma^t r(s_t, a_t)
\]

with policy \(\pi_\theta(a \mid s)\). COMPASS introduces a continuous latent variable \(z \in \mathbb{R}^{16}\), sampled uniformly over \([-1,1]^{16}\), so the conditional policy becomes \(\pi_\theta(a \mid s, z)\). During training, for each instance, \(N\) latent vectors \(\{z_i\}_{i=1}^N\) are sampled, and each \(\pi_\theta(\cdot \mid z_i)\) is rolled out to produce a trajectory. The “best” latent, \(z_{i^*}\), is the one yielding the highest trajectory reward. Only \(z_{i^*}\) is used to update parameters via REINFORCE:

\[
\nabla_\theta J_{\rm compass} = \mathbb{E}_{\rho \sim \mathcal{D}} \, \mathbb{E}_{z_1, ..., z_N \sim p(z)} \, \mathbb{E}_{\tau_i \sim \pi_\theta(\cdot | z_i)}\left[\nabla_\theta \log \pi_\theta(\tau_{i^*} | z_{i^*}) (R_{i^*} - \mathcal{B})\right]
\]

where \(\mathcal{B}\) is a learned baseline. This best-of-\(N\) strategy ensures regions of latent space specialize towards different kinds of instances, promoting implicit policy diversity.

## 3. Policy Architecture and Parameterization of Diversity

The policy is realized as a Transformer-based encoder-decoder (adopting the architecture of POMO):

- **Encoder:** Maps problem-specific features (e.g., coordinates, graph structure) to node embeddings.
- **Decoder:** Generates solutions by attending to embeddings and the current partial solution.

Latent conditioning is integrated by concatenating the sampled \(z\) or its affine transformation to each key, query, and/or value vector across decoder attention layers. Therefore, each distinct latent vector defines a different policy within an infinite family parameterized by \((\theta, z)\).

At inference, the space \([-1, 1]^{16}\) is explored, where each \(z\) yields a different specialized heuristic. This formulation enables the representation of a vast repertoire of implicitly learned heuristics, facilitating adaptation beyond the representational reach of a single global policy.

## 4. Inference-Time Latent Space Search: Covariance Matrix Adaptation

For a new problem instance \(\rho\), the goal is to efficiently identify the latent \(z\) producing the highest-quality solution using \(\pi_\theta(\cdot | z)\). COMPASS applies parallel Covariance Matrix Adaptation Evolution Strategy (CMA-ES) components in the latent space:

- **Initialization:** Multiple components \(\{(\mu_j, \Sigma_j)\}\) are initialized at Voronoi centroids of \([-1,1]^{16}\).
- **Iterative Search:**
  1. For each component \(j\), sample \(M\) candidates \(z_{j,m} \sim \mathcal{N}(\mu_j, \Sigma_j)\).
  2. Evaluate rollouts \(\pi_\theta(\cdot|z_{j,m})\) and receive rewards \(R_{j,m}\).
  3. Update \((\mu_j, \Sigma_j)\) using standard CMA-ES update based on \(R_{j,m}\).
- **Selection:** After a fixed rollout budget \(B\), return the best solution found.

The computational bottleneck is dominated by \(O(B)\) forward model rollouts; CMA-ES logic is negligible in wall time. This search protocol enables rapid adaptation per instance, substantially improving solution quality and generalization.

#### Inference-Time CMA-ES Pseudocode

```
Initialize components { (μ_j, Σ_j) }_{j=1}^C via Voronoi centroids.
best_solution ← ∅
for budget_step = 1…B/C do
  for each component j=1…C:
    sample {z_{j,m}_{m=1…M} ∼ N(μ_j,Σ_j)
    for m=1…M: evaluate R_{j,m} = rollout(πθ(·|z_{j,m}))
    update (μ_j,Σ_j) by CMA−ES using {R_{j,m}
    record best_solution if any R_{j,m} better
return best_solution
```

## 5. Empirical Evaluation and Benchmarks

COMPASS was evaluated against standard RL baselines (POMO, Poppy, EAS—active search fine-tuning) and industrial solvers (Concorde/LKH3 for routing, OR-Tools for scheduling) on eleven canonical benchmarks, each using an equal rollout budget of 1,600 samples per instance. Table 1 summarizes primary results for representative instance sizes:

| Task                 | Best Industrial | POMO     | Poppy    | EAS      | COMPASS  |
|----------------------|----------------|----------|----------|----------|----------|
| TSP\(_{100}\)        | 7.765 (0.000%) | 7.779 (0.185%) | 7.766 (0.013%) | 7.778 (0.161%) | 7.765 (0.002%) |
| CVRP\(_{100}\)       | 15.650 (0.000%)| 15.713 (0.399%)| 15.663 (0.084%)| 15.663 (0.081%)| 15.594 (−0.361%)|
| JSSP \(10\times10\)  | 807.6 (0.00%)  | 862.1 (6.7%)   | 849.7 (5.2%)   | 858.4 (6.3%)   | 845.5 (4.7%)   |

For larger instances (TSP 125/150/200, CVRP 125/150/200, JSSP 15×15/20×15), COMPASS consistently achieved the best performance among RL-based methods and often closed much of the gap to industrial solvers.

## 6. Distribution Shift and Robustness Analysis

Robustness is evaluated on 180 out-of-distribution (OOD) test sets generated by applying nine geometric or structural mutation operators (explosion, implosion, cluster, rotation, linear/axis projection, expansion, compression, grid), each at ten mutation strengths. COMPASS outperformed all compared methods on all OOD distributions.

For example, under the strongest TSP mutation (power = 0.9), average tour lengths and optimality gaps were:

| Solver     | Tour Length | Gap (%) |
|------------|-------------|---------|
| LKH3       | 6.732       | 0.000   |
| POMO       | 6.770       | 0.557   |
| Poppy 16   | 6.744       | 0.176   |
| EAS        | 6.753       | 0.308   |
| COMPASS    | 6.738       | 0.091   |

COMPASS’s degradation in gap under increasing mutation power was substantially slower than other RL-based approaches, indicating superior generalization under strong distribution shift.

## 7. Key Insights, Limitations, and Prospective Extensions

- **Conditional Policy Family:** Training with a continuous latent space produces an implicit infinite family of specialized policies, each accessible by selecting an appropriate \(z\).
- **Instance Specialization:** The best-of-\(N\) objective induces structure in latent space, with neighborhoods corresponding to problem instance sub-types.
- **Efficient Adaptation:** CMA-ES in latent space enables principled, computationally efficient search for high-performing policies per instance, precluding the need for online backpropagation.

**Limitations:**

- Diversity is only implicit: No explicit diversity regularization is incorporated during training, potentially limiting latent coverage.
- Training inefficiency: Uniform latent sampling can, in principle, miss optimal conditionings for instances.
- Noisy search landscape: The latent space can be noisy; Bayesian optimization methods were less effective than CMA-ES within rollout budgets.

**Future research directions include:**

- Incorporating an unsupervised diversity bonus (e.g., VAE-style entropy regularization) to further encourage latent space coverage.
- Learning conditional priors or amortized inference networks \(q(z|\rho)\) to propose promising latent candidates per instance.
- Applying latent-space regularization such that smooth changes in \(z\) yield gradual policy variations.
- Combining COMPASS’s search with fine-tuning (as in EAS) for enhanced performance, particularly on large-scale instances.

COMPASS represents a methodological advance for RL-based combinatorial optimization by synthesizing a parametric diversity of heuristics during training and leveraging efficient black-box search during inference to realize superior performance both in-distribution and under substantial distribution shifts.

Source: https://www.emergentmind.com/topics/compass-combinatorial-optimization-with-policy-adaptation-using-latent-space-search