TDR: Task Distribution Relevance in RL

• TDR is defined as the Jensen gap between per-task maximization and averaged Q-values, capturing optimal action disagreement across tasks.
• It reveals that even subtle reward-function differences can lead to significant performance gaps when policies cannot identify the active reward model.
• Empirical estimation via finite task samples allows practical assessment of task difficulty, justifying the need for reward-informed policies in RL.

Searching arXiv for the specified paper and a closely related topic for context. Task Distribution Relevance (TDR) is a state-dependent quantity introduced to characterize task heterogeneity in reinforcement learning settings where tasks share similar dynamics but differ in reward functions. In “Task Aware Dreamer for Task Generalization in Reinforcement Learning,” TDR serves as the quantitative object that links cross-task disagreement in optimal control to the failure modes of policies that cannot identify which reward function is currently active (Ying et al., 2023). Within that framework, TDR is used to justify reward-informed policies and to explain why general world models alone are insufficient when tasks differ significantly in their optimal action-values.

1. Formal definition within task-generalization reinforcement learning

The paper considers a distribution over Markov decision processes, denoted by $\mathcal{T}$, where each task is an MDP $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$. For each task $\mathcal{M}$, let $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$ be its optimal action-value function. For any state $s\in\mathcal{S}$, Task Distribution Relevance is defined as (Ying et al., 2023)

$$D_{\mathrm{TDR}}(\mathcal{T},s) = \mathbb{E}_{\mathcal{M}\sim\mathcal{T}} \bigl[\max_{a}Q^*_{\mathcal{M}}(s,a)\bigr] - \max_{a}\, \mathbb{E}_{\mathcal{M}\sim\mathcal{T}} \bigl[Q^*_{\mathcal{M}}(s,a)\bigr].$$

This definition compares two operations. The first term selects the best action separately for each task and then averages the resulting optimal values across tasks. The second term averages the task-specific $Q^*$ functions first and then selects a single action that is optimal only with respect to that average. The difference between these two quantities is the TDR at state $s$.

The construction is explicitly local in state space. TDR is not defined as a single scalar property of a task family unless one subsequently aggregates it across states. In the paper, that locality is important because policies experience task ambiguity only in the subset of states visited under their discounted occupancy measures.

2. Intuition: Jensen gap, task disagreement, and action inconsistency

The paper interprets TDR as a Jensen gap between “choose the best action per task, then average” and “average $Q$-values, then choose one action” (Ying et al., 2023). This interpretation makes the metric directly about disagreement among tasks in optimal action choice.

If all tasks in $\mathcal{T}$ share exactly the same preferred action at $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$0, then the maximizer $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$1 is identical across tasks. In that case,

$$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$2

so $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$3. The zero value therefore corresponds to a state at which the task family is decision-theoretically aligned.

If different tasks prefer different actions at $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$4, the first term remains large because each task is permitted to use its own maximizing action, whereas the second term is smaller because it forces a single action to be selected after averaging. The gap therefore grows with disagreement among the tasks’ optimal $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$5-functions.

A common misconception is to treat TDR as a generic statistical distance between tasks. The definition is narrower and more operational: it quantifies disagreement in optimal control at a given state, not arbitrary discrepancy in rewards, transitions, or observations. This suggests that TDR is best understood as a control-relevant notion of task difference rather than a general-purpose divergence.

3. Theoretical properties and the performance-gap bound

The paper establishes several basic properties of TDR (Ying et al., 2023). First, nonnegativity follows from Jensen’s inequality: $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$6 Second, the zero-gap condition holds iff all tasks in $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$7 agree on at least one common maximizer at $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$8. This characterizes precisely when a single state-conditioned action can be optimal for every task in the distribution.

The most consequential result is the performance-gap bound. For any Markovian policy $$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{\mathcal{M}},\gamma)$$9, if $\mathcal{M}$0 is its discounted-state occupancy in task $\mathcal{M}$1, then Theorem 3 gives

$\mathcal{M}$2

This bound identifies high-TDR states as the states at which policies that cannot distinguish tasks lose return. The lower bound is formulated in terms of state occupancies, so the impact of TDR is not merely a property of the task distribution in abstraction; it is mediated by where the policy actually spends time. A plausible implication is that two task families with similar average TDR could induce different empirical difficulty if their high-TDR states lie in regions of state space with very different occupancies.

4. Practical estimation and empirical aggregation

The paper also specifies how TDR can be computed from a finite sample of tasks and evaluated states (Ying et al., 2023). Given sampled tasks $\mathcal{M}$3 and evaluated states $\mathcal{M}$4, the procedure is:

For each task $\mathcal{M}$5, estimate or approximate its optimal $\mathcal{M}$6 for all $\mathcal{M}$7. In small MDPs this may be done with dynamic programming; in deep-RL settings one may train a critic network to convergence.

Then compute

$\mathcal{M}$8

From these,

$\mathcal{M}$9

and the empirical TDR is

$$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$0

The paper notes that one can aggregate across states by averaging or by taking a weighted sum under a state-occupancy distribution. That aggregation choice matters conceptually. A uniform average emphasizes geometric properties of the state space, whereas occupancy weighting emphasizes the effective decision surface encountered by a policy. This suggests that empirical TDR is not a single canonical scalar unless the state-weighting convention is fixed.

5. Examples and empirical interpretation across task pairs

The paper illustrates TDR with several task pairs whose reward structures induce different levels of action agreement (Ying et al., 2023). Walker-stand and Walker-walk both reward “height above 1.2,” so they share the same optimal action and have $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$1. Acrobot-swingup and swingup_sparse involve the same intended motion with dense versus sparse reward, and likewise have $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$2. Reacher-easy and Reacher-hard both reward “reach within radius $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$3,” but with different radii, yielding a small nonzero TDR. Walker-walk and Walker-prostrate reward nearly opposite behaviors—being tall versus being very low—so their optimal $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$4 functions are nearly opposite and TDR is very large.

These examples clarify that reward-function differences do not automatically imply large TDR. Dense-to-sparse modifications of the same behavioral objective can preserve the optimal action ranking and therefore induce near-zero TDR, whereas reward changes that invert action preferences at the same state induce large TDR. A common misunderstanding is to identify task similarity with superficial reward-shape similarity; the paper’s examples instead define similarity through the alignment of optimal action-values.

In the experiments, Markovian or history-only policies, denoted $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$5 and $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$6, perform well when $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$7 but degrade catastrophically as TDR grows, whereas Task Aware Dreamer, denoted $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$8, remains robust. The reported pattern is therefore not simply that harder tasks have larger TDR, but that higher TDR specifically degrades policies lacking reward-informed task identification.

6. Role in Task Aware Dreamer and reward-informed policies

Task Aware Dreamer extends world models to the task-generalization setting by integrating reward-informed features that identify consistent latent characteristics across tasks (Ying et al., 2023). Within that method, the world model optimizes the variational lower bound of sample-data log-likelihood together with a new term designed to differentiate tasks using their states. TDR is introduced to demonstrate the advantages of the reward-informed policy in this framework.

The central argument is that high TDR means tasks differ significantly in their optimal $$Q^*_{\mathcal{M}}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$$9. A policy that chooses actions based only on current state, or on state-action history, cannot infer which reward function is in effect, so it must average over tasks. The performance-gap theorem then formalizes the associated loss. By contrast, a policy that includes reward histories as context can identify the current task and select the corresponding optimal action.

TDR therefore functions as more than a descriptive metric. It is the paper’s formal bridge between task distribution structure and architectural necessity. In low-TDR regimes, a general world model and a Markovian decision rule may be adequate because the tasks already agree on what action should be taken. In high-TDR regimes, reward-informed context becomes necessary because the ambiguity is irreducibly task-conditional. This suggests that TDR can be read as a criterion for when task inference must enter the policy class, rather than remaining implicit in the world model alone.

7. Scope, interpretation, and conceptual limits

Within the paper, TDR is defined using the optimal action-value functions $s\in\mathcal{S}$0, so it is fundamentally an oracle quantity (Ying et al., 2023). In practice it must be estimated through approximate critics or dynamic programming, which makes empirical TDR dependent on the quality of $s\in\mathcal{S}$1 approximation. The paper addresses this operationally by giving a finite-sample estimation recipe, but the definition itself remains idealized.

TDR is also state-specific and policy-relevant rather than universally task-global. High task relevance at one state does not imply high relevance everywhere, and a task family may contain both benign and highly ambiguous regions. This is important for interpretation: a low aggregated TDR may conceal localized states at which Markovian policies incur substantial regret if those states are heavily occupied.

Finally, TDR should not be conflated with partial observability in the usual dynamics-centered sense. In the formulation of the paper, the critical ambiguity arises because the reward mechanism varies across tasks while dynamics may be shared. The resulting difficulty is not that the environment state is hidden, but that the active task cannot be recovered from state alone when different tasks assign different optimal actions to the same state. TDR quantifies that specific obstruction and thereby explains why reward-informed policies are essential whenever tasks exhibit significant task distribution relevance.