🤖 machine learning

AETDICE: Unified Framework and Offline Optimization for Nonlinear Multi-Objective RL

This paper introduces AETDICE, a unified framework and offline RL algorithm that bridges the divide between Scalarized Expected Return (SER) and Expected Scalarized Return (ESR) paradigms to enable tractable sample-based optimization for nonlinear multi-objective reinforcement learning using static datasets.

Original authors: Woosung Kim, Youngjun Suh, Jinho Lee, Jongmin Lee, Byung-Jun Lee

Published 2026-07-01

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Woosung Kim, Youngjun Suh, Jinho Lee, Jongmin Lee, Byung-Jun Lee

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 you are training a robot to drive a taxi. But this isn't just about getting from Point A to Point B; the robot has to balance multiple, conflicting goals at the same time. Maybe it wants to be fast, but also save energy. Or maybe it needs to serve two different groups of passengers equally, rather than just focusing on the ones who pay the most.

In the world of robotics and AI, this is called Multi-Objective Reinforcement Learning (MORL). The challenge is: how do you tell the robot what "good" looks like when "good" means different things to different people?

The Old Problem: Two Different Rulesets

For a long time, researchers had to choose between two very different ways to teach the robot, and they couldn't mix them:

The "Average" Approach (SER): You tell the robot, "On average, over many trips, I want you to be balanced." This is like a school principal saying, "We don't care if you fail math one day, as long as your average grade across the whole year is good." The robot learns to take risks, sometimes failing a specific trip to get a better average later.
The "Every Single Trip" Approach (ESR): You tell the robot, "Every single trip must be balanced." This is like a strict parent saying, "You must get an A on every test." If the robot fails one math test, the whole strategy is ruined. This is much harder because the robot has to remember exactly what happened in the past (like, "I already served Group A, so now I must serve Group B") to make the right choice for the current moment.

The problem was that existing AI methods could handle the "Average" approach or the "Every Single Trip" approach, but not both at once. And worse, no one had figured out how to teach the robot these complex rules using only a fixed dataset (like a library of past driving logs) without letting the robot drive around and make mistakes in the real world.

The New Solution: AETDICE

The authors of this paper built a new framework called AETDICE. Think of it as a universal translator that can understand any combination of these rules.

Here is how they did it, using a simple analogy:

1. The "Memory Backpack" (Augmented State)

The biggest headache with the "Every Single Trip" rule is that the robot needs to remember its history. If it just looks at the current street corner, it doesn't know if it already served the first passenger group.

The Fix: The authors gave the robot a backpack. Every time the robot takes a step, it writes down its "score so far" in the backpack. Now, when the robot looks at a street corner, it doesn't just see "Corner"; it sees "Corner + Backpack Score."
Why it helps: This turns a confusing, memory-dependent problem into a standard one. The robot can now make decisions based on "Corner + Backpack" just like a normal robot makes decisions based on "Corner."

2. The "Rewriting the Map" (Transformed Rewards)

Usually, you tell a robot, "Go to the store, you get 10 points." But with complex rules (like "balance the scores"), the points aren't fixed; they depend on the backpack.

The Fix: The authors changed the map. Instead of giving points at the end of the trip, they calculated how much each single step contributed to the final goal and gave the robot that tiny slice of points immediately.
Why it helps: This allows the robot to learn from the fixed dataset (the library of past trips) without needing to guess what will happen in the future. It turns a complex, global puzzle into a series of simple, local steps.

3. The "Global Balancer" (DICE Optimization)

Once the robot has its backpack and the new map, there's still a tricky part: ensuring the overall balance across all trips is fair, not just locally.

The Fix: They used a mathematical tool called DICE (Distribution Correction Estimation). Imagine you have a bag of marbles representing all the past trips in your dataset. Some marbles are "good" for your specific goal, and some are "bad." DICE acts like a smart filter that re-weights the marbles, telling the robot: "Ignore the trips where you were too greedy; focus on the trips where you were balanced."
Why it helps: This allows the robot to find the perfect strategy even if the original data was messy or biased.

What Did They Discover?

When they tested this new system, they found some fascinating behaviors that old methods couldn't achieve:

The "Stochastic" Mix: Sometimes, the best strategy isn't to be 100% consistent. The robot learned to be a "coin flipper." In 50% of trips, it would focus entirely on Passenger Group A, and in the other 50%, it would focus entirely on Group B. On average, this was the most balanced outcome. Old methods forced the robot to be either a "specialist" (always A) or a "generalist" (always split), missing this clever middle ground.
The "History-Aware" Driver: For the "Every Single Trip" rule, the robot learned to change its behavior based on its backpack. If it had already served Group A, it would aggressively hunt for Group B, even if the street looked the same. This is a behavior that standard AI usually can't learn from static data.

The Bottom Line

This paper introduces a new way to train AI agents to handle complex, competing goals using only past data. It bridges the gap between "average performance" and "perfect consistency," allowing robots to learn strategies that are fair, balanced, and adaptable to different situations—all without needing to interact with the real world during training.

In short: They built a universal training kit that lets robots learn to juggle multiple goals perfectly, using only a library of old videos, by giving the robot a memory backpack and a smarter way to read the map.

Technical Summary: AETDICE: Unified Framework and Offline Optimization for Nonlinear Multi-Objective RL

1. Problem Statement

Multi-objective reinforcement learning (MORL) aims to optimize sequential decisions under multiple, often competing objectives. While linear scalarization (weighted sums) reduces MORL to standard single-objective RL, nonlinear scalarization is essential for capturing complex preferences such as risk aversion, fairness, and diminishing returns.

Historically, nonlinear MORL has been bifurcated into two distinct paradigms based on the ordering of scalarization and expectation:

Scalarized Expected Return (SER): Applies nonlinearity to the expected return vector ( $u(\mathbb{E}[R])$ ). This requires global-level optimization because the marginal contribution of an action depends on the policy's overall expected return, often necessitating stochastic policies.
Expected Scalarized Return (ESR): Applies nonlinearity to each trajectory's return ( $\mathbb{E}[u(R)]$ ). This creates a non-Markovian problem where the optimal action depends on the accumulated reward history, requiring state augmentation to restore Markovianity.

The Gap: Existing offline RL methods have addressed linear scalarization and SER objectives. However, ESR remains unaddressed in the offline setting. Trajectory-level nonlinearity makes optimal actions dependent on accumulated rewards, and learning such reward-conditioned policies from static datasets has not been explored. Furthermore, no existing method addresses objectives that combine both trajectory-level and expectation-level nonlinearities, nor do they handle the specific trade-offs induced by such combinations.

2. Methodology

The AET Framework

The authors propose the Aggregation–Expectation–Transformation (AET) framework to unify nonlinear MORL objectives. AET decomposes the scalarization process into three components:

Transformation (F): A trajectory-level transformation $F: \mathbb{R}^m \to \mathbb{R}^n$ that maps the return vector to $n$ utility dimensions (e.g., capturing diminishing returns via $\log$ or efficiency via Cobb-Douglas).
Expectation (E): The expectation over trajectories.
Aggregation (G): An expectation-level aggregation $G: \mathbb{R}^n \to \mathbb{R}$ (typically concave) that combines the expected utilities.

The AET objective is defined as:
$J_{AET}(\pi; F, G) = G(\mathbb{E}_{\tau \sim \pi}[F(R(\tau))])$

This framework subsumes Linear MORL, SER-MORL, and ESR-MORL as special cases while enabling a richer class of objectives (AET-MORL) that combines nonlinear $F$ and concave $G$ .

AETDICE Algorithm

To optimize AET objectives offline, the authors propose AETDICE (Offline AET optimization via DIstribution Correction Estimation). The method addresses the two sources of nonlinearity separately:

1. Handling Nonlinear Transformation $F$ (Augmented MOMDP):
To handle the trajectory-level nonlinearity of $F$ , the authors construct an augmented finite-horizon MOMDP $\tilde{M} = (\tilde{S}, A, \tilde{P}, H, \tilde{p}_0, \tilde{r})$ .

State Augmentation: The state is augmented with the accumulated reward vector: $\tilde{s}_t = (s_t, R^{acc}_t, t)$ . This restores Markovianity for ESR objectives.
Transformed Reward: A new reward function is defined to capture the incremental contribution of each action to the trajectory utility $F$ :
$\tilde{r}(\tilde{s}_t, a_t) = F(R^{acc}_t + r(s_t, a_t)) - F(R^{acc}_t)$
This creates a telescoping sum where $\sum \tilde{r} = F(R(\tau))$ . This transformation reduces ESR-MORL to standard single-objective RL in the augmented space and AET-MORL to SER-MORL in the augmented space.

2. Handling Nonlinear Aggregation $G$ (Convex Optimization via DICE):
With $F$ absorbed into the reward, the problem becomes maximizing a concave function $G$ over expected utilities in the augmented space. The authors extend the Distribution Correction Estimation (DICE) framework:

Visitation Distribution: The objective is reformulated as a convex optimization over the visitation distribution $d_t(\tilde{s}, a)$ in the augmented space.
Slack Variable: A slack variable $k$ is introduced to decouple the expectation from the concave aggregation $G$ , allowing for unbiased sample-based estimation.
Regularization: A $\phi$ -divergence penalty is added to correct for distribution shift between the learned policy and the offline dataset.
Dual Formulation: The primal problem is converted to a dual problem solvable via stochastic gradient descent on neural networks representing the Lagrangian multipliers ( $\nu$ ) and the slack variable ( $\mu$ ). The optimal policy is extracted via weighted behavior cloning using the learned visitation ratios.

3. Key Contributions

Unified Framework (AET): The paper introduces AET, a principled framework that unifies SER and ESR and defines a new class of objectives (AET-MORL) combining trajectory-level and expectation-level nonlinearities.
First Offline ESR Optimization: By defining a transformed reward in an augmented state space, the authors provide the first method to optimize ESR objectives from static datasets, reducing the problem to standard offline RL.
AETDICE Algorithm: The authors propose AETDICE, a sample-based offline algorithm that handles both nonlinear $F$ (via augmentation) and nonlinear $G$ (via DICE-style distribution correction).
Empirical Validation: The framework is validated across diverse objectives, including Fair MORL (interpolating between ESR and SER via $\alpha$ -fairness), convex utilities, and non-convex Cobb-Douglas utilities with safety constraints.

4. Experimental Results

Experiments were conducted on Fair-Taxi, MO-PointMaze, and D4MORL (MO-Ant, MO-Walker2d, MO-HalfCheetah) environments.

Fair MORL (Fair-Taxi & PointMaze): AETDICE successfully optimized ESR, SER, and Balanced Scalarized Return (BSR) objectives.
- Behavioral Differences: ESR-optimal policies produced balanced returns within every trajectory (non-Markovian behavior dependent on history). SER-optimal policies produced stochastic mixtures of specialized trajectories (Markovian, balancing only in expectation). BSR policies smoothly interpolated between these behaviors.
- Performance: AETDICE achieved high scores on their respective target objectives, outperforming baselines like FairDICE (limited to SER) and ESR-IQL (limited to ESR).
General AET Objectives (D4MORL):
- Convex $F$ : The method successfully optimized objectives prioritizing high-performing dimensions (convex $F$ ) while maintaining balance via concave $G$ , resulting in stochastic mixtures of specialized behaviors.
- Non-Convex $F$ (Cobb-Douglas): AETDICE optimized efficiency trade-offs (e.g., speed vs. energy). Crucially, by adding a safety utility via concave aggregation, the method prevented agent collapse (a failure mode of pure efficiency optimization) while maintaining high efficiency.
Ablation: The study on the regularization coefficient $\beta$ showed that smaller values are required for nonlinear MORL to permit the necessary distribution shift, yet the method remains stable and effective.

5. Significance and Claims

The paper claims that AETDICE resolves long-standing barriers in offline MORL by unifying previously fragmented optimization strategies.

Unification: It bridges the divide between SER and ESR, providing a single framework for general nonlinear MORL.
Novelty: It is the first method to address offline ESR optimization and the broader class of AET objectives.
Behavioral Insight: The framework reveals distinct policy behaviors (non-Markovian strategies, stochastic mixtures) arising from the interaction of trajectory-level and expectation-level nonlinearities, which existing methods fail to capture or address.
Practicality: By enabling offline optimization of complex fairness and efficiency objectives, the method supports real-world decision-making where policies must balance competing criteria both within and across episodes.

The authors note limitations regarding the scalability of the augmented state space as the number of objectives increases and the stringent data coverage requirements inherent to offline MORL with distribution shift. However, they assert that AETDICE significantly advances the field by making a broader class of nonlinear objectives tractable in the offline setting.