🤖 machine learning

Momba: Network Modernization Improves Multi-Objective Reinforcement Learning

This paper introduces Momba, a network modernization approach for multi-objective reinforcement learning that integrates observation normalization, weight normalization, and distributional returns with entropy regularization to significantly improve solution set quality without altering the underlying algorithms.

Original authors: Adam Štafa, Santeri Heiskanen, Petr Novotný, Joni Pajarinen

Published 2026-08-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Adam Štafa, Santeri Heiskanen, Petr Novotný, Joni Pajarinen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers learn to play games or control robots not by being told exactly what to do, but by trying things out and learning from their mistakes. This is the realm of Reinforcement Learning (RL), a branch of artificial intelligence where an "agent" explores an environment to maximize a reward. Think of it like a dog learning tricks: it gets a treat for sitting and a gentle "no" for jumping, eventually figuring out the best way to behave.

But what happens when the dog has to choose between sitting and staying, or running fast and saving energy? In the real world, goals often clash. You can't always be the fastest runner and the most energy-efficient one at the same time. This is the challenge of Multi-Objective Reinforcement Learning (MORL). Instead of finding just one "best" way to do something, MORL tries to find a whole menu of different strategies, each balancing the conflicting goals in a unique way. It's like a chef trying to create a menu where every dish offers a different, perfect balance of spicy and sweet, rather than just one "best" recipe.

For a long time, researchers trying to solve these complex balancing acts focused on inventing new, complicated math recipes (algorithms) to help the computer learn. They assumed the problem was the math itself. However, a new paper titled "Momba: Network Modernization Improves Multi-Objective Reinforcement Learning" suggests that maybe the math wasn't the only thing holding things back. The authors, a team from Masaryk University and Aalto University, wondered: What if the computer's "brain" (its neural network architecture) was just too simple? They decided to test if upgrading the brain's design, using modern tricks from single-goal learning, could make multi-goal learning much better without changing the core math rules.

The Brain Upgrade

The researchers took an existing, solid algorithm called CAPQL (which is already good at finding these balanced strategies) and gave it a serious makeover. They didn't rewrite the rules of the game; instead, they swapped out the player's equipment for a high-tech, modern version. They introduced three specific upgrades to the neural network, the part of the AI that does the thinking:

Normalization (The Equalizer): Imagine trying to learn a new language where some words are shouted in a whisper and others are screamed at you. It's confusing. The first upgrade, observation and feature normalization, acts like a volume knob. It ensures that all the information the AI receives (like speed, energy, or position) is on a similar scale, so the AI doesn't get overwhelmed by the "loud" numbers and ignore the "quiet" ones.
Weight Normalization (The Balanced Diet): In a neural network, connections between neurons have "weights" that determine how important a signal is. Sometimes, these weights can get too big or too small, throwing off the whole system. The second upgrade, weight normalization, keeps these connections in check, ensuring the AI's internal logic stays balanced and doesn't go haywire.
Distributional Critic (The Crystal Ball): This is the star of the show. Traditional AI usually guesses the average reward it will get. But in a world of trade-offs, the average isn't always enough. The distributional critic is like a crystal ball that doesn't just predict a single number; it predicts the entire range of possible outcomes. It understands that a move might have a 50% chance of being great and a 50% chance of being okay, rather than just saying "it's a 7." This helps the AI understand risk and uncertainty much better.

The Results: A Menu of Masterpieces

The team tested their upgraded AI, which they named Momba, on seven different continuous control tasks. These are complex simulations where robots (like a cheetah, a human, or a swimmer) have to move efficiently while balancing multiple goals, such as speed versus energy consumption.

The results were striking. By simply upgrading the architecture, Momba didn't just keep up with the best existing methods; it left them in the dust.

In terms of the quality of the solutions found (measured by a metric called Hypervolume), Momba improved performance by about 132% compared to the original, un-upgraded version of the algorithm.
Even compared to the second-best method in the field, Momba showed a 35% improvement.
Perhaps most impressively, Momba was much more sample-efficient. This means it learned faster, needing far fewer attempts to reach a high level of performance. In some tests, it reached the performance level of a competitor that had trained for 48 million steps in just 1 million steps.

The authors ran these experiments across 10 different random seeds (essentially, 10 different starting conditions) to ensure the results weren't just luck. They found that Momba consistently produced a wider, more diverse, and higher-quality set of solutions. For example, in a simulation of a robot ant, Momba could generate a smooth curve of solutions that covered every possible trade-off between moving fast in the X direction and the Y direction, whereas other methods left big gaps.

What It Means

The paper makes a clear point: you don't always need to invent a brand-new, complex mathematical algorithm to solve a hard problem. Sometimes, you just need to give the existing algorithm a better brain. The authors explicitly argue against the idea that the only way to improve Multi-Objective RL is through complex new update rules or preference selection strategies. Instead, they show that modernizing the neural network architecture is a powerful, often overlooked path to success.

They also tested which of their three upgrades mattered most. Using a statistical tool called Shapley values, they found that observation normalization was the biggest hero, contributing about 55% of the improvement. The distributional critic and weight normalization also played crucial roles, working together to make the whole system sing.

In short, the paper suggests that the future of teaching AI to balance complex, conflicting goals might not lie in writing harder math, but in building smarter, more robust neural networks. By giving the AI a better way to see the world, normalize its inputs, and understand the full range of possibilities, we can get it to learn faster and make better decisions, all without changing the fundamental rules of how it learns.

Technical Summary: Momba – Network Modernization Improves Multi-Objective Reinforcement Learning

Problem Statement

Multi-Objective Reinforcement Learning (MORL) aims to discover a set of policies that balance trade-offs among conflicting objectives, often represented by a Pareto front. While recent advances in single-objective Deep Reinforcement Learning (DRL) have demonstrated that improving neural network architectures (e.g., normalization techniques, distributional critics) can yield substantial gains in sample efficiency and asymptotic performance without altering underlying algorithms, MORL research has predominantly focused on algorithmic innovations.

Current MORL approaches, particularly general policy methods that learn a preference-conditioned policy $\pi(\cdot|s, \omega)$ , typically rely on simple feedforward networks to represent value functions and policies. The paper posits that these architectures may be insufficiently expressive to handle the significant differences in optimal policies and value functions across varying trade-offs. The central question addressed is whether integrating modern neural network design principles into MORL can improve performance without requiring complex preference selection algorithms or specialized update rules.

Methodology

The authors propose Momba, an algorithm that integrates recent architectural advancements from single-objective DRL (specifically the SimbaV2 architecture) into an existing entropy-regularized MORL algorithm, CAPQL. The methodology involves three primary architectural modifications:

Observation and Feature Normalization:
- Observation Normalization: The input observations are normalized using running statistics (Running Statistic Normalization) to zero mean and unit variance.
- Feature Normalization: Hidden features are mapped onto a unit hypersphere (L2 normalization) to stabilize learning dynamics and prevent plasticity loss. A learnable shift constant is concatenated to preserve magnitude information.
Weight Normalization: Network weights are normalized after every gradient update to ensure they remain on the unit hypersphere, improving the conditioning of the optimization landscape.
Distributional Critic Adaptation:
- Standard distributional critics (e.g., C51) model the distribution of returns. In MORL, returns are multivariate vectors.
- Instead of modeling the complex multivariate joint distribution (which is computationally expensive and often limited to tabular cases), Momba proposes learning the scalarized return distribution directly.
- The critic $Z_\theta(s, a, \omega)$ predicts the categorical distribution of the scalarized return $\omega^T G$ , where $G$ is the vector of returns.
- The critic is trained by minimizing the cross-entropy loss between the predicted distribution and the target distribution (computed via TD bootstrap on the scalarized reward $\omega^T r$ ).
- The policy is updated to maximize the expected value of this scalarized distribution plus an entropy term.

Preference Selection: To isolate the effects of architectural changes, the authors employ a simple preference selection strategy: a new preference vector $\omega$ is sampled once per episode from a static uniform distribution over the preference space, rather than using sophisticated dynamic selection strategies found in other MORL baselines.

Key Contributions

Architectural Integration in MORL: The paper demonstrates that employing more expressive neural network architectures (SimbaV2 components) with a distributional critic in a MORL algorithm substantially improves performance on continuous control tasks. This is achieved without altering the core MORL algorithmic framework or requiring complex preference selection mechanisms.
Scalarized Distributional Critic: The authors propose a pragmatic adaptation of the categorical critic to the multi-objective domain. By directly learning the scalarized return distribution under the linear scalarization assumption, they avoid the complexity of modeling multivariate joint distributions while retaining the stability benefits of distributional RL.
Empirical Validation of Design Choices: Through extensive ablation studies, the paper identifies observation normalization as the most significant contributor to performance gains (accounting for ~55% of improvements via Shapley values), followed by feature and weight normalization. The distributional critic is also identified as a central component for performance improvement.

Results

The proposed method was evaluated across 7 continuous control benchmarks (Ant, HalfCheetah, Hopper, Humanoid, Swimmer, Walker2D) with 2 or 3 objectives.

Asymptotic Performance: Momba outperforms state-of-the-art baselines, including PGMORL (a multi-policy method) and CAPQL (the underlying algorithm). On aggregate metrics, Momba achieved approximately 35% improvement in Hypervolume (HV) and 16% improvement in Expected Utility Metric (EUM) over the runner-up (PGMORL). Compared to the base CAPQL, improvements were approximately 132% in HV and 35% in EUM.
Sample Efficiency: Momba reaches the final performance of PGMORL in environments like Ant and Humanoid using only a fraction of the training steps (100k and 200k steps, respectively, compared to PGMORL's 48M and 120M steps).
Comparison with GPI-LS: When compared to GPI-LS, a method known for best-in-class sample efficiency that uses sophisticated preference selection, Momba (using a simple static preference sampling) outperformed GPI-LS in normalized HV and EUM after 200k steps when paired with similar Update-to-Data (UTD) ratios.
Solution Quality: Visual analysis of the generated solution sets (Pareto fronts) shows that Momba produces sets with good coverage and density, often dominating baselines in environments like Walker2D and Hopper.

Significance and Claims

The paper claims that the performance of existing MORL algorithms can be significantly boosted by adopting modern neural network design principles, specifically normalization techniques and distributional critics. The significance lies in the finding that architectural improvements alone, without complex algorithmic modifications or specialized preference selection strategies, can outperform current state-of-the-art MORL methods.

The authors emphasize that the distributional critic is a major component of this success, as it stabilizes learning dynamics and improves the conditioning of the optimization problem. They also note that while their approach relies on linear scalarization, the proposed distributional critic adaptation is a simple yet effective tactic that avoids the need for modeling complex multivariate distributions.

Limitations Acknowledged:

The study is restricted to linear scalarization; nonlinear or learned scalarization functions are not addressed.
Benchmarks are limited to continuous control tasks with 2 or 3 objectives; scalability to discrete spaces or higher-dimensional objectives is not evaluated.
The more expressive function approximators may increase wall-clock training time, though this can be mitigated by system-level advances (e.g., JIT compilation).