SAME: Learning Generic Language-Guided Visual Navigation with State-Adaptive Mixture of Experts

In plain words

A navigation agent gets a sentence such as "walk to the kitchen" or "find a chair", then moves through a building by looking around at each step. Earlier work built a separate agent for each kind of instruction. The authors show that one agent trained on all of them at once often does worse on the harder ones, because the kinds of instruction pull the agent in different directions.

Their fix is a set of small specialist sub-networks, called experts, inside the agent. At each step a small router looks at what the agent currently sees and what it was told, then picks a couple of experts to use for that step. Easy exploration and careful instruction-following can then use different specialists, while the shared parts of the agent keep the common knowledge.

The result is a single model that handles seven navigation benchmarks. The authors report that it matches or beats models built for one task, though not on every metric.

TL;DR

• SAME is a mixture-of-experts layer on the visual query of the cross-attention in the DUET navigation agent §3.2 §4 app., DUET revisit.
• A router picks the top experts per navigation step from the mean of the panorama views and the instruction [CLS] token §3.2.
• Trained with DAgger on seven VLN datasets in the discrete Matterport graph, with a load-balancing loss weighted 0.8 §4.
• Author-reported: best unified-model results on several splits, but separate single-task models still lead on REVERIE tab. 4.
• Code is MIT-licensed and mostly matches the paper, with three notable differences in routing softmax, balance-loss scaling and the data mix (see Paper vs code).

Problem

Language-guided navigation tasks differ in how detailed their instructions are. Fine-grained tasks (R2R) describe every step, coarse-grained tasks (REVERIE) name a target object and a room, and ObjectNav gives only a category §2 fig. 1. Training one model on a mix of these datasets gives inconsistent results: the authors report that mixing data cuts success rate on R2R by 2-3 points and on REVERIE by 6-7 points in their held-out tests §2.1.

The paper also finds that common mixture-of-experts routing does not fit this setting. Routing by token [shazeer2017outrageously] or by task identity gave lower REVERIE success than the state-adaptive routing in single-run tests tab. 2, so the authors design routing from the agent's current state §3.2.

Method

The base model is DUET, a dual-scale graph transformer. It builds a local branch on the current panorama and a global branch on the topological map of visited and candidate nodes, and both branches use cross-attention over the instruction app., DUET revisit. SAME does not change the graph or the action head. It changes the cross-attention layers.

1. Experts in a sparse MoE layer §3.1. A layer holds N experts and a linear router. The router turns a routing feature into one score per expert, a softmax gives probabilities, and the output is the probability-weighted sum over the top-k experts §3.1, MoE formulation. The paper gives the formula but does not state N or k; the code uses N = 8 and k = 2 [code: src/configs/default.yaml#L60-L61].

2. Routing feature (SAME routing) §3.2. The routing feature merges the mean of the visual features over all views with the instruction [CLS] token, through one linear layer W_m that maps 2h to h. A variant adds a task embedding. In the single-run ablations this variant has lower SPL than the plain routing on all three tasks §3.3 tab. 2.

3. Where the experts sit §3.4 fig. 2. The experts replace the query projection of the visual cross-attention in every cross-attention layer of DUET's local and global branches vilmodel.py. The paper compares this with experts on the key and value projections, and with experts in the feed-forward network. Experts on the visual query gave the best SPL on all three ablation tasks tab. 3.

4. Training objective §4. DAgger, not plain imitation: at each step the agent samples its action from its own predicted distribution, and the teacher's action supervises that choice §4 §4.2, training schema. The loss is the DAgger navigation loss plus lambda times the MoE load-balancing loss, with lambda = 0.8 §4 tab. 7. The load-balancing loss follows the Switch Transformer form [fedus2022switch] §3.1.

5. Data mix §4. The training sets are ScaleVLN, R2R, RxR-EN, CVDN, REVERIE, SOON and Habitat-Web (ObjectNav demonstrations), sampled in the ratio 10:1:1:1:1:1:2. A batch never mixes datasets §4. The model starts from ScaleVLN pre-trained weights §3.3 §4.

# SAME routing for one navigation step  [§3.2]
# Input:  views (B, V, h), instruction_cls (B, h)
# Output: weights (B, k), experts (B, k), logits (B, N)

def route_step(views, instruction_cls):
    # Average the view features into one state vector  [§3.2]
    visual_mean = mean(views, over=V)                     # (B, h)
    # Merge the visual state with the instruction token  [§3.2]
    routing_feature = W_m(concat(visual_mean, instruction_cls))
    # Score every expert for each sample  [§3.1]
    logits = W(routing_feature)                           # (B, N)
    # Turn scores into probabilities over the experts  [§3.1]
    probs = softmax(logits, over=N)
    # Keep the k most likely experts and their weights  [§3.1]
    weights, experts = top_k(probs, k)
    return weights, experts, logits

# Visual query with MoE  [§3.1] [§3.4]
# Input:  x (B, L, h), weights (B, k), experts (B, k)
# Output: query (B, L, h)

def moe_query(x, weights, experts, expert_linears):
    # Sum the weighted outputs of the selected experts  [§3.1]
    query = zeros_like(x)
    for i in range(N):
        # Samples that routed to expert i, with that sample's weight  [§3.1]
        rows, slots = where(experts == i)
        query[rows] += weights[rows, slots] * expert_linears[i](x[rows])
    return query

# Training loss for one batch of episodes  [§4] [§3.1]
# Input:  dagger_losses (T,), balance_terms (T, N), lam = 0.8
# Output: scalar

def batch_loss(dagger_losses, balance_terms, lam):
    # Navigation loss summed over steps, as the paper's objective  [§4]
    nav_loss = sum(dagger_losses)
    # Load-balance term from the routing decisions  [§3.1]
    balance = sum(balance_terms)
    return nav_loss + lam * balance

Balance loss, as the paper defines it §3.1. For each routing batch of K inputs with N experts, F is the fraction of inputs whose top-1 expert is i, and D is the mean router probability of expert i. The loss is N times the sum over experts of F_i times D_i. With uniform routing it equals 1. The code uses a different count, see Paper vs code.

# MoE load-balance term for one routing batch  [§3.1] [eq. balance]
# Input:  probs (K, N) router probabilities
# Output: scalar

def balance_term(probs, num_experts):
    # Fraction of inputs whose top-1 expert is each expert  [§3.1]
    top1 = argmax(probs, over=N)
    fraction = mean(one_hot(top1, num_experts), over=K)   # (N,)
    # Mean router probability given to each expert  [§3.1]
    mean_prob = mean(probs, over=K)                       # (N,)
    return num_experts * sum(fraction * mean_prob)

The paper's loss equation is labelled equ:balance_loss in the LaTeX source; the equation numbers in the PDF were not checked, so the card cites the section.

Hyperparameters

Results

Author-reported. Setup: discrete Matterport3D graph, DUET backbone initialised from ScaleVLN, multi-task training on seven datasets §4.

Main results (discrete environment) tab. 4. SAME row against the unified baselines and the best single-task model, val unseen unless marked.

† the authors' multi-task run of ScaleVLN. The full-results table in the appendix gives SAME's TL and NE for each benchmark app., full-results table.

Zero-shot on continuous Habitat tab. 5. The same checkpoint, evaluated with a waypoint predictor and no training in the continuous environment §4.

For context, the paper lists ScaleVLN trained on R2R-CE at SR 55 and SPL 51, and Habitat-Web at MP3D SR 35 and SPL 10 tab. 5.

Routing comparison, single runs tab. 2. Setup: the Section 3 experiments, experts on visual query, pre-trained from ScaleVLN. Columns: TL, NE, SR, SPL.

Expert position, single runs tab. 3.

Training schema, single runs tab. 6.

Balance coefficient, single runs tab. 7.

Data mixture for DUET, held-out results tab. 1. SR (SPL in brackets) on R2R, REVERIE and ObjectNav MP3D when DUET is trained on the listed sets. Entries marked (zs) are zero-shot, meaning the dataset was not in training. Habitat-rendered images are used for this experiment.

The single-run ablations (tables 2, 3, 6, 7) use the Section 3 setup and are not directly comparable with the main results in table 4.

Paper vs code

Repository: https://github.com/GengzeZhou/SAME, MIT, at commit b74c57bbb189a0ed8c908e90b53d02dfbe417528. The authors are listed in the paper's author block, so this is their repository.

• Routing softmax runs over the batch, not the experts. Paper: softmax over the N experts for each sample §3.1. Code: the task router applies softmax along dimension 0, which is the batch dimension, and then takes the top-k vilmodel.py. Per-sample ranking can change because the normaliser differs per expert.
• Top-k weights are renormalised. Paper: the output uses the unnormalised router probabilities of the top-k experts §3.1. Code: the top-k weights are divided by their sum, so they add to 1 vilmodel.py.
• Balance loss scale. Paper: F uses the top-1 expert, and the loss is 1 at uniform routing §3.1. Code: the one-hot mask covers all k selected slots, and the sum runs over those slots, so the loss is k at uniform routing (2 with k = 2) duet_agent.py [code: src/tasks/agents/duet_agent.py#L323-L324]. The code's default balance weight of 0.8 therefore has an effective scale about twice the paper formula.
• Balance term is summed over navigation steps. Paper: one loss over the routing batch §3.1. Code: a balance term is added at every step of every episode, then multiplied by 0.8 and divided by the batch size, the same way as the navigation loss duet_agent.py duet_agent.py.
• Extra teacher-forcing term. Paper: DAgger loss plus lambda times the balance loss §4. Code: a teacher-forcing (imitation) loss with coefficient 0.2 is added to the DAgger loss in each step default.yaml base_trainer.py.
• Data mix. Paper: seven sets at 10:1:1:1:1:1:2 §4. Code main config: nine sources at 10:1:1:1:1:1:1:1:2, adding R2R_PREVALENT and REVERIE_SCALEVLN main_multi_q.yaml. The default config uses 20:1:1:10:1:1:1:1:2 default.yaml.
• Training length and batch. Paper: 20k iterations at batch 16 §4. Code default: 500000 iterations at batch 8 default.yaml. The main multi-task configs set batch 16 and do not override the iteration count.
• Experts N and k. Not stated in paper; code uses 8 and 2 [code: src/configs/default.yaml#L60-L61].
• Pre-trained checkpoints. Paper: initialised from ScaleVLN §3.3. Code: each MoE position loads its own pre-trained checkpoint, for example Attnq_pretrained_ckpt.pt for the visual-query variant [code: src/configs/main_multi_q.yaml#L17]. The README says these are released on HuggingFace README.md. The paper does not describe how they were produced.
• Paper claims against its own tables. The text says lambda = 0.2 "performs worse than all other variants" §4.2 tab. 7. On R2R SPL and ObjectNav SPL, lambda = 0.2 is the best row or close to it (65.73 and 43.51); on REVERIE SR it is the lowest (42.86) tab. 7. The claim holds only for some metrics.
• Average SPL. The text reports an average SPL of 48.31% §3.3. The mean of the three SPL values in the routing-comparison table for SAME is 48.75 tab. 2. The paper does not say how the average was taken.

No discrepancies were found for the navigation policy structure, the cross-attention placement, or the routing feature merge (mean views plus instruction [CLS] through a linear layer) vilmodel.py.

Pitfalls

• Routing details. The softmax axis and the top-k renormalisation matter for reproduction. Follow the code if you want to match the reported numbers, and the paper if you want the stated formulation. Report which one you used.
• Balance coefficient scale. Because the code's balance term is k times the paper formula, lambda = 0.8 in code is not the same number as lambda = 0.8 in the paper formula. Matching the paper's ablations needs the paper's form.
• N and k are unstated. The paper gives no number of experts and no k. The code's values (8 and 2) are the only source default.yaml.
• Pre-training is a prerequisite. Table 2 shows that without VLN pre-training the SR drops to 59.05 on R2R, from 73.69 with it tab. 2. The paper's pre-training is ScaleVLN, and the MoE-position checkpoints in the code are not described in the paper.
• Ablations are single runs. Tables 2, 3, 6 and 7 are single runs of the Section 3 setup. Differences of about one point are within noise. The paper does not report seeds or variance.
• Mixing data hurts DAgger. Table 6 shows that batch mixing reduces DAgger results sharply (R2R SR 62.36 against 73.69 sequential) tab. 6. Keep datasets in separate batches, as the paper does §4.
• Continuous results are zero-shot. The continuous Habitat numbers come from a discrete-trained model with a waypoint predictor §4 tab. 5. They are not a continuous-environment training result.
• Missing details for a faithful rebuild. Not stated in the paper: hidden sizes, the exact teacher labels for DAgger, gradient clipping, the teacher-forcing weight, dropout, max action length, and the scheduler. The code has values for all of these (see Hyperparameters), but the paper does not say they were used.
• Compute. Single 80 GB A100 for training §4. The paper does not report training time, memory or inference cost for SAME against DUET.
• Visual data. The ablations use Matterport3D-rendered images for R2R and REVERIE §3.3, while the Table 1 study uses Habitat-rendered images tab. 1. Keep the rendering source in mind when comparing numbers.

Lineage

• Builds on: DUET (Chen et al., CVPR 2022, "Think Global, Act Local"), the base navigation model, extended with MoE on its cross-attention app., DUET revisit.
• Builds on: ScaleVLN, used for pre-training and data §3.3 §4.
• Builds on: mixture-of-experts routing from Shazeer et al. (2017) and Switch Transformers (Fedus et al.), used for the balance loss §3.1.
• Compared against: MT-RCM, NaviLLM, ScaleVLN† (multi-task), and separate single-task models such as ScaleVLN, BEVBert, VER and GOAT tab. 4.
• Superseded by: not known from content already in this wiki.