NoMaD: Goal Masked Diffusion Policies for Navigation and Exploration

In plain words

A robot that is sent into an unfamiliar building has two jobs. It must wander to discover where things are, and it must walk to a place once someone shows it a photo of that place. Usually two separate systems handle these jobs, one for searching and one for reaching a goal.

NoMaD uses one neural network for both. It looks at the robot's recent camera frames and, optionally, at a photo of the goal. When the goal photo is hidden, the network explores. When the photo is shown, the network heads toward the place in the photo. The network's output is a short planned path. That path comes from a diffusion model, which starts from random noise and refines it step by step into a plausible sequence of moves. This lets the model keep several reasonable options open, for example at a junction, instead of blending them into one unhelpful average.

The authors use this single policy inside a simple map of places the robot has visited. They report higher exploration success and fewer collisions than the alternatives they tested, and goal reaching on par with the best of them.

TL;DR

• One Transformer policy, trained with the goal token masked for half of the samples, does both goal-reaching and undirected exploration §IV-A.
• Actions are a diffusion model over short 2D sequences, sampled with a 10-step denoiser conditioned on the observation context §IV-B §IV-C.
• Author-reported real-robot results beat the best published baseline on exploration success and collisions, with fewer parameters tab. 1.
• Several numbers and details are absent from the paper and must be taken from the code. The code also departs from the paper in a few places. See the Pitfalls and Paper vs code sections.

Problem

Goal-conditioned visual navigation policies work well when the target is known, but they cannot explore on their own §I §III. Earlier systems pair a goal policy with a separate high-level module that proposes subgoals, such as a generative image model, a planner or a latent-variable model §I §II. The authors say this adds complexity and often needs task-specific parts §I. Their closest baseline, ViNT, pairs a goal-conditioned policy with a 300M-parameter image diffusion model that proposes subgoal images §II; Table I lists the whole subgoal system at 335M parameters tab. 1.

Exploration needs a policy that can represent several possible behaviours at once, for example turning left or right at an intersection, and avoid obstacles at the same time §IV-B. The authors argue that regression-style or simple latent-variable policies struggle with such multimodal action distributions §II §V-A.

Method

NoMaD has two parts §IV. The first is attention-based goal masking, which decides whether the goal image reaches the rest of the network §IV-A fig. 2. The second is a diffusion policy that models the distribution of future actions given the observation context §IV-B. The architecture follows ViNT: EfficientNet-B0 encoders turn the observation frames and the goal into tokens, and a small Transformer fuses them into a context vector §III §IV-C fig. 2.

Section and figure numbers below follow the paper's print order: sections use Roman numerals (I–VI) with subsections A–C, tables and figures are numbered in the order they appear.

Topological planning around the policy. The policy is used inside the ViKiNG setup. The robot builds a topological graph of visited observations, and a frontier-based planner picks exploration targets. The policy supplies the local actions §III. The paper does not describe the planner's internals beyond this, so the cited ViKiNG setup is needed for full detail.

Goal masking §IV-A. The goal token is a feature from a goal fusion encoder. The paper defines a binary mask m, where m = 1 hides the goal token from attention and m = 0 keeps it. During training m is drawn from a Bernoulli distribution with probability p_m = 0.5. At test time m = 1 gives undirected exploration and m = 0 gives goal reaching. The masked token is excluded from downstream attention rather than multiplied by zero.

Diffusion policy §IV-B. The policy models p(a_t | c_t), where c_t is the observation context after goal masking. Training adds noise to ground-truth action sequences and trains a noise predictor to recover the noise. Sampling starts from Gaussian noise and runs K denoising steps. The paper writes one update as eq. 1, with α, γ and σ described only as functions of the noise schedule §IV-B:

a^{k-1}_t = \alpha \cdot \bigl(a^k_t - \gamma\,\epsilon_\theta(c_t, a^k_t, k) + \mathcal{N}(0, \sigma^2 I)\bigr) eq. 1

The noise predictor ε_θ is conditioned on c_t and is a 1D conditional U-Net §IV-C. The paper models the action distribution conditionally and leaves the context out of the denoised output, which the authors say allows real-time control and end-to-end training of the encoder §IV-B.

Training objective §IV-C. The loss has two terms: a noise-prediction MSE and a temporal-distance MSE weighted by λ = 10^-4 eq. 2. Here f_d is the distance head and d is the temporal distance between observation and goal:

\mathcal{L} = \mathrm{MSE}\bigl(\varepsilon^k,\; \varepsilon_\theta(c_t,\, a^0_t + \varepsilon^k,\, k)\bigr) + \lambda \cdot \mathrm{MSE}\bigl(d(o_t, o_g),\, f_d(c_t)\bigr) eq. 2

# One NoMaD training step  [§IV-C] [eq. 2]
# Input:  obs_images (B, 3P, H, W), goal_image (B, 3, H, W)
#         actions (B, T, 2) positions, distance (B,) in steps
# Output: scalar loss

def training_step(obs_images, goal_image, actions, distance):
    # Mask the goal token for about half of the samples  [§IV-A]
    goal_mask = bernoulli(p=0.5, size=B)
    context = vision_encoder(obs_images, goal_image, goal_mask)

    # Predict the temporal distance; use goal-conditioned samples only
    dist_pred = distance_head(context)
    dist_loss = masked_mean(mse(dist_pred, distance), 1 - goal_mask)

    # Per-step displacements, normalized to [-1, 1]  [§IV-C]
    steps = normalize(diff(actions))
    # Sample a denoising step and Gaussian noise  [§IV-B] [§IV-C]
    k = randint(0, K)
    noise = normal(0, 1, shape=steps.shape)
    noisy_steps = scheduler.add_noise(steps, noise, k)

    # Predict the noise and compare it with the true noise  [eq. 2]
    noise_pred = noise_net(noisy_steps, k, context)
    diffusion_loss = mse(noise_pred, noise)

    # Weighted sum; the paper's lambda is 1e-4  [eq. 2] [§IV-C]
    return lam * dist_loss + (1 - lam) * diffusion_loss

# Undirected action sampling with the goal masked  [§IV-B] [eq. 1]
# Input:  obs_images (1, 3P, H, W), num_samples N
# Output: N trajectories (N, H_pred, 2)

def sample_actions(obs_images, num_samples):
    # Goal mask = 1 hides the goal, so the placeholder image is ignored
    context = vision_encoder(obs_images, placeholder, goal_mask=1)
    context = repeat(context, num_samples)
    # Start from Gaussian noise  [§IV-B]
    steps = normal(0, 1, shape=(num_samples, H_pred, 2))
    # Denoise over K steps  [§IV-B] [eq. 1]
    for k in scheduler.timesteps:
        noise_pred = noise_net(steps, k, context)
        steps = scheduler.step(noise_pred, k, steps)
    # Undo the normalization, then integrate displacements into positions
    return cumsum(unnormalize(steps), axis=1)

The second pseudocode block is a reading of the released deployment code, because the paper gives only the schematic update in eq. 1; the code's steps are cited in Paper vs code. Batch loops and the distance-based mask scaling are simplified.

Hyperparameters

Values in the paper are cited to the paper. Values in the repository are cited to the code, and the two columns are given separately in the Notes column where they differ.

Results

Author-reported. Real-world experiments on a LoCoBot platform, indoor and outdoor, in unseen environments fig. 4 §V-A. The paper says six environments in its evaluation overview and in the benchmarking section, but its results text says five §V §V-A. Each baseline was trained on the GNM and SACSoN mix for 20 epochs with minimal tuning §V-A. Number of trials is not stated. Collisions are the mean per experiment §V-A.

Exploration and navigation tab. 1.

Unified versus dedicated policies tab. 2.

Encoder and goal-masking ablation tab. 3.

The CNN variants combine the goal with the observations either before or after the final MLP layers, and the paper does not define the two names beyond that §V-C.

• Late-fusion CNN: 52% success, 3.2 collisions.
• Early-fusion CNN: 68% success, 1.5 collisions.
• ViT with the goal patches blocked by attention masks: 32% success, 2.5 collisions.
• NoMaD (ViNT encoder, attention masking): 98% success, 0.2 collisions.

The paper also reports that the unified model matches the best single-behaviour policies in each mode §V-B tab. 2, and that predicted action distributions stay multimodal in undirected mode fig. 3 fig. 5.

Paper vs code

Code: the official repository robodhruv/visualnav-transformer (MIT, owner account of author Dhruv Shah), at commit dca79815b704e5aa9c6bdc3082351f9e3b2848c2. The arXiv abstract links the project page, whose "Code" link points to a different account (PrieureDeSion/visualnav-transformer). That account is not an author, so it was not used. The same project page's "Data" link points to the authors' repository, which is the one used here (arXiv abstract; project page).

Differences between the paper and the released code. Each entry gives both citations.

• Training data. Paper: GNM and SACSoN datasets, over 100 hours §IV-C. Code: the config lists recon, go_stanford, cory_hall, tartan_drive and sacson, with scand and seattle commented out as private nomad.yaml. The paper does not name the datasets, so the config list is the one a reproduction can follow.
• Epochs. Paper: 30 epochs §IV-C. Config: 100 epochs nomad.yaml.
• Warmup. Paper: warmup with a cosine schedule, length not stated §IV-C. Code: 4 warmup epochs wrapped around a cosine schedule with T_max set to the epoch count nomad.yaml train.py train.py.
• Loss form. Paper: a plain sum of the diffusion MSE and λ times the distance MSE eq. 2. Code: alpha times the distance loss plus (1 − alpha) times the diffusion loss train_utils.py. The relative weight α/(1−α) is about 1e-4, so the paper's λ is kept in effect.
• Loss masking. Paper: no masking in eq. 2. Code: the distance loss is averaged only over goal-conditioned samples [code: train/vint_train/training/train_utils.py#L626-L627], and the diffusion loss is averaged only over samples with valid action masks train_utils.py.
• Noisy input. Paper: the noise predictor sees a^0 + ε^k eq. 2. Code: the noisy input comes from the DDPM add_noise call with schedule coefficients train_utils.py.
• Action parameterization. Paper: the policy outputs future actions a_t, with no statement of coordinates §IV-B. Code: the diffusion model predicts normalized per-step displacements in two dimensions, and inference integrates them with a cumulative sum train_utils.py train_utils.py.
• Normalization constants. Code uses fixed per-dimension min and max values for the displacements data_config.yaml. The paper does not state a normalization.
• Sampler and clipping. Paper: eq. 1 with α, γ and σ left unspecified §IV-B. Code: the diffusers DDPM scheduler with epsilon prediction and clip_sample set to True explore.py train.py.
• Transformer name. Paper: a Transformer decoder with 4 layers and 4 heads §IV-C. Code: nn.TransformerEncoder with self-attention, followed by mean pooling over tokens nomad_vint.py nomad_vint.py.
• Normalization layers. Paper: silent. Code: BatchNorm is replaced by GroupNorm in the vision encoder nomad_vint.py train.py.
• Distance head. Paper: silent on the head. Code: a three-layer MLP that narrows the embedding by a factor of 4 at each layer, from 256 to 1 with ReLU nomad.py.
• Distance units. Paper: temporal distance d eq. 2. Code: the label is the number of waypoint steps between observation and goal, and negative goals get the maximum distance vint_dataset.py.
• Action mask. Paper: silent. Code: negatives and distances outside the action bounds are excluded from the action loss vint_dataset.py nomad.yaml.
• Samples and waypoint. Paper: silent. Code: 8 samples by default explore.py, waypoint index 2 by default navigate.py, and the chosen waypoint is scaled by the speed limit divided by a rate navigate.py.
• EMA. Paper: silent. Code: an exponential moving average of the weights is maintained with power 0.75 and saved as the EMA model train_eval_loop.py train_eval_loop.py.
• Goal masking at test time. Paper: m = 1 for undirected exploration §IV-A. Code: the same, with a random placeholder image that the mask hides explore.py. This matches.
• Context and horizon. Paper: the number of past observations and the horizon are not stated §IV-A §IV-B. Code: context 3 (comment says 5) and horizon 8 nomad.yaml nomad.yaml.

Nothing else was found to differ on the parameters the paper does state, such as batch size, learning rate, optimizer, goal-mask probability, denoising steps, encoder width, head count and layer count.

Pitfalls

• The headline gain is not 25 points. The text says "over 25%" and the table caption says "by 25%" §V-A tab. 1. From Table I, 98% against 77% is 21 points, or about 27% relative. The "15x" fewer parameters claim is the paper's figure; Table I's 335M against 19M is about 17.6x.
• Environment count differs. Six environments in the evaluation overview and benchmarking section §V §V-A, five in the results text §V-A.
• eq. 1 is schematic. The coefficients α, γ and σ are never given §IV-B. Use the DDPM square-cosine schedule and epsilon prediction as in the code train.py.
• Goal masking must hide the token from attention and rescale the pooled output. The code removes the goal token from attention and multiplies the pooled output by a mask that also rescales the unmasked tokens nomad_vint.py nomad_vint.py. A plain mean over all tokens, including the goal, would change the undirected distribution.
• Undirected mode still runs the goal encoder. The deployment code passes random noise as the goal image and sets the mask to 1 [code: deployment/src/explore.py#L117-L118]. Only the mask makes the image irrelevant.
• The λ of 1e-4 makes the distance gradient small. The code also masks the distance loss to goal-conditioned samples train_utils.py. This reading is an inference from the code, not a claim of the paper: check the distance head separately if you retrain.
• Clipping is applied at every reverse step. With clip_sample True, the denoiser's prediction of the clean sample is clipped to the normalized range at each step train.py. The paper does not mention this.
• Normalization constants come from the data. Changing datasets changes the min and max used for scaling data_config.yaml. Retrain or recompute them when the dataset changes.
• Sampling is stochastic. Every control step draws new noise and new samples. The paper gives no per-step latency, so measure latency on your hardware explore.py.
• The U-Net is not in this repository. The conditional 1D U-Net is imported from an external diffusion-policy package train.py, so the paper's "15 convolutional layers" cannot be checked from the repository.
• Unstated values. The number of past frames, the horizon, the image size, the sample count and the waypoint choice are not stated in the paper. Use the config values above and say so in any reproduction.
• Published numbers are author-reported. Trial counts and variance are not given, so differences of a few percentage points between methods are not established by the paper.

Lineage

• Builds on: the ViNT goal-conditioned Transformer policy with EfficientNet tokens, cited in the paper's bibliography as shah2023vint §III; the ViKiNG topological-memory and frontier setup, cited as shah2022viking §III; the diffusion policy and its 1D U-Net, cited as chi2023diffusionpolicy §II §IV-C; and the square-cosine noise schedule, cited as pmlr-v139-nichol21a §IV-C.
• Compared against: VIB (latent goal model for exploration, shah2021rapid); an autoregressive policy over discretized actions; Masked ViNT; Random Subgoals and Subgoal Diffusion, both ViNT systems; the Diffusion Policy and ViNT Policy as single-behaviour baselines §V-A §V-B.
• Superseded by: not stated in this paper.