Paper deep dive
Hereditary Geometric Meta-RL: Nonlocal Generalization via Task Symmetries
Paul Nitschke, Shahriar Talebi
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 94%
Last extracted: 7/20/2026, 5:59:47 AM
Summary
The paper introduces Hereditary Geometric Meta-RL, a framework that leverages Lie group symmetries to enable non-local generalization in reinforcement learning. By treating the task space as having 'hereditary geometry' induced by system symmetries, the agent can reuse policies learned from training tasks by transforming states and actions via Lie group actions. This approach converts Meta-RL into a symmetry discovery problem, improving sample efficiency and numerical stability compared to standard smoothness-based methods, as demonstrated on 2D navigation tasks.
Entities (10)
Relation Signals (6)
Paul Nitschke → affiliatedwith → Harvard University
confidence 99% · P. Nitschke is with the Paulson School Of Engineering And Applied Sciences, Harvard University
Shahriar Talebi → affiliatedwith → UCLA
confidence 99% · S. Talebi is with the UCLA Samueli School of Engineering and Applied Science
Hereditary Geometric Meta-RL → uses → Lie Group
confidence 95% · the agent reuses a policy learned at the train time by transforming states and actions through actions of a Lie group.
SO(2,R) → exampleof → Lie Group
confidence 92% · A hereditary geometry G in the 2-D navigation Example 1 is given by G = SO(2,R)
Hereditary Geometric Meta-RL → appliedto → 2-D Navigation Task
confidence 90% · Empirically, on a two-dimensional navigation task, our method efficiently recovers the ground-truth symmetry
Hereditary Geometric Meta-RL → employs → Differential Symmetry Discovery
confidence 88% · To learn these structures, we develop a differential symmetry discovery method.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Meta-Reinforcement Learning (Meta-RL) commonly generalizes via smoothness in the task encoding. While this enables local generalization around each training task, it requires dense coverage of the task space and leaves richer task space structure untapped. In response, we develop a geometric perspective that endows the task space with a "hereditary geometry" induced by the inherent symmetries of the underlying system. Concretely, the agent reuses a policy learned at the train time by transforming states and actions through actions of a Lie group. This converts Meta-RL into symmetry discovery rather than smooth extrapolation, enabling the agent to generalize to wider regions of the task space. We show that when the task space is inherited from the symmetries of the underlying system, the task space embeds into a subgroup of those symmetries whose actions are linearizable, connected, and compact--properties that enable efficient learning and inference at the test time. To learn these structures, we develop a differential symmetry discovery method. This collapses functional invariance constraints and thereby improves numerical stability and sample efficiency over functional approaches. Empirically, on a two-dimensional navigation task, our method efficiently recovers the ground-truth symmetry and generalizes across the entire task space, while a common baseline generalizes only near training tasks.
Tags
Links
- Source: https://arxiv.org/abs/2603.00396v1
- Canonical: https://arxiv.org/abs/2603.00396v1
Trouble viewing inline? Open PDF directly →
Full Text
43,182 characters extracted from source content.
Expand or collapse full text
Hereditary Geometric Meta-RL: Nonlocal Generalization via Task Symmetries Paul Nitschke 1 Shahriar Talebi 2 Abstract—Meta-Reinforcement Learning (Meta-RL) com- monly generalizes via smoothness in the task encoding. While this enables local generalization around each training task, it requires dense coverage of the task space and leaves richer task space structure untapped. In response, we develop a geometric perspective that endows the task space with a “hereditary geometry” induced by the inherent symmetries of the underlying system. Concretely, the agent reuses a policy learned at the train time by transforming states and actions through actions of a Lie group. This converts Meta-RL into symmetry discovery rather than smooth extrapolation, enabling the agent to generalize to wider regions of the task space. We show that when the task space is inherited from the symmetries of the underlying system, the task space embeds into a subgroup of those symmetries whose actions are linearizable, connected, and compact—properties that enable efficient learning and inference at the test time. To learn these structures, we develop a differential symmetry discovery method. This collapses functional invariance constraints and thereby improves numerical stability and sample efficiency over functional approaches. Empirically, on a two-dimensional nav- igation task, our method efficiently recovers the ground-truth symmetry and generalizes across the entire task space, while a common baseline generalizes only near training tasks. Index Terms—meta-reinforcement learning, symmetry, Lie groups, invariance, geometric learning, geometric task embed- dings, policy transfer, kernel methods, navigation tasks I. INTRODUCTION Reinforcement Learning (RL) continues to face fundamental challenges, particularly in terms of generalization and sample efficiency [1]. Meta-RL addresses this by training the agent on multiple tasks sampled from a task space M to generalize to “similar” unseen tasks [2]. A prevalent approach is memory- based Meta-RL [3] which posits a smooth-manifold hypothesis on M: Jointly learn a task encoder ψ : M → Z ⊆ R d and a joint policy for all tasks conditioned on z ∈ Z . While plausible, memory based agents typically generalize only locally around the training tasks [4] – thereby requiring a dense coverage of M with training tasks to generalize successfully. We argue that memory-based methods generalize locally because they generalize via smoothness in Z . The encoder ψ commonly maximizes the mutual information (MI) between task trajectories and embeddings [5], either via contrastive learning (CL) [6] or variational auto-encoders [7]. The MI ob- jective is insensitive to the global geometry ofZ . Specifically, the unique global optima of CL are equidistant embeddings, 1 P. Nitschke is with the Paulson School Of Engineering And Applied Sciences, Harvard University, Cambridge, US. (paul.nitschke@outlook.de) 2 S. Talebi is with the UCLA Samueli School of Engineering and Applied Science, University of California, Los Angeles, US. (s.talebi@ucla.edu) where all training-task encodings lie equally far apart [8]. This construction preserves only tangential geometry and thereby yields faithful encodings for precisely those tasks that are infinitesimally close to the training set, while potentially neglecting richer structure present in M. In practice, departures from CL’s idealized regime—e.g., smaller embedding dimen- sion or fewer negatives samples—can induce limited non- tangential structure [9], [10], [11]. However, these effects are challenging to tune and remain fundamentally local. Likewise, standard actor–critic backbones such as PPO and SAC [12], [13] for policy optimization exhibit primarily local generaliza- tion around training tasks [1], also leaving broader task-space transfer unresolved. This motivates our central question: Can we endow M with a structure richer than the smooth manifold hypothesis that enables generalization beyond local smoothness? In this work, we introduce a framework that captures symmetry-induced geometry in the task representation Z . Inspired by biological agents and their case-based reasoning approach to generalization [14], we propose retrieve & reuse to generalize non-locally: Given a test task, retrieve a sim- ilar training task and roll out its policy after transforming states and actions via left actions of a Lie group G–for an introduction to Lie Groups we refer to [15] and also lectures [16]. During training, the agent discovers G, while the specific task inference g ∈ G is performed at test time. This enables efficient knowledge transfer across wider regions of the task space M. Our main contributions are as follows: • After formalizing the problem setup (§I,) we propose our central symmetry hypothesis for Meta-RL, called a hereditary geometry (§I.) • We identify a salient regime in which the geometry of Z is inherited from the symmetries of the system. This endows M with a hereditary geometry and subsumes many settings where tasks are intuitively perceived as “similar” (§IV.) • We formulate hereditary geometry discovery as a con- crete learning problem, estimating inherent symmetries from trajectory data and performing inference over group actions (§V.) Finally, we empirically validate our approach on a 2-D naviga- tion benchmark in §VI and compare against existing methods. We provide concluding remarks in §VII. arXiv:2603.00396v1 [cs.LG] 28 Feb 2026 I. PROBLEM FORMULATION Herein, we first define the standard Meta-RL problem setup adapted from [17]. A task space M is a set of elements M ∈ M, each representing a Markov decision process M = S,A,R M ,T M ,γ. All tasks share the same state and action spaces S and A and discount rate γ ∈ (0, 1), but differ in their reward R M : S × A × S → [0, 1] and transition functions T M : S × A → ∆ S , where ∆ S denotes the probability simplex on S. Then, the Meta-RL problem proceeds in two stages: At the meta-train time, the agent is given black-box access to N train + 1 uniformly sampled tasks M 0 ,...,M N train ∼ U (M) to train a policy π which aims to minimize the mean incurred regret R at the meta-test time arg min π∈Π E H X h=K E τ :h ∼(π(τ :h−1 ),M) [R (π(τ :h ) ;M)] , (1) where the outer expectation is over [M ∼ U (M)], H ∈ N denotes the number of episodes, τ :h all trajectories collected up to episode h. The policy π(τ :h ) = π(τ :h ;M) can, in general, be non-Markovian, that is depend on all past trajectories τ :h−1 , and generally depends on the latent task M. The value function associated with the policy π(τ :h ;M) in the task M is V π(τ :h ;M) M (s 0 ) : = E " ∞ X t=0 γ t R M (s t ,a t ,s t+1 ) s 0 = s 0 # , (2) where the expectation is over [a t ∼ π(a t | s t ,τ :h ;M),s t+1 ∼ T M (s t+1 | s t ,a t )], and V ∗ M (s 0 ) denotes the value function of a policy that maximizes (2), both evaluated at the initial state s 0 ∈ S. Then, the regret associated with the policy π in the taskM is defined asR(π;M) : = V ∗ M (s 0 )−V π M (s 0 ). The set Π represents a family of admissible policies, usually a dense neural network, and the “shot parameter” K ∈ N controls the number of free exploration episodes. As π(τ :h ;M) generally depends on M, a successful meta- agent must both identify the new test task and quickly general- ize its knowledge from the training tasks to roll-out an optimal policy in its belief of the test task. While these two objectives may generally interfere, this work focuses on the latter; that is, the regime in which the generalization performance dominates. We therefore set the shot parameter sufficiently high to allow one to obtain a relatively reliable estimate of the task before incurring any regret. A common formalization of Meta-RL is memory-based Meta-RL [18], [19], [3] which endows M with a manifold hypothesis: For every M ∈ M, there exists an unknown task encoding z ∈ R d such that for all s, s ′ ∈ S, a∈ A R M (s,a) = R(s,a;z) = : R z (s,a), T M (s ′ |s,a) = T (s ′ |s,a;z) = : T z (s ′ |s,a). We denote the set of all task encodings by Z ∋ z. Given black-box access to training tasks M 0 ,...,M N train ∈ M, memory-based approaches jointly train a task encoder ψ and a global policy π a | s;z . The policy conditions on the s 0 z 0 z 1 z 2 Fig. 1. Illustration of the 2-D navigation task. After learning to navigate from the origin s 0 to the goal positions z 0 and z 1 , the agent aims to generalize their knowledge to navigate to the unseen location z 2 at the test time. current state s and the task embedding z of the current task M. Example 1 (2-D Navigation) A popular Meta-RL example is 2-D navigation [3], [20]: Starting from the origin s 0 = 0, the agent aims to navigate to different unknown goal positions z ∈ S 1 : = Z on the unit circle (cf. Figure 1) where the reward and transition function are given by R(s,a;z) =−||s− z|| 2 2 , T (s ′ | s,a;z) = s + a. Given a step-size α > 0 (which we set to 1 in the following for notational simplicity and to α = 0.1 in the sample implementation), the optimal policy of a task z is taking a step towards the goal location z: π ∗ (a| s;z) = α(z− s). Assumption 1 We assume that reward R(s,a;z), transition T (s ′ |s,a;z), and optimal policy π(a|s;z) are Lipschitz con- tinuous in s,a and z where the Lipschitz continuities in R and T are respectively quantified by the L 2 norm and the Wasserstein distance while task-distances are measured in Z and denoted by d(M,M ′ ). This is a standard assumption in the Meta-RL literature [21] which reads that similar tasks (measured in Z ) have similar optimal policies, see [21] for further discussion. Exploiting the Lipschitz continuity of the policy in z, an agent faithfully encoding M can then generalize their policy to test tasks that are “Lipschitz close” to a training task in Z . More precisely, for every training taskM i , i∈0,...,N train , there exists some constant ε i > 0 such that the agent generalizes to M ∈ M if d(M,M i ) < ε i for some training task M i . While possible, such an approach is highly inefficient: It requires the training tasks to densely cover M and, as a result, discards potentially global structures of the task space M by replacing it with a purely local approximation. In response, we pose the following problem: Train an agent to expand its generalization beyond local smoothness to non- local test tasks, that is, test tasks that are not necessarily close to the training set in terms of their Lipschitz continuity captured by each ε i . Problem statement. Consider a black-box access to training tasks M 0 ,...,M N train ∈ M that do not (necessarily) densely cover M ∃M∈ M s.t. d(M,M i ) > ε i ∀i∈1,...,N train , where ε i > 0 represents the prior generalization constant. Learn a policy π that generalizes non-locally, and uniformly within M R(π;M) < ε∀M∈ M, for a uniform constant ε > 0 comparable with ε i (up to the Lipschitz continuity factor of π in z.) While these requirements are harder to satisfy, they promote utilization of richer structures in M beyond Lipschitz continu- ity, which current memory-based settings do not induce. In this work we endow M with a geometric hypothesis that enables such non-local generalization. Then, our goals are threefold: (i) To formalize this geometric hypothesis; (i) To show why we expect common Meta-RL applications to enjoy such geometric structure; and (i) To learn such a policy purely from trajectory data samples. I. HEREDITARY GEOMETRIES Biological agents commonly generalize by first retrieving a similar, known situation and then reusing it, also called case-based-reasoning [14]. For example, an ice skater can rollerblade by first recalling their ice-skating experience, then reusing it by applying the same movements but on wheels and asphalt rather than blades and ice. We formalize this intuition by positing that the optimal policy π ∗ (a| s;z) in the test task coincides with an optimal policy of some training task M 0 after transforming S and A via left actions L g : S → S and K g : A→ A of a Lie group G: ∀z ∈Z ∃ g ∈ G s.t. π ∗ (a| s;z) = K −1 g (π ∗ (a| L g · s;z 0 )), ∀s∈ S,a∈ A. (3) Here, the action K g on the distribution π ∗ over A is defined via inversion (see Equation 12 in [22] for a similar definition regarding equivariant policies) and we callM 0 the “base task.” For instance, we have π 0 = π ice skating , L g asphalt ground→ icy ground and K g (a) = a in the ice skater example. Assuming a group structure in the transformations L and K allows the agents to efficiently generalize at test time by inferring G, L g and K g at train-time and only inferring the respective group element g ∈ G at test-time. Equation (3) converts the conventional meta-RL optimiza- tion problem into a supervised symmetry discovery problem. To this end, we make a standard assumption from symmetry discovery [23], [24] that G’s left actions can be linearized: Definition 1 (Linear left actions) The left actions L g and K g are called linear if there exist diffeomorphisms φ : S → ̃ S and η : A→ ̃ A such that for all g ∈ G φ◦ L g ◦ φ −1 : = ̃ L g ∈ GL + (|S|,R), η◦ K g ◦ η −1 : = ̃ K g ∈ GL + (|A|,R), where GL + (d,R) denotes the general linear group of d- dimensional, real-valued matrices with strictly positive deter- minant. Then, we also call φ and η representations of L g and K g . Linearization allows a simple parametrization of L g and K g in the learning problem and always holds true locally under regularity assumptions by the rank theorem (see Theorem 4.12 in [15]). Merging the two above conditions yields our geometric hypothesis: Definition 2 (Hereditary Geometry) We call the geometry of M hereditary if there exists a task encodingZ and a Lie group G with linear left actions L g and K g that fulfill (3). A property of an object is called hereditary if it is inherited by all of its subobjects. Conjugating the left-action in Equation (3) reads that the geometry in M is inherited by any task in M, implying that the geometry in M is hereditary–hence the name. Next, we show that the hereditary geometry in π ∗ naturally traces back to R and T . This model based approach lays the ground for our learning problem in §V. Lemma 1 Assume there exists a Lie group G with linear left actions L g : S → S and K g : A→ A such that ∀ z ∈Z ∃ g ∈ G s.t. ( R(s,a;z) = R(L g · s,K g · a;z 0 ) T (s ′ | s,a;z) = T (L g · s ′ | L g · s,K g · a;z 0 ) , (4) for all s,s ′ ∈ S,a∈ A. Then, the geometry in M is hereditary. Proof.We first show that the hereditary geometry in R and T translates into a hereditary geometry in the optimal Q-function Q ∗ (s,a;z) which is the unique fixed point of the Bellman operator B : R |S|×|A|×|Z| → R |S|×|A|×|Z| (BQ)(s,a;z) : = E s ′ R(s,a;z) + γ max a ′ ∈A Q(s ′ ,a ′ ;z) , for s ∈ S,a ∈ A and z ∈ Z and the expectation is taken over [s ′ ∼ T (s ′ |s,a;z)]. Now, recall value iteration (VI) which recursively defines Q t+1 (s,a,z) : = (BQ t )(s,a,z) and Q 0 (s,a,z) = 0 for all s ∈ S,a ∈ A,z ∈ Z and converges to the optimal Q-function: lim t→∞ Q t (s,a,z) = Q ∗ (s,a,z) [25]. Then, we show via induction over t ∈ N that the claim holds for every iterate of VI. Taking the limit in t on both sides concludes the claim for Q ∗ . Finally, the lemma follows by defining π ∗ (a|s;z) as the greedy policy.□ Example 2 A hereditary geometry G in the 2-D navigation Example 1 is given by G = SO(2,R) L g = g −1 · s K g = g −1 · a, where SO(2,R) is the two-dimensional special-orthogonal group which acts on states and actions via standard matrix multiplication. Proof.Define z 0 = (1, 0) T as the base encoding. As Z = S 1 , each element z ∈Z can be written as z = B·z 0 for some B ∈ SO(2,R). Then, we have for the reward function R R(s,a;z) = (s− B· z 0 ) T (s− B· z 0 ) = (s− B· z 0 ) T B T (s− B· z 0 ) = B T · s− B T B· z 0 T B T · s− B T B· z 0 = (L g · s− z 0 ) T (L g · s− z 0 ) = R(L g · s,K g · a;z 0 ) as B T = B −1 for B ∈ SO(2,R) and similarly for T .□ IV. HEREDITARY GEOMETRIES VIA SYMMETRIES Observe the high-level structure in the preceding example: After endowingZ with the geometry SO(2,R), we embedded SO(2,R) into S and A while leaving R and T invariant, that is, we embed into the symmetries of the base task. We now formalize this intuition and subsequently show in Theorem 1 that the geometry in M is then hereditary. This provides one general, concrete setting where M has a hereditary geometry, namely if the geometry of the task space arises from the symmetries of the system. We first define the geometry in the task space and the symmetries of a task, starting with the former. Assumption 2 (Geometric task space) There exists a compact and connected Lie group H acting on a base task encoding z 0 ∈Z with a left action J h :Z →Z that spans Z [ h∈H J h · z 0 =Z. Combining the memory-based setting with Assumption 2 implies that the entire task space M collapses to the tuple (M 0 ,H) which we call a geometric meta-MDP [ g∈G (S,A,R J g ·z 0 ,T J g ·z 0 ,γ) = [ z∈Z (S,A,R z ,T z ,γ) = [ M∈M (S,A,R M ,T M ,γ) = M. We now define the symmetries of a single task. Meta-RL ap- plications commonly enjoy rich, high-dimensional symmetries as they often originate in physical systems, such as robotics, that naturally exhibit strong symmetries. 1 A task M is called symmetric with respect to a compact and connected Lie group N if N leaves M invariant [22]: Definition 3 (Symmetric MDP) Let N act on S and A via linear left actions L n : S → S and K n : A → A. Then, the tuple (M,N ) is called a symmetric MDP if R and T are invariant with respect to N for all s,s ′ ∈ S,a∈ A R(L n · s,K n · a;z) = R(s,a;z),(5) T (L n · s ′ | L n · s,K n · a;z) = T (s ′ | s,a;z).(6) 1 Or as Philip Anderson puts it: It is only slightly overstating the case to say that physics is the study of symmetry [26]. Example 3 (Symmetry in 2-D navigation) Asymmetryin each navigation task from Example 1 is given by N = SO(2,R), φ(s) = s− z, η(a) = a, L g · s = g· s,K g · a = g· a, as the two-norm in the reward function is rotation invariant and the transition function is linear. In a symmetric MDP, the N -invariance of R and T trans- lates into an equivariance in M’s optimal policies [22]: Lemma 2 (Symmetric MDPs admit equivariant optimal poli- cies) Let (M,N ) denote a symmetric MDP. Then, its optimal policy π ∗ fulfills an equivariance property with respect to L n and K n : For all n∈ N,s∈ S,a∈ A π ∗ (a| L n · s;z) = K −1 n · π ∗ (a| s;z),(7) i.e., the pushforward of the measure π ∗ (.| s;z) by K n . We now formalize the embedding of H into N from the beginning of §IV which, in turn, enables endowing M with a hereditary geometry in Theorem 1. To make the embedding independent of the chosen base task, we require that all tasks share compatible symmetries: Definition 4 (Compatible Symmetry) We say that N is com- patible with the symmetries of a geometric Meta-MDP (M 0 ,H) if for any other M ∈ M induced by h ∈ H , N is a symmetry with representation (φ h ,η h ) which changes equivariantly in N ; that is, N induces linear left actions J n : S → S and ̃ J n : ̃ S → ̃ S such that 1) denoting the representation of N in M 0 by (φ 0 ,η 0 ), ∃n∈ N : φ h (s) = ̃ J n ◦ φ 0 ◦ J −1 n (s) ∀s∈ S,(8) 2) and ̃ J −1 n ◦ ̃ L n ◦ ̃ J n combines to a valid left action: ∀n∈ N ∃ n ′ ∈ N : ̃ J −1 n ◦ ̃ L n ◦ ̃ J n = ̃ L n ′ .(9) Alternatively, the above reads that 1) all tasks have “similar” symmetries and 2) their symmetries are related through com- patible representations. The 2-D navigation task indeed has compatible symmetry. Example 4 (Compatible Symmetries in 2-D navigation) Let z = B · z 0 , B ∈ SO(2,R) denote a task encoding from Example 1. Example 3 showed that a symmetry embedding φ h is given by φ h (s) = s− z. Define left actions J n : = B· s and ̃ J n : = B· ̃s. Then, (8) holds ̃ J n ◦ φ 0 ◦ J −1 n · s = B· (B −1 · s− z 0 ) = s− z = φ h (s). Further, (9) holds as the product of rotation matrices is a rotation matrix. Finally, we show that Compatible Symmetry gives rise to a class of systems with hereditary geometry. This is formalized in the following result: Theorem 1 (Hereditary Geometry from Symmetry) Let (M 0 ,H) be a geometric Meta-MDP that induces a task space M whose symmetries are compatible with H —in the sense of Def. 4. Let M ∈ M denote a task induced by h ∈ H and assume that there exists n∈ N such that π ∗ (a| s;J h · z 0 ) = π ∗ (a| J −1 n φ −1 h ̃ L n φ h (s);z 0 ).(10) Then, the geometry in M is hereditary. Proof.We have that: π ∗ (a| s;J h · z 0 ) = π ∗ (a| J −1 n φ −1 h ̃ L n φ h (s);z 0 ) = π ∗ (a| J −1 n J n φ −1 0 ̃ J −1 n ̃ L n ̃ J n φ 0 J −1 n · s;z 0 ) = π ∗ (a| φ −1 0 ̃ L n ′ φ 0 J −1 n · s;z 0 ) = η −1 0 ̃ K n ′ η 0 π ∗ (a| J −1 n · s;z 0 ),(11) where the equalities respectively follow from (10), (8), (9) and (7). Finally, define the Lie group G : = N . Then, G induces left actions L g : S → S and K g : A→ A given by L g : = J −1 n and K g : = η −1 0 ̃ K n ′ η 0 , where n ′ is the normal left action from (9). Plugging L g and K g into (11) and observing that L g and K g are linear left actions completes the proof.□ Finally, the condition in (10) reads that the change in the optimal policy for a new task (induced by J h ) can be expressed with a symmetry transformation within M which is pulled back into M 0 (via J −1 n .) We exemplify this below. Example 5 Equation (10) holds in the 2-D navigation task π ∗ (a| J −1 n φ −1 h ̃ L n φ h (s);z 0 ) = z 0 − B −1 · (B(s− z) + z) = (−1)· (s− B· z 0 ) = π ∗ (a| s;J h · z 0 ). V. LEARNING PROBLEM Herein, we develop our learning problem, split into meta- train and meta-test. At meta-train time, we aim to learn the hereditary geometry G, that is its left actions L g and K g and their representations (φ,η). At meta-test time, we aim to infer the left action of the test task. Adopting the model based approach from Lemma 1, we now show how to convert (4) into a differentiable optimization objective. The two main characteristics of our learning problem are that it is (i) symmetry based and (i) defined on the differential of the underlying system. Fix a task M 0 ∈ M 1 ,...,M N train induced by h 0 ∈ H and assume w.l.o.g. that h 0 = e. 2 Given another task M ∈ M 1 ,...,M N train induced by some h ∈ H , we now aim to discover left actions L g and K g such that for all s∈ S,a∈ A R(s,a;J h · z 0 ) = R(L g · s,K g · a;z 0 ), T (s ′ | s,a;J h · z 0 ) = T (L g · s ′ | L g · s,K g · a;z 0 ).(12) In the final step, we then learn the left actions L g 0:N tasks and K g 0:N tasks of all tasks independently and generators W S and W A of L g and K g . (We define these generators momentarily.) Mapping (s,a) 7→ (L −1 g · s,K −1 g · a), and recalling that g ∈ G⇔ g −1 ∈ G, the above equivalently reads R h (L g ·s,K g ·a) : = R(L g ·s,K g ·a;J h ·z 0 ) = R(s,a;z 0 ) = R(s,a;J h 0 · z 0 ) = : R 0 (s,a). (13) 2 Otherwise define a new group ̃ H where each left action is shifted by (h 0 ) −1 . T h (L g ·s ′ | L g ·s,K g ·a) : = T (L g ·s ′ | L g ·s,K g ·a;J h ·z 0 ) = T (s ′ | s,a;z 0 ) = T (s ′ | s,a;J h 0 · z 0 ) = : T 0 (s ′ | s,a). (14) Given black box sampling access to R h ,R 0 ,T h and T 0 , we could, in theory, approximate all functions with neural networks and minimize mean squared errors of (13) and (14). Instead of comparing the functions everywhere, we only com- pare the reward functions within their respective symmetries as we expect the hereditary geometry to arise from these. For the transition functions we use the approach described before and defer categorizing symmetries as in (6) to future work. To this end, we now aim to better understand the symmetries of R to only search therein. Note that the symmetry of R is precisely the conventional invariance statement. A. From functional to differential symmetries Let M =S,A,R,T,γ denote a task with symmetry N . We now provide a necessary condition for the N -invariance of R via the n-invariance of its differential dR. n denotes the Lie algebra of N which is the set of all smooth left-invariant vector fields on N that is isomorphic to the tangent space of N at the neutral element e by the differential of left actions. This enables studying the symmetries induced by N by only considering tangent space at a single group element. In §VI we will show that this approach is thereby substantially more sample efficient and stable. The Lie algebra n is a d-dimensional vector space with basis W = W 1 ,...,W d where d is the dimension of N . Assuming a linear group structure in N , the left action L n has a representation ρ S : n → gl(|S|,R) in the Lie algebra gl(|S|,R) of the general linear group GL(|S|,R). Matrices ρ S (W i ) : = W i S ∈ R |S|×|S| are called the infinitesimal gener- ators of L n which we stack into a tensor W S ∈ R d×|S|×|S| . If N is connected and compact, the set of all left actions L n for n ∈ N can be characterized with W S and the matrix- exponential exp (hence the name) [ n∈N L n = [ t∈R d exp(t· W S ),(15) where · denotes standard vector-tensor multiplication and similarly for K n and its generator W A ∈ R d×|A|×|A| . Then, existing symmetry discovery approaches parameterize W S ∈ R d×|S|×|S| and W A ∈ R d×|A|×|A| and minimize a mean- squared-error type loss of (5) by randomly sampling weights t∈ R d [23], [24]. We now show that one can forego such sampling and evaluate the invariance only at a single group element by con- sidering the differential dR of R. So, differential symmetries of R are captured by its kernel distribution: Definition 5 (Kernel distribution) The kernel distribution D R of R contains the directional derivatives along R’s level sets D R (s,a) : =v ∈ T (s,a) (S,A)| dR (s,a) (v) = 0 (s,a)∈ S× A D R : = S (s,a)∈S×A D R (s,a) , where T s,a (S,A) denotes the tangent space of S and A at (s,a) which is, roughly speaking, the best linear approxima- tion of S× A at (s,a). Then, we have the following result: Lemma 3 (Differential symmetry implies functional symmetry) If the generators of L n and K n preserve R’s kernel distribution, that is for all s∈ S,a∈ A,i = 1,...,d W i S 0 0 W i A · (s,a)∈ D R (s,a) ,(16) then R is N -invariant: R(L n s,K n a) = R(s,a),∀n∈ N . Proof.We begin by differentiating the invariance constraint in (5) w.r.t. (s,a). Then, we argue that the N -invariance of R only has to hold for group elements n≈ e which are infinites- imally close to the identity e as (5) holds globally in S× A. Thereby, L n and K n can be expressed with the generators W S and W A via L n = Id +c·W S and K n = Id +c·W A for some constant c and Id denotes the identity mapping. Assuming that the left action is linear, that is W S and W A are constant in s and a, and using that the derivative operator is linear, one can then rewrite the differential of the invariance constraint into the stated claim.□ Compare how the functional invariance in (5) must be evaluated on S × A× N while the differential invariance in (16) only on S×A. Note that the above provides a necessary condition that, as we will see, can be turned into a learning problem from samples. Even though this approach is very efficient, the sufficiency of these conditions must be studied which is deferred to our future work. B. A symmetry based approach for hereditary geometry dis- covery We now learn the hereditary geometry by only comparing R h and R 0 along their differential symmetries, that is within the kernel distribution. This can be achieved by establishing the following necessary condition for the invariance property in (13). For (12) we directly compare the two transition functions T h and T 0 . Lemma 4 A necessary condition for R h (L g · s,K g · a) = R 0 (s,a) to hold on S × A is that the differentials dL g and dK g of L g and K g push the differential symmetries D R 0 of R 0 into the differential symmetries D R h of R h , that is for all s∈ S,a∈ A,v ∈ E R 0 s,a d(L g ,K g ) s,a [v]∈ span(E R h L g ·s,K g ·a ),(17) where E R denotes a basis of D R , also called a frame. Note that as group actions are invertible, the above also implies that L −1 g ,K −1 g push the symmetries of R h into the symmetries of R 0 . Then, the R-symmetries of M 0 and M are also called G-related. (cf. Page 182 in [15]) and we see them as the same up to G. Proof. [Proof of Lemma 4] The claim is a direct consequence of the following implications R h (L g · s,K g · a) = R 0 (s,a) ⇒ dR h L g ·s,K g ·a ◦ d(L g ,K g ) s,a [v] = dR 0 s,a [v],∀v ∈ T s S× T a A ⇒ dR h L g ·s,K g ·a ◦ d(L g ,K g ) s,a [v] = dR 0 s,a [v] = 0,∀v ∈ D R 0 s,a ⇔ d(L g ,K g ) s,a [v]∈ D R h L g ·s,K g ·a ,∀v ∈ D R 0 s,a ⇔ d(L g ,K g ) s,a [v]∈ span(E R h L g ·s,K g ·a ),∀v ∈ E R 0 s,a , which hold for all s∈ S,a∈ A.□ As L g and K g can be linearized by φ and η, (12) and (13) can be cast into an optimization problem with Objectives O1a and O1b as follows: Find ̃ L g ∈ GL(|S|,R) and ̃ K g ∈ GL(|A|,R) and diffeomorphisms (encoders/decoders) φ : S → ̃ S and η : A → ̃ A (and correspondingly φ −1 : ̃ S → S and η −1 : ̃ A→ A) such that for all s,s ′ ∈ S,a∈ A,v ∈ E R 0 s,a d(L g ,K g ) s,a [v]∈ span(E R h L g ·s,K g ·a )(O1a) T h (L g · s ′ | L g · s,K g · a) = T 0 (s ′ | s,a)s(O1b) where φ −1 ̃ L g φ : = L g and η −1 ̃ K g η : = K g to keep the notation consistent. Finally, using (15), learning generators W S of ̃ L g and W A of ̃ K g can be cast into another optimization with Objectives O2a and O2b: Learn differential generators W S ∈ R d×|S|×|S| and W A ∈ R d×|A|×|A| that contain the left actions ( ̃ L g ) 1:N tasks and ( ̃ K g ) 1:N tasks : ( ̃ L g ) 1:N tasks ∈ exp(span(W S ))⇔ log(( ̃ L g ) 1:N tasks )∈ span(W S ) (O2a) and similarly log(( ̃ K g ) 1:N tasks )∈ span(W A )(O2b) where log denotes the matrix logarithm. C. Practical implementation of (O1) and (O2) We smoothen the binary span constraints a ∈ span(b) of tensor a w.r.t. the vector space with basis b from (O1a), (O2a) and (O2b) by minimizing the orthogonal complement, denoted by a b ⊥ , yielding our loss functions L R geo ( ̃ L g , ̃ K g ,φ,φ −1 ,η,η −1 | E R h ,E R 0 ) = E s,a∈S×A h X v∈E R 0 (s,a) d(L g ,K g ) s,a [v] (E R h (L g ·s,K g ·a)) ⊥ 2 + λ 1 ( φ −1 φ(s)− s 2 + η −1 η(a)− a 2 ) i + λ 0 ∥φ,φ −1 ,η,η −1 ∥ 1 , (18) L T geo ( ̃ L g , ̃ K g ,φ,φ −1 ,η,η −1 | T h ,T 0 ) = E s,a∈S×A ||T h (L g · s ′ | L g · s,K g · a)− T 0 (s ′ | s,a)||, (19) L gen (W □ | log(□ g )) =∥ log( ̃ □ g ) (W □ /∥W □ ∥ F ) ⊥ ∥ F + λ 0 ∥W □ ∥ 1 , (20) where the sum of v ∈ E R 0 (s,a) ranges over the basis vectors from the frame E R 0 (s,a), W □ denotes a generator (which we set to W S or W A ) and ̃ □ g a linear left action induced by W □ (which we set to L g or K g ), λ 0:1 > 0 weights and ∥·∥ F the Frobenius norm. Additionally, we added Lasso regularizers on the weights of all decision variables and a reconstruction regularizer for the representations φ and η. 10 ° 10 - 1 Differential Agent 10 - 2 0 5 10 15 20 25 30 35 40 Gradient Step xl0 3 0 5 10 15 20 25 30 35 40 Gradient Step xl0 3 10 ° 10 2 Functional Loss Differential Loss Functional Agent Oracle Fig. 2. Differential (green) and Functional (blue) symmetry discovery agents against an Oracle (red) evaluating either the differential loss (left) or the functional loss (right) over time—lower is better. The differential symmetry discovery (green) is an order of magnitude more sample efficient and stable. To simplify the optimization, we trivialize the constrained domain ̃ L g ∈ GL(|S|,R) ⊊ R |S|×|S| by learning the dif- ferential log( ̃ L g ) ∈ gl(|S|,R) = R |S|×|S| and similarly for ̃ K g . As a practical benefit, it is far easier to compute the exponential of a matrix than its logarithm. (See [27] for a discussion.) Finally, we sample from the replay buffer using a Boltzmann distribution with a tuneable temperature to sample more uniformly from the support of the occupation measure. We now state our final learning problem. Given estimates E R 0 ,...,E R N train and T 0 ,...,T N train of the reward function kernels and transition functions, we minimize the sum of the three loss functions (18), (19) and (20) min data geo ,W S ,W A L gen (W S | logL g i ) +L gen (W A | logK g i ) + N tasks X i=1 L R geo (data geo | E R 0 ,E R i ) +L T geo (data geo | T h i ,T 0 ) where data geo = (log ̃ L g i , log ̃ K g i ,φ,φ −1 ,η,η −1 ) and we convert log ̃ L g i and log ̃ K g i into ̃ L g i and ̃ K g i via the matrix exponential. D. Meta-test At meta-test time, we sample a replay buffer in the test task M according to a random policy to estimate the kernel and transition function of M. Then, we minimize min c∈R d L R geo (L g ,K g | data geo ) +L T geo (L g ,K g | data geo )(21) where L g = exp(c·W S ) and K g = exp(c·W A ) and data geo = (W S ,W A ,φ,φ −1 ,η,η −1 ,E R ,E R 0 ,T h ,T 0 ). Note that we can use the geometric structure and its representation discovered at the training time, only searching within the spans of the generators W S and W A . VI. EMPIRICAL EVALUATION The following questions are addressed to empirically vali- date our approach: (a) how a kernel frame E R can be learned from samples, (b) how differential and functional symmetry discovery compare, and (c) whether geometric Meta-RL can improve generalization. As a test bench, we use the 2-D navigation task in Exam- ple 1. Concerning (a), we adopt a simple numerical differ- entiation approach. Concerning (b), we compare the sample efficiency of our proposed differential symmetry discovery to Augerino [23] which minimize the functional symmetry constraint in (5). Concerning (c), we uniformly sample only 4 tasks M∼ U (M) at train time—for an emphasis on the fact that the meta-dataset is not necessarily densely sampled from the task space. After an exploration period at test time, we roll-out the agent without further training. Then, we compute the regret against an optimal policy in the test task and plot it against the distance to the closest training task. We compare our approach against Contrastive learning augmented Context- based Meta-RL (CCM) [28] which combines contrastive learn- ing for task encoding and SAC for policy training. We begin by learning a frame of a kernel distribution. Let B =s i ,a i ,R(s i ,a i ),T (s i ,a i ) N i=1 denote a replay buffer of some training task and ε 1 > 0 denote a hyperparameter that determines when (s,a) and (s ′ ,a ′ ) are ”close”. Then, compute for each sample (s,a)∈B its neighbors N (s,a): N (s,a) =(s ′ ,a ′ )∈B |||(s,a)− (s ′ ,a ′ )|| 2 ≤ ε 1 Fix ε 2 > 0, and define the local level set ˆ L (s,a) of a point (s,a)∈ S× A: ˆ L (s,a) : =(s ′ ,a ′ )∈N (s,a)| R(s ′ ,a ′ )− R(s,a) 2 ≤ ε 2 By Taylor expansion of R around (s,a) ∈ B, we have R(s ′ ,a ′ ) ≈ R(s,a) + dR (s,a) (s ′ − s,a ′ − a) so (s ′ − s,a ′ − a) ∈ D R ((s,a)) for (s ′ ,a ′ ) ∈ ˆ L (s,a) . An orthonormal basis E R (s,a) for each (s,a) can then be computed via PCA on (s ′ − s,a ′ − a) for all (s ′ ,a ′ )∈ ˆ L (s,a) . Finally, we smoothly connect the point-wise bases E R (s,a) for all (s,a) ∈ B into a frames E R via a dense neural network. Next, we define the differential loss L diff to be the 2-norm of the orthogonal complement of (16): L diff (W S ,W A |B,E R ) : = E (s,a)∼B " d X i=1 W i S 0 0 W i A s a (E R (s,a)) ⊥ 2 # , and the functional loss L func as the mean squared error of (5): L func (W S ,W A |B) : = E (s,a)∼B E t∼N (0 d×1 ,1 d×d ) [||R(exp(t· W S )· s, exp(t· W A )· a;z)− R(s,a;z)|| 2 ], where ”time”-constants t∈ R d are sampled to generate differ- ent left actions. (See [24] for further discussion.) Then, Figure 2 shows loss over time (lower is better) of the differential (green) and the functional (blue) symmetry discovery agents, evaluated with the differential loss (left) or the functional loss (right)—with results averaged across 10 runs with identical initializations. In both cases, differential symmetry discovery agent converges an order of magnitude faster (2.5 vs. 25k steps) with lower variance. The overall final lower loss of the functional approach did not result in any meaningful changes 100 ---j x CCM x Geometric agent Trend CCM Trend Geometric Agent w 6 0 +---+------+------t----------t------'--::;;;�-------'--'---+------------1 C'> QJ C:: 4 0 +---+------+--------,.,,.-----=---'f--------'�------!------+...,,.---------j 0.0 0.2 0. 4 0.6 0.8 D istance of test task to the closest training task (radians) Fig. 3. Generalization in the 2-D navigation task: The regret versus the distance to closest training task for our geometric approach (green) and CCM (blue)—lower is better. The CCM agent generalizes well to nearby tasks but collapses for tasks distant to the training set while the geometric agent generalizes within the entire task space. in the learned generator; Both methods converged against the ground-truth SO(2,R) symmetry. Finally, Figure 3 shows the regret R(π;M) as a func- tion of distance to closest training task (lower is better) for our proposed geometric approach (in green) versus CCM (in blue). CCM generalizes by local smoothness where the regret increases monotonically in distance to the closest training task. In contrast, the geometric agent can general- ize within the entire task space. The code-base for these simulations can be found at https://github.com/PaulNitschke/ Hereditary-Geometries. VII. CONCLUSIONS In this paper we studied non-local generalization in Meta- RL via symmetries which we formalized as a hereditary ge- ometry. We argued that hereditary geometries commonly arise from the symmetries of the system and proposed a symmetry- based learning problem that is empirically validated in a 2-D navigation task. We conclude this paper by discussing some shortcomings and potential future work. First, in this study we assumed a model based perspective to discover G which only provides a sufficient but not a necessary requirement, that is we may converge to sub-symmetries. Second, even though in the implementation we only consider invariance for simplicity (which is sufficient for the 2-D Navigation task,) we emphasize that the policy generalization by equivariance may be considered similarly for an effective outcome in more general settings. REFERENCES [1] R. Kirk, A. Zhang, E. Grefenstette, and T. Rockt ̈ aschel, “A Survey of Zero-shot Generalisation in Deep Reinforcement Learning,” J. Artif. Int. Res., vol. 76, May 2023. [2] J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick, “Learning to reinforcement learn,” Jan. 2017, arXiv:1611.05763 [cs]. [3] K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen, “Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Vari- ables,” in Proceedings of the 36th International Conference on Machine Learning. PMLR, May 2019, p. 5331–5340. [4] Z. Mandi, P. Abbeel, and S. James, “On the Effectiveness of Fine-tuning Versus Meta-reinforcement Learning,” Advances in Neural Information Processing Systems, 2022. [5] L. Li, H. Zhang, X. Zhang, S. Zhu, Y. Yu, J. Zhao, and P.-A. Heng, “Towards an Information Theoretic Framework of Context-Based Of- fline Meta-Reinforcement Learning,” Advances in Neural Information Processing Systems, 2024. [6] L. Li, R. Yang, and D. Luo, “FOCAL: Efficient Fully-Offline Meta- Reinforcement Learning via Distance Metric Learning and Behavior Regularization,” in Int. Conf. on Learning Representations, Oct. 2020. [7] L. Zintgraf, S. Schulze, C. Lu, L. Feng, M. Igl, K. Shiarlis, Y. Gal, K. Hofmann, and S. Whiteson, “VariBAD: variational Bayes-adaptive deep RL via meta-learning,” J. Mach. Learn. Res., Jan. 2021. [8] F. Graf, C. Hofer, M. Niethammer, and R. Kwitt, “Dissecting Super- vised Contrastive Learning,” in Proceedings of the 38th International Conference on Machine Learning. PMLR, Jul. 2021. [9] S. Damrich and F. A. Hamprecht, “On UMAP’ s True Loss Function,” in Advances in Neural Information Processing Systems, 2021. [10] J. N. B ̈ ohm, P. Berens, and D. Kobak, “Attraction-Repulsion Spectrum in Neighbor Embeddings,” Journal of Machine Learning Research 23 (2022) 1-32, 2022. [11] S. Lee, M. Cho, and Y. Sung, “Parameterizing Non-Parametric Meta- Reinforcement Learning Tasks via Subtask Decomposition,” Advances in Neural Information Processing Systems, vol. 36, 2023. [12] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal Policy Optimization Algorithms,” arXiv:1707.06347 [cs]. [13] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor,” Aug. 2018, arXiv:1801.01290 [cs]. [14] S. K. Pal, T. S. Dillon, and D. S. Yeung, Eds., Soft computing in case based reasoning. Berlin, Heidelberg: Springer-Verlag, Sep. 2000. [15] J. M. Lee, Introduction to Smooth Manifolds, ser. Graduate Texts in Mathematics. Springer-Verlag, 2012, vol. 218. [16] F.Schuller,“GeometricAnalysisofTheoreticalPhysics,” 2015.[Online].Available:https://w.youtube.com/playlist?list= PLPH7f 7ZlzxTi6kS4vCmv4ZKm9u8g5yic [17] J. Beck, R. Vuorio, E. Z. Liu, Z. Xiong, L. Zintgraf, C. Finn, and S. Whiteson, “A Survey of Meta-Reinforcement Learning,” Aug. 2024, arXiv:2301.08028 [cs]. [18] F. Doshi-Velez and G. Konidaris, “Hidden Parameter Markov Decision Processes: A Semiparametric Regression Approach for Discovering Latent Task Parametrizations,” Aug. 2013, arXiv:1308.3513 [cs]. [19] A. Hallak, D. D. Castro, and S. Mannor, “Contextual Markov Decision Processes,” Feb. 2015, arXiv:1502.02259 [stat]. [20] S. Lee and S.-Y. Chung, “Improving Generalization in Meta-RL with Imaginary Tasks from Latent Dynamics Mixture,” in Advances in Neural Information Processing Systems, Nov. 2021. [21] H. Fu, J. Yao, O. Gottesman, F. Doshi-Velez, and G. Konidaris, “Per- formance Bounds for Model and Policy Transfer in Hidden-parameter MDPs,” in The Eleventh International Conference on Learning Repre- sentations, Sep. 2022. [22] E. van der Pol, D. Worrall, H. van Hoof, F. Oliehoek, and M. Welling, “MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning,” in Advances in Neural Information Processing Systems, vol. 33. Curran Associates, Inc., 2020, p. 4199–4210. [23] G. Benton, M. Finzi, P. Izmailov, and A. G. Wilson, “Learning Invari- ances in Neural Networks from Training Data,” in Advances in Neural Information Processing Systems, vol. 33. Curran Associates, Inc., 2020, p. 17 605–17 616. [24] J. Yang, N. Dehmamy, R. Walters, and R. Yu, “Latent space symmetry discovery,” in Proceedings of the 41st International Conference on Machine Learning, ser. ICML’24, vol. 235. Vienna, Austria: JMLR.org, Jul. 2024, p. 56 047–56 070. [25] R. Sutton and A. Barto, Reinforcement Learning: An Introduction, 2nd ed. Cambridge: MIT press, 2018. [26] P. W. Anderson, “More Is Different,” Science, vol. 177, no. 4047, p. 393–396, Aug. 1972. [27] Contributors, “PyTorch Issues: [feature request] Add matrix functions \#9983,” Jul. 2018. [Online]. Available: https://github.com/pytorch/ pytorch/issues/9983 [28] H. Fu, H. Tang, J. Hao, C. Chen, X. Feng, D. Li, and W. Liu, “Towards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 35, May 2021, p. 7457–7465.