Causal dynamics learning has recently emerged as a promising approach to enhancing robustness in reinforcement learning (RL). Typically, the goal is to build a dynamics model that makes predictions based on the causal relationships among the entities. Despite the fact that causal connections often manifest only under certain contexts, existing approaches overlook such fine-grained relationships and lack a detailed understanding of the dynamics. In this work, we propose a novel dynamics model that infers fine-grained causal structures and employs them for prediction, leading to improved robustness in RL. The key idea is to jointly learn the dynamics model with a discrete latent variable that quantizes the state-action space into subgroups. This leads to recognizing meaningful context that displays sparse dependencies, where causal structures are learned for each subgroup throughout the training. Experimental results demonstrate the robustness of our method to unseen states and locally spurious correlations in downstream tasks where fine-grained causal reasoning is crucial. We further illustrate the effectiveness of our subgroup-based approach with quantization in discovering fine-grained causal relationships compared to prior methods.

翻译：因果动力学学习近期成为增强强化学习（RL）鲁棒性的重要途径。其核心目标是构建基于实体间因果关系的动力学预测模型。尽管因果关联往往仅在特定上下文中显现，现有方法却忽视了这类细粒度关系，导致对动力学的理解缺乏深度。本文提出一种新型动力学模型，通过推断细粒度因果结构并将其用于预测，从而提升强化学习的鲁棒性。关键创新在于联合学习带有离散潜变量的动力学模型，该潜变量将状态-动作空间量化为子群，从而识别出具有稀疏依赖性的有意义的上下文。训练过程中，每个子群均能学习到对应的因果结构。实验结果表明，在需要细粒度因果推理的下游任务中，本方法对未见状态和局部虚假关联均展现出鲁棒性。进一步通过量化策略，我们验证了基于子群的方法在发现细粒度因果关系方面相较于现有方法的有效性。