Technical Summary: Matrix-Space Reinforcement Learning for Reusing Local Transition Geometry

Problem Statement

Sequential decision-making requires agents to identify which components of prior rollouts remain useful for new tasks to achieve compositional generalization. While existing reinforcement learning (RL) methods successfully reuse skills, predictive models, or compressed states, they often overlook rich local transition geometry and dynamics. Current approaches struggle to determine how to leverage trajectories with different appearances (e.g., varying horizon lengths, task realizations, or action representations) to improve value and policy learning without inducing infeasible behaviors. The core challenge is quantifying local transition dynamics in a way that exposes shared hidden structures and supports algebraic composition for transfer across related tasks.

Methodology: Matrix-Space Reinforcement Learning (MSRL)

The paper proposes Matrix-Space Reinforcement Learning (MSRL), a geometric abstraction that represents trajectory segments through positive semidefinite (PSD) matrix descriptors.

1. The Matrix Proxy $M(\tau)$

The fundamental unit of MSRL is the trajectory-segment matrix $M(\tau)$. For a segment $\tau = ((o_t, a_t, r_t, o_{t+1}))_{t=0}^{T-1}$, the method defines a fixed-step "lift" $\psi_t(\tau)$ that aggregates:

• The embedded transition midpoint ($c_t$) and directed displacement ($\delta_t$).
• Action-conditioned midpoints and displacements (via Kronecker products with action features).
• The observed reward ($r_t$) and a constant channel.

The matrix proxy is defined as the sum of outer products of these lifted vectors:
$$M(\tau) = \sum_{t=0}^{T-1} \psi_t(\tau)\psi_t(\tau)^\top \in \mathbb{S}^m_+$$
This construction results in a covariance-like matrix that aggregates second-order statistics of lifted one-step transitions. Crucially, because $M(\tau)$ is a commutative sum, it discards the full chronological sequence (which varies across tasks) while retaining low-order additive statistics essential for value learning.

2. Algebraic Composition and Gluing

The method establishes an algebra on these descriptors:

• Composition: Valid concatenation of segments $\tau = \tau^{(1)} \star \tau^{(2)}$ corresponds exactly to matrix addition: $M(\tau) = M(\tau^{(1)}) + M(\tau^{(2)})$.
• Gluing: Overlapping local restrictions can be assembled into a global descriptor using inclusion-exclusion operators, allowing replay libraries to store and recombine partial segments.

3. Obstruction Filtering

To ensure that composed matrices correspond to realizable trajectories, the method introduces an obstruction score. This metric measures the distance of a candidate matrix from the set of reachable matrices in a specific environment. Positive obstruction serves as a certificate of inconsistency, filtering out implausible compositions before they are used for planning.

4. Value Learning and Transfer

MSRL conditions value functions on the trajectory-segment matrix. The authors prove that under specific assumptions (lifted Bellman compatibility and local smoothness), the value function admits a first-order smooth approximation in the matrix space. Specifically, for small compositional perturbations $H_0, H_1$:
$$\bar{Q}^\star(o, H_0 + H_1, a) \approx \bar{Q}^\star(o, H_0, a) + \bar{Q}^\star(o, H_1, a) - \bar{Q}^\star(o, 0, a)$$
This "centered compositional value law" enables source-to-target transfer: a matrix-to-value branch ($M \to V$) is pretrained on source environments and used to initialize the critic in target tasks. The target policy and action heads are trained from scratch on target data, while the matrix-value branch provides a structural inductive bias.

Key Contributions

1. Geometric Descriptor: Introduction of the trajectory-segment matrix $M(\tau)$ as a sufficient, low-order descriptor that captures local geometry and dynamics while enabling algebraic composition.
2. Foundational Properties: Proofs that $M(\tau)$ is:
   • Well-defined up to coordinate gauge (invariant under aligned coordinate changes).
   • Complete for the induced low-order additive signal class (capturing all relevant first-order and quadratic statistics).
   • Additive under valid segment composition.
   • Minimally sufficient among admissible additive descriptors.
   • Separable from environment realizability via obstruction scores.
3. MSRL Theory: Development of a matrix-conditioned value law showing that value functions are locally smooth and approximately additive in the matrix space, justifying the use of source-trained $M \to V$ mappings to bootstrap learning in new tasks.
4. Empirical Validation: Demonstration that MSRL achieves superior performance in few-shot settings compared to standard RL and transfer baselines.

Experimental Results

The method was evaluated on a suite of simple source environments (GraphMotif, MicroGrid, PointRooms, Dubins-Reacher) and complex target control tasks (AntMaze, Reacher/Pusher, Walker2d, Hopper, Humanoid).

• Source Diagnostics: The matrix descriptor successfully recovered interpretable structures, including directed motifs, coordinate-aligned matching, and action-conditioned transport.
• Transfer Performance: In the complex target suite, MSRL initialized with a source-pretrained matrix-value branch achieved a normalized few-shot AUC of 0.73. This outperformed:
  • MSRL trained from scratch (0.65).
  • TD-MPC-PT+FT (0.63).
  • Standard TD-MPC (0.57).
• Learning Speed: The source-initialized approach demonstrated faster early learning convergence compared to baselines, indicating that the transfer of matrix-value knowledge effectively reduces the sample complexity of the target task.

Significance and Limitations

The paper argues that MSRL provides a transparent algebraic interface for compositional generalization, transforming trajectory reuse into matrix operations. By abstracting local transition geometry into PSD matrices, the method enables agents to identify and recombine structural components across tasks without relying on exact state matching or full chronological history.

The authors maintain a modest stance on the method's universality. They explicitly note that:

• The approach requires an aligned transition lift across source and target settings; it does not automatically align arbitrary environments.
• The descriptor preserves only low-order additive affine-quadratic information and does not encode arbitrary chronological order beyond what is visible in the lifted vectors.
• The approximate additive value rule is a local modeling rationale dependent on smoothness assumptions, not an exact dynamic programming identity for arbitrary matrix sums.
• The obstruction certificate is exact only for the true reachable set; in practice, it relies on a learned surrogate filter.

Despite these limitations, the work establishes a rigorous geometric framework for reusing local transition structures, offering a plug-in compatible solution for both model-free and model-based RL methods to improve compositional generalization.