Matrix-Space Reinforcement Learning for Reusing Local Transition Geometry

TL;DR

This paper introduces Matrix-Space Reinforcement Learning (MSRL), a geometric abstraction that captures local transition dynamics for improved transfer and generalization in sequential decision-making tasks.

Contribution

MSRL provides a novel matrix-based descriptor of trajectory segments that supports algebraic composition, transfer, and smooth value function approximation, enhancing reinforcement learning methods.

Findings

MSRL achieves a target AUC of 0.73, outperforming baseline methods.

The descriptor is well-defined up to coordinate gauge and complete for certain additive signals.

Conditioning value functions on MSRL descriptors enables effective transfer in new tasks.

Abstract

Compositional generalization in sequential decision-making requires identifying which parts of prior rollouts remain useful for new tasks. Existing methods reuse skills or predictive models, but often overlook rich local transition geometry and dynamics. We propose Matrix-Space Reinforcement Learning (MSRL), a geometric abstraction that represents trajectory segments through positive semidefinite matrix descriptors aggregating first- and second-order statistics of lifted one-step transitions. These descriptors expose shared hidden structure, support algebraic composition in an abstract matrix space, and reveal opportunities for transfer. We prove that the descriptor is well defined up to coordinate gauge, complete for the induced low-order additive signal class, additive under valid segment composition, and minimally sufficient among admissible additive descriptors. We further show that conditioning value functions on the trajectory-segment matrix yields a first-order smooth approximation of action values, enabling source-learned matrix-to-value mappings to bootstrap learning in new tasks. MSRL is plug-in compatible with standard model-free and model-based methods, while obstruction filtering rejects implausible compositions. Empirically, MSRL achieves the best average finite-budget target AUC of 0.73, outperforming MSRL from scratch (0.65), TD-MPC-PT+FT (0.63), and TD-MPC (0.57).