LocoMamba: Vision-Driven Locomotion via End-to-End Deep Reinforcement Learning with Mamba

In plain words

Legged robots can walk over rough ground more reliably when they can see the ground ahead, not only feel their own joints. LocoMamba gives a simulated four-legged robot a single controller that takes two inputs: the robot's own body and joint sensor readings, and a depth image from a camera. The controller is trained with trial-and-error reinforcement learning, so the robot learns its own walking behaviour instead of following hand-written rules.

The part that gives the paper its name is the sequence-processing core. It uses Mamba, a recent design that processes a long sequence at a cost that grows roughly in step with its length. Attention, used by Transformer models, grows much faster with length. In the authors' simulated tests, this controller travelled farther and hit fewer obstacles than a Transformer-based version trained the same way. Those results are the authors' own, and they come from simulation only.

TL;DR

• Policy for a 12-joint quadruped, trained end to end with PPO on simulated terrain with obstacles §3.3 §4.1.
• Proprioception becomes one token through an MLP. Depth images become spatial tokens through a CNN. Both feed two stacked Mamba (selective state-space) layers §3.1 §3.2 tab. 1.
• Author-reported: higher return, fewer collisions and longer distance than a Transformer fusion baseline on the training terrain; the same ordering holds zero-shot on unseen terrain tab. 4 tab. 9 tab. 10.
• Only a simulator (PyBullet) was used; real-robot tests were not done §6.
• Official code at the cited commit contains only a README and licence, with no model or training code README.md.

Problem

Blind controllers that use only proprioception react only after contact, so they struggle to avoid obstacles or place feet on irregular ground §1. Adding vision fixes this but has a cost: Transformer-based fusion scales quadratically with the number of tokens, which limits sequence length, spatial resolution and training efficiency §1 §2.3. The paper assumes that a recurrent state-space backbone can keep long-range context at lower cost, and that PPO with domain randomisation and a curriculum can train the whole pipeline end to end §1.

Method

The policy takes the current proprioceptive vector and the four most recent depth frames, encodes each modality into tokens, fuses them with stacked Mamba layers, pools the result into one feature and feeds policy and value heads trained by PPO §3 fig. 1.

Step 1: observation §3.1 §4.2. The observation is the proprioceptive vector together with four depth frames eq. 1. The paper says temporal aggregation over these frames is handled by the backbone §3.1; see Pitfalls.

Step 2: tokenisation §3.1 eq. 2–8. An MLP maps the proprioceptive vector to one token eq. 2. A compact CNN turns the depth image into N spatial tokens eq. 3–4. Both are projected to a shared width d eq. 5–6 and concatenated, with the proprioceptive token first eq. 7. Learned spatial index and modality embeddings are added, followed by LayerNorm eq. 8.

Step 3: Mamba fusion §3.2 eq. 9–11 fig. 2. Each layer scans the token sequence. For each token it computes input-dependent parameters and updates a hidden state eq. 9–10. The state is carried from the end of one time step to the start of the next §3.2. Layers are stacked with residual connections and LayerNorm eq. 11. The paper uses two layers tab. 1.

Step 4: pooling and heads §3.2 eq. 12. The proprioceptive output token is concatenated with the mean of the visual output tokens, then passed through a 2-layer MLP to give the feature h_t eq. 12 tab. 1. Separate MLP heads produce the Gaussian policy mean and log-standard-deviation and the value estimate §3.3.

Step 5: action §3.3 §4.2 eq. 14–15. The policy samples a 12-dimensional unbounded vector and squashes it with tanh to actuator limits. Actions are joint position targets, which a PD controller converts into torques §4.2.

Step 6: PPO §3.3 eq. 16–20 alg. 1. Training uses the standard clipped PPO objective with GAE advantages, a value regression term and an entropy bonus eq. 16–20. Each iteration collects a fixed number of samples, then runs several epochs of shuffled minibatch updates §4.4 alg. 1.

Step 7: curriculum and randomisation §4.4 tab. 3 alg. 1. Physics parameters are sampled per episode tab. 3. Depth frames receive salt-like pixel saturation §4.4. Obstacle density rises linearly from sparse to dense over training §4.4.

Reward §4.2 eq. 21–24. The reward is a weighted sum of forward velocity, an actuation penalty and a per-step alive bonus eq. 21.

Key equations:

\tilde{a}_t \sim \mathcal{N}\big(\mu_\theta(h_t),\ \mathrm{diag}(\sigma_\theta^2(h_t))\big), \qquad a_t = a_{\max}\tanh(\tilde{a}_t) eq. 14 eq. 15

\mathcal{J}(\theta,\phi) = -\mathcal{L}_{\text{clip}}(\theta) + \beta_V\,\mathcal{L}_V(\phi) - \beta_H\,\mathbb{E}[\mathcal{H}_t] eq. 20

R_t = 1.0\,R^{\text{fwd}}_t - 0.005\,\|\tau_t\|_2^2 + 0.1\,R^{\text{alive}}_t eq. 21 eq. 22 eq. 23 eq. 24

# Cross-modal Mamba feature  [§3.1] [§3.2] [eq. 2-12] [fig. 1]
# Input:  prop (93,); depth (4, 64, 64)  [§4.2]
# Output: feature h (d_h,); updated recurrent state

def fused_feature(prop, depth, state):
    # Proprioceptive token from a 2-layer MLP  [eq. 2] [tab. 1]
    z_prop = mlp_prop(prop)
    # Depth to N spatial tokens; CNN and patch size not stated  [eq. 3-4]
    patches = cnn_patchify(depth)                      # (N, d_v)
    # Project both modalities to width d  [eq. 5] [eq. 6]
    token_prop = linear_prop(z_prop)                   # (d,)
    token_vis = linear_vis(patches)                    # (N, d)
    # Proprioceptive token first, then visual tokens  [eq. 7]
    tokens = concat(token_prop, token_vis)             # (1 + N, d)
    # Add spatial and modality embeddings, then LayerNorm  [eq. 8]
    tokens = layer_norm(tokens + pos_embed + mod_embed)
    # Two stacked Mamba layers with residual and LayerNorm  [§3.2] [eq. 11]
    for layer in mamba_layers:
        out, state = selective_scan(layer, tokens, state)
        tokens = layer_norm(out + tokens)
    # Mean-pool the visual tokens, keep the proprio token  [eq. 12]
    pooled = concat(tokens[0], mean(tokens[1:], axis=0))
    return head_mlp(pooled), state

# One selective state-space layer, sequential form  [§3.2] [eq. 9-10]
# Input:  tokens (1 + N, d); state (h,) carried from the previous step
# Output: outputs (1 + N, d); final state

def selective_scan(layer, tokens, state):
    outputs = []
    for token in tokens:
        # A, B, C, D come from affine maps of this token; form not stated
        a_bar, b_bar, c_bar, d_bar = layer.gates(token)
        # Hidden-state update  [eq. 9]
        state = a_bar * state + b_bar * token
        # Output read-out  [eq. 10]
        outputs.append(c_bar * state + d_bar * token)
    return stack(outputs), state

# Generalized advantage estimation  [eq. 18]
# Input:  rewards (T,); values (T + 1,)  (bootstrap value at the end)
# Output: advantages (T,)

def gae(rewards, values, gamma, lam):
    advantages = zeros(T)
    running = 0.0
    for t in reversed(range(T)):
        # One-step TD error  [eq. 18]
        delta = rewards[t] + gamma * values[t + 1] - values[t]
        running = delta + gamma * lam * running
        advantages[t] = running
    return advantages

# PPO loss for one minibatch  [eq. 16] [eq. 17] [eq. 19] [eq. 20]
# Inputs: batch with features h, actions a, old log-probs, advantages,
#         returns; coefficients eps, beta_v, beta_h (see Hyperparameters)

def ppo_loss(batch, policy, value, eps, beta_v, beta_h):
    # Likelihood ratio against the behaviour policy  [eq. 16]
    ratio = exp(policy.logp(batch.a, batch.h) - batch.logp_old)
    # Clipped surrogate, to be maximised  [eq. 17]
    surrogate = minimum(
        ratio * batch.adv,
        clip(ratio, 1 - eps, 1 + eps) * batch.adv,
    )
    # Value regression to the empirical return  [eq. 19]
    value_loss = (value(batch.h) - batch.returns) ** 2
    # Entropy of the Gaussian policy  [eq. 19]
    entropy = policy.entropy(batch.h)
    # Total loss to minimise  [eq. 20]
    return (
        -mean(surrogate)
        + beta_v * mean(value_loss)
        - beta_h * mean(entropy)
    )

Algorithm 1 alg. 1 is the full loop: sample a curriculum density and textures per iteration, collect episodes with randomised physics and depth noise until the iteration has 16,384 samples, compute GAE, normalise the advantages, run 3 epochs of minibatch updates, then advance the curriculum.

Hyperparameters

Results

Author-reported. Setup: PyBullet, PPO with the same schedule for all methods, training on the Thin Obstacle terrain, zero-shot tests on Rugged and Dynamic terrain, 10 runs each §4 §5.1 §5.5. Baseline is Transformer Proprio-Vision §5.1.

Ablation on visual token count (LocoMamba, Thin Obstacle tab. 7): 32 tokens return 740.55, collisions 57.88, distance 32.93; 64 tokens (the default) 762.34, 72.53, 32.41; 256 tokens 797.46, 106.95, 34.74. Memory grows from 200 MB (32 tokens) to 359.6 MB (256 tokens) tab. 8. On rugged terrain, 256 tokens collide far more often (4166.64 vs 470.27 for 64) tab. 7.

Other headline ratios the paper reports, all derived from the rows above: vs Transformer Proprio-Vision on the training terrain, 48.9% more return, 48.9% fewer collisions and 30.4% more distance §5.2.

Paper vs code

The official repository is allen-quad-robot/locomamba, linked from the abstract. At code_commit it contains only README.md, LICENSE and media/ README.md LICENSE. There is no model, environment, training or configuration code, so no architecture, PPO setting or reward term can be compared with code. The README lists the paper and a follow-up paper with a Mamba2 backbone; the follow-up's method is not covered by this card README.md. Repository licence is MIT LICENSE.

Two discrepancies come from the paper itself, not from code. First, the introduction and the abstract say safety is enforced by penalties for collisions and falls §1. The reward equations contain no collision or fall penalty; falls and unrecoverable collisions only end the episode, which stops the +0.1 per-step alive bonus eq. 21 eq. 24. Second, the README BibTeX lists the authors as "Yinuo Wang and Xiaowen Tao" README.md; the arXiv metadata used for this card lists "Gavin Tao" as the second author.

Pitfalls

• Recurrent state is underspecified. The paper carries the SSM state from the last token of one time step to the first token of the next §3.2. It does not say whether the state resets at episode start or at the 999-step horizon, whether gradients flow through the carried state during PPO updates, or how the state is recomputed for shuffled minibatches §4.4 alg. 1. Choose and document an option.
• Depth-frame handling disagrees across sections. Section 3.1 says the four depth frames are aggregated by the backbone §3.1. Section 4.2 stacks them as one observation §4.2. The paper does not say whether each frame becomes its own set of tokens or whether frames are stacked as channels before the CNN.
• Token count is not derived in the text. N = 64 is labelled "Ours" in the token ablation tab. 7, but the main text never states N, the patch size P or d_v eq. 4. Whether N counts one depth frame or all four is not stated.
• Safety penalties are not in the reward. See Paper vs code. Implement exactly what eq. 21–24 say, and check the intro's claim against the authors' evaluation before relying on it.
• Alive term at termination. R^alive is 1 "until termination" eq. 24. Whether the terminal step itself gets the bonus is not stated.
• Bootstrapping and truncation. GAE sums to the end of the rollout without a termination mask, and the bootstrap value at the end of a 999-step horizon is not specified eq. 18 alg. 1.
• Entropy and value coefficients. Table 2 gives an entropy coefficient of 0.005. The name is mapped to β_H here by inference. β_V is not stated, and gradient clipping thresholds are not stated eq. 20 §4.4.
• Action scale and PD gains. The action bound a_max is not given eq. 15. Nominal PD gains are not given; Table 3 only gives randomisation ranges.
• Curriculum is not numeric. The density endpoints, the length of the ramp and the "empirically determined" schedule are not stated §4.4.
• Variance is large. Several baselines have standard deviations larger than their means, for example Transformer Proprio-Vision on rugged terrain (257.38 ± 431.52) tab. 9. The paper reports no significance test, and its rugged-terrain gap sits inside those spreads. Treat the zero-shot gains as suggestive.
• Efficiency claims are narrower than the text. Section 5.3 calls the FLOPs "comparable" §5.3, but the policy FLOPs per step differ by about 18% (4.11M vs 5.01M) tab. 6. LocoMamba reserves slightly more memory (3160 vs 3156 MB) tab. 6. Throughput was measured on a 4 GB laptop GPU, and the throughput gap is about 6%.
• Naming drift. The same models appear as "Mamba Proprio-Vision" and "Transformer Vision-State" in the text §5.2 §5.3. Use the table names.
• Mamba vision-only. The paper attributes this variant's weak results to recurrent-state drift under partial observability §5.3. That is the authors' explanation, not a tested result.
• Real-world transfer is untested. No hardware experiments were done; latency and safety on a real robot are listed as future work §6.

Lineage

• Builds on: Mamba, a selective state-space model with input-dependent parameters and a hardware-aware parallel scan, as cited §2.3 [gu2023mamba]; S4, as cited §2.3 [gu2021efficiently]; PPO §3.3 [schulman2017proximal].
• Compared against: Proprio-Only and vision-only MLP fusion variants §5.1; Transformer Proprio-Vision and Transformer Vision-Only, following the cross-modal Transformer policy of [yang2021learning] §5.1.
• Related but not card-backed: the same authors' follow-up with an SSD-Mamba2 backbone is named in the repository README README.md; no card.
• Superseded by: not known from wiki content.