Dynamics-Aligned Shared Hypernetworks for Contextual RL under Discontinuous Shifts

TL;DR

This paper introduces DMA*-SH, a dynamics-aligned shared hypernetwork framework for zero-shot generalization in contextual reinforcement learning, especially under discontinuous shifts in latent contexts.

Contribution

The paper proposes a novel hypernetwork-based method trained via dynamics prediction that effectively handles discontinuous context shifts in RL environments.

Findings

DMA*-SH achieves zero-shot generalization on the Actuator Inversion Benchmark.

Outperforms domain randomization by 58.1% on AIB tasks.

Surpasses standard context-aware baselines by 11.5% on average.

Abstract

Zero-shot generalization in contextual reinforcement learning remains a core challenge, particularly when the context is latent and must be inferred from data. A canonical failure mode arises when latent context discontinuously changes how actions affect the environment, requiring incompatible control responses across contexts. We propose DMA*-SH, a framework where a single hypernetwork, trained solely via dynamics prediction, generates a small set of adapter weights shared across the dynamics model, policy, and action-value function. This shared modulation imparts an inductive bias matched to discontinuous context-to-dynamics shifts, while input/output normalization and random input masking stabilize context inference, promoting directionally concentrated representations. We provide theoretical support via expressivity separation results for hypernetwork modulation, and a variance decomposition with policy-gradient variance bounds that formalize how within-mode compression improves learning under non-overlapping contexts. For evaluation, we introduce the Actuator Inversion Benchmark (AIB), a suite of environments designed to isolate challenging context-to-dynamics interactions, including actuator inversion, actuator permutations, and weakly non-overlapping continuous dynamics. On AIB's held-out tasks, DMA*-SH achieves zero-shot generalization, outperforming domain randomization by 58.1% and surpassing a standard context-aware baseline by 11.5% on average.