Dynamic Multi-Strategy Reward Distillation

• The paper introduces a dual reward decomposition method that distills a robust shared task reward while capturing strategy-specific nuances from heterogeneous demonstrations.
• It employs adversarial IRL with dynamic strategy creation to update policies continuously in lifelong and federated learning settings.
• Empirical results demonstrate significant gains, including up to 100% strategy precision and a 77% improvement in policy returns on simulated and real-world tasks.

Dynamic Multi-Strategy Reward Distillation (DMSRD) is a framework for inverse reinforcement learning (IRL) that enables the extraction of a robust, shared task reward and simultaneous discovery or adaptation of multiple diverse strategies from heterogeneous demonstrations. DMSRD addresses challenges inherent to learning from demonstration (LfD), specifically reward ambiguity and heterogeneity among demonstrators, and extends previous multi-strategy reward distillation approaches to dynamic, lifelong, and federated settings (Chen et al., 2020, Jayanthi et al., 2022).

1. Problem Formulation and Motivation

In the standard IRL paradigm, the goal is to infer a reward function that explains observed behavior, typically under the assumption that all demonstrations optimize a single latent objective. However, in LfD scenarios, demonstrations for a single task often display significant heterogeneity, reflecting different user strategies and preferences. Naïve aggregation of such data produces reward ambiguity, where a single reward cannot faithfully represent all styles and may collapse to trivial or incoherent solutions (Jayanthi et al., 2022). Prior approaches either disregard heterogeneity or require explicit clustering and retraining, failing to exploit shared structure across strategies.

DMSRD aims to (1) distill a robust task reward common to all strategies, resolving ambiguity; (2) learn strategy-specific reward components explaining demonstrator preference heterogeneity; and (3) update these models dynamically as new data arrives, allowing for online personalization, federated learning, and scalability to large, evolving datasets (Chen et al., 2020, Jayanthi et al., 2022).

2. Model Architecture and Reward Decomposition

The core modeling assumption in DMSRD is that each strategy's reward can be decomposed into two parts:

$$R^{(i)}(s,a) = R^{(0)}(s,a) + \alpha_i\,\widetilde R^{(i)}(s,a)$$

where $R^{(0)}$ is a shared task reward, $\widetilde R^{(i)}$ is a strategy-only residual, and $\alpha_i > 0$ balances strategy uniqueness versus shared objective (Chen et al., 2020, Jayanthi et al., 2022).

Network architectures are multi-layer perceptrons (MLPs) for both the shared task reward and each strategy residual, with non-shared weights. Policies are maintained per strategy and trained under their respective combined rewards. This “two-column” distillation, inspired by Distral-style decomposition, enables concurrent learning of both global and idiosyncratic objectives.

3. Learning Objective and Algorithmic Procedure

DMSRD unifies adversarial IRL (via AIRL discriminators) with reward distillation. For each demonstration-strategy pair, the following loss is optimized:

$$L_D = \sum_{i=1}^N \left[\, \mathbb{E}_{(s,a)\sim \tau_{\rm exp}^{(i)}}[-\log D_{\theta_i, \theta_0}(s,a)] + \mathbb{E}_{(s,a)\sim \tau_{\rm gen}^{(i)}}[-\log(1-D_{\theta_i, \theta_0}(s,a))] - \alpha_i\,\mathbb{E}_{(s,a) \sim \pi^{(i)}}\left[\|\widetilde R^{(i)}_{\theta_i}(s,a)\|^2\right] \,\right]$$

The discriminator $D_{\theta_i,\theta_0}(s,a)$ combines the current strategy and task reward, regularizing residuals to zero unless necessary. These updates are alternated with policy improvement steps (using TRPO or PPO). The approach naturally distills a robust task reward and isolates strategy-specific nuances (Chen et al., 2020, Jayanthi et al., 2022).

In the fully dynamic DMSRD framework (Jayanthi et al., 2022), demonstrations arrive sequentially. Each new demonstration either (A) is explained by a convex mixture of existing strategy policies or (B) triggers creation of a new strategy if the KL divergence between the demonstration’s state-distribution and the closest mixture exceeds a threshold. The assignment step uses random search over strategy weights, minimizing empirical KL divergence, and, if no satisfactory mixture is found, creates and trains a new strategy module via AIRL. Strategy and task modules are updated by the composite AIRL+MSRD loss plus “between-class discrimination” constraints to enforce consistent mixture semantics.

4. Strategy Discovery, Policy Mixtures, and Scalability

A key innovation of DMSRD is its approach to strategy discovery and compositional policy generation (Jayanthi et al., 2022). New demonstrations are first evaluated for compatibility with mixtures of existing strategies, with policies expressed as Gaussian-action mixtures $\pi_w(\cdot|s) = \mathcal{N}(\mu_w(s), \sigma^2)$, where $\mu_w(s) = \sum_j w_j \cdot \mu_{\pi_{\phi_j}}(s)$.

When a new “style” cannot be well-explained by existing policies (KL divergence exceeds tolerance), a new base strategy is instantiated and absorbed into the strategy set. Between-class discrimination (BCD) losses are used to ensure that strategies remain distinct and mixture weights have well-interpretable behavioral meaning. This dynamic regime ensures efficiency and prevents proliferation of redundant strategy modules, supporting lifelong scaling and minimizing model complexity as the demonstration corpus grows.

Federated variants of DMSRD enable each user or agent to maintain private strategy policies and residuals, sharing only updates to the global task reward network and assignment indices with a central server. This design supports decentralized, privacy-aware learning without catastrophic forgetting, as new or evolved strategies are continually incorporated (Jayanthi et al., 2022).

5. Experimental Evidence and Key Results

DMSRD demonstrates empirical superiority over prior approaches across simulated and real domains:

• In simulated control tasks such as the inverted pendulum and Hopper, DMSRD achieves higher Pearson correlation with ground-truth task rewards ($r=0.998$ for Pendulum, $r=0.943$ for Hopper) compared to AIRL baselines ($r=0.51$, $r=0.89$) (Chen et al., 2020).
• In physical robot table tennis, the method recovers a task reward that reliably distinguishes successful from failed returns; strategy-only rewards assign highest scores to trajectories of matching strokes, validating disentanglement (Chen et al., 2020).
• In lifelong control scenarios (e.g., lunar lander, Pendulum), DMSRD improves policy returns by 77% over AIRL/MSRD, increases task-reward correlation (Pendulum: $r=0.953$ vs. $r=0.74$), and achieves perfect (100%) precision in assigning demonstrations to correct strategy modules (Jayanthi et al., 2022).

A summary of selected metrics:

Task/Metric	DMSRD Result	Baseline Result (AIRL/MSRD)
Pendulum: $r$ corr	$0.998$	$0.51$
Hopper: $r$ corr	$0.943$	$0.89$
Table Tennis Retrn	$83\%$ (topspin: $90\%$)	--
Lifelong Policy Retn	$+77\%$	--
Task Reward Corr	$0.953$ (Pend), $0.614$ (Lander)	$0.74/0.50$ (Pend/Lander)
Strategy Precision	$100\%$ (Pend), $100\%$ (Lander)	$57\%$, $83\%$

DMSRD also demonstrated scalable strategy discovery, distilling 5–9 base strategies from 100 demonstrations without performance degradation (Jayanthi et al., 2022).

6. Hyperparameters, Regularization, and Theoretical Considerations

Key hyperparameters include the strategy deviation weight $\alpha$ (best results for $\alpha\approx 0.1$), network depth (2–3 layers, 32–64 units), and mixture assignment KL threshold $\epsilon$ (usually $0.05$) (Chen et al., 2020, Jayanthi et al., 2022). Excessively small $\alpha$ causes strategy networks to collapse into the task reward, while excessively large $\alpha$ impedes distillation, leading to over-fragmented models.

Optimization is via Adam for reward modules and PPO/TRPO for policies, with AIRL loss regularized by the $L_2$ norm of residuals. Training steps per demonstration are typically 50,000 frames for new AIRL strategy creation; 20 MSRD updates per demonstration ensure continual adaptation.

While a formal proof of convergence is not provided, the method inherits local convergence properties of GAN-style min-max optimization and trust-region policy algorithms, such as monotonic improvement given accurate reward estimation (Chen et al., 2020). The two-column reparameterization used in MSRD/DMSRD accelerates convergence relative to vanilla distillation by enabling direct gradient flow from discriminator losses to the task reward module.

7. Strengths, Limitations, and Future Directions

Strengths:

• Resolves reward ambiguity by leveraging multiple strategies for robust task distillation.
• Simultaneously models both shared (global) and personalized (residual) objectives.
• Supports dynamic, federated, and lifelong learning settings, with effective scaling and privacy preservation.
• Demonstrated in both simulated and real-world, high-dimensional robot control domains with superior empirical performance (Chen et al., 2020, Jayanthi et al., 2022).

Limitations:

• Requires labeled or identifiable strategies for pure strategies; entirely unsupervised discovery is not addressed.
• Decomposition is constrained to the linear form $R^{(i)} = R^{(0)} + \alpha_i \widetilde R^{(i)}$, precluding more complex reward interactions.
• Min-max adversarial optimization may require careful tuning to ensure stability (Chen et al., 2020, Jayanthi et al., 2022).

Potential Extensions:

• Integration of unsupervised structure discovery (e.g., Dirichlet-process priors for strategies).
• Non-linear or multiplicative reward decompositions to capture richer dependencies.
• Theoretical analysis of joint min-max and policy optimization loops.
• Application to multi-task and continual learning scenarios.

This suggests DMSRD provides a modular, scalable, and robust solution to reward ambiguity and strategy heterogeneity in IRL and LfD. Its dynamic and federated mechanisms enable efficient and privacy-compliant adaptation across evolving, diverse user populations (Chen et al., 2020, Jayanthi et al., 2022).