SKooP: Symmetric Koopman Predictions for Faster and More Generalizable Legged Robot Locomotion with Reinforcement Learning

SKooP: Symmetric Koopman Predictions for Faster and More Generalizable Legged Robot Locomotion with RL

Paper: SKooP: Symmetric Koopman Predictions for Faster and More Generalizable Legged Robot Locomotion with Reinforcement Learning

Authors: Evelyn D'Elia, Weishu Zhan, Giulio Turrisi, Giulio Romualdi, Giuseppe L'Erario, Raffaello Camoriano, Wei Pan, Daniele Pucci (IIT / Univ. of Manchester / Politecnico di Torino / Newcastle Univ.)

Links: Project Page · arXiv:2607.11624

Code: Built on open-source libraries DynamicsHarmonicsAnalysis, MorphoSymm, ESCNN

---

One-Sentence Summary

SKooP combines morphological symmetries with a Koopman autoencoder, learning a symmetric linearized latent dynamics model online and using multi-step Koopman predictions as privileged observations for the critic, enabling faster convergence and better symmetric generalization on bipedal locomotion tasks—a quadruped's mirrored push-door success rate improves from 0% to 77.7%.

---

Background and Motivation

RL algorithms classically suffer from poor sample efficiency. In robotics, a recent line of work addresses this by encoding physics priors in the learning process. However, most approaches are validated on well-defined, low-dimensional benchmark systems rather than high-dimensional robots with complex nonlinear dynamics.

Legged robots have complex, nonlinear dynamics where learning effective policies remains challenging. Purely data-driven model-free RL relies on costly trial and error; traditional model-based approaches exploit known physics but require expert knowledge and extensive manual tuning. Fusing model-based control with RL holds promise for overcoming their respective limitations.

SKooP leverages two physics priors. The first is morphological symmetry—legged robots typically have $\mathbb{C}_2$ reflection symmetry, and symmetry constraints can reduce required samples and produce more generalizable policies. The second is Koopman theory—for any nonlinear system, there exists an infinite-dimensional function space where dynamics are globally linear. Finite approximations can be learned via autoencoders, lifting and linearizing high-dimensional state spaces at relatively low cost.

The key innovation is combining these two priors: learning a symmetric Koopman model via an equivariant autoencoder and using Koopman-predicted latent states as privileged observations for the critic. This enables the critic to learn based on smoother, more informative features, while equivariant constraints ensure policy generalization across symmetric tasks. No prior robotic control approach directly injects Koopman predictions into the policy learning process.

---

Method

Equivariant Network Architecture

SKooP implements symmetry priors in two ways. First, through structural constraints on actor and critic networks. The actor is constrained to be $\mathbb{G}$-equivariant and the critic to be $\mathbb{G}$-invariant. In the robot's joint space, due to $\mathbb{C}_2$ symmetry, a state $(q, \dot{q})$ can be reflected to obtain the symmetric state $(g_s \triangleright q, g_s \triangleright \dot{q})$. The symmetry rules are:

$$\pi_\theta(g \triangleright x) = g \triangleright \pi(x), \quad V_\xi(g \triangleright x) = V_\xi(x)$$

$$f_\phi(g \triangleright x) = g \triangleright z, \quad f_\phi^{-1}(g \triangleright z) = g \triangleright x$$

where $\pi_\theta$ is the actor, $V_\xi$ is the critic, $f_\phi$ and $f_\phi^{-1}$ are encoder and decoder. These networks are termed equivariant multilayer perceptrons (EMLPs).

Symmetric Koopman Autoencoder

SKooP trains an equivariant autoencoder concurrently with the policy, avoiding pretraining and ensuring training data relevance. This Equivariant Controlled Dynamics Auto-Encoder (ecDAE) linearizes dynamics as:

$$z_{k+H} = A^H f_\phi(x_k) + \sum_{i=0}^{H-1} A^{H-1-i} B u_{k+i}$$

$$x_k = f_\phi^{-1}(z_k)$$

where $A \in \mathbb{R}^{n \times n}$, $B \in \mathbb{R}^{n \times m}$ are linear system matrices, $H$ is the prediction horizon. Training uses three standard losses:

$$\mathcal{L}_{sr} = \|x_k - f_\phi^{-1}(f_\phi(x_k))\|_2^2 \quad \text{(state reconstruction)}$$

$$\mathcal{L}_{lp} = \|f_\phi(x_{k+H}) - \mathcal{K}^H f_\phi(x_k)\|_2^2 \quad \text{(latent prediction)}$$

$$\mathcal{L}_{sp} = \|x_{k+H} - f_\phi^{-1}(\mathcal{K}^H f_\phi(x_k))\|_2^2 \quad \text{(state prediction)}$$

Latent State Privileged Observations

The Koopman-predicted latent state $z_{k+1}$ is provided to the critic as a privileged observation. This enables the critic to learn a more precise value function based on the linear representation of system dynamics, simplifying the problem of predicting expected future returns. The actor only receives proprioceptive observations, not depending on Koopman predictions, ensuring no online Koopman inference is needed at deployment.

Online Training

The ecDAE is trained concurrently with PPO using a Prioritized Experience Replay (PER) buffer of 10,000 data points. Each task uses 8,192 parallel environments over 5 random seeds. The latent-to-state dimensionality ratio is 3, with prediction horizon $H=5$. A single SKooP training iteration takes 34.0% longer than baseline PPO, but faster convergence compensates for this overhead.

Architecture Pipeline

graph TB
    subgraph Policy Training
        OBS["Proprioceptive Obs x_k"] --> ACTOR["Equivariant Actor
π_θ(x)→u_k"]
        ACTOR --> ENV["Environment
Isaac Gym"]
        ENV --> NEXT["Next State x_k+H"]
        OBS --> ENC["Equivariant Encoder f_φ
x_k→z_k"]
        ENC --> KOOP["Koopman Prediction
z_k+H=A^H*z_k+ΣB*u"]
        KOOP --> CRITIC["Invariant Critic
V_ξ(x,z_k+H)"]
        OBS --> CRITIC
        CRITIC --> VALUE["Value Estimate"]
        NEXT --> DEC["Equivariant Decoder f_φ⁻¹
z_k+H→x_k+H"]
        ENC --> DEC
    end
    subgraph Autoencoder Training
        ENC --> LOSS["Loss L_sr+L_lp+L_sp"]
        DEC --> LOSS
        NEXT --> LOSS
        LOSS --> ENC
    end

---

Experimental Results

The paper tests three bipedal locomotion tasks on the Cyberdog2 quadruped: stand dance, walk slope, and push door. All require the quadruped to transition to an upright bipedal posture and walk on hind legs. Push door is the primary generalization evaluation—models train only on right-opening doors and test on left-opening (mirrored) doors.

Figure 1: SKooP architecture overview—equivariant Koopman autoencoder trained concurrently with actor-critic policy, predicting latent states as critic privileged observations.

Push Door Ablation

Method	Right SR ↑	Left SR ↑	SI ↓	OOD Right ↑	OOD Left ↑	OOD SI ↓
PPO	70.01±6.65	0.00±0.00	199.99	42.27±7.65	0.01±0.02	199.94
PPOeqic	39.65±5.07	39.12±6.31	13.82±6.88	21.43±4.44	21.72±4.10	17.68±13.44
SKooP-NoSym-NoPred	84.18±4.31	0.02±0.02	199.91	45.95±6.56	0.02±0.02	199.82
SKooP-NoSym	83.36±3.86	0.05±0.02	199.74	41.23±4.84	0.01±0.01	199.89
SKooP-NoPred	74.37±12.88	68.45±11.32	8.03±6.89	44.94±7.61	43.43±5.90	9.14±10.57
SKooP	82.77±2.96	77.70±6.66	8.49±4.26	47.65±2.59	46.18±4.37	9.06±9.24

Table 1: Push door ablation study (% mean±std). Data from paper Table I.

Full SKooP maintains high right-door success (82.77%) while improving left-door (mirrored) success from PPO's 0% to 77.70%. The symmetry index drops from 199.99 to 8.49 (lower is more symmetric). OOD performance is equally strong—left-door success of 46.18% vs PPO's 0.01%. The ablation reveals two key findings: symmetry constraints are necessary for mirrored generalization (no-symmetry variants achieve near-zero left-door success), and Koopman predictions further improve training reward and convergence speed.

Figure 2: Training reward ablation across three bipedal tasks—dance, slope, and push door, averaged over 5 seeds.

Koopman Model Stability and Controllability

Metric	cDAE (no symmetry)	ecDAE (with symmetry)
$\mathcal{L}_{sr}$ MSE	1.503±0.153	1.245±0.152
5-step $\mathcal{L}_{sp}$ MSE (Right)	2.144±0.078	1.778±0.153
5-step $\mathcal{L}_{sp}$ MSE (Left)	3.150±0.143	1.848±0.196
Invariant eigenvalue $|\lambda|$	N/A	1.000±0.003

Table 2: cDAE vs ecDAE MSE comparison (push door, 5 seeds). Data from paper Table IV.

Symmetry constraints enable ecDAE to achieve nearly identical state prediction error on the mirrored task as on the training task (1.848 vs 1.778), while cDAE degrades by ~50% on the mirrored task (3.150 vs 2.144). The ecDAE learns an invariant eigenvalue $|\lambda| \approx 1.0$ representing the invariance of system dynamics across symmetry groups. Removing this invariant mode significantly increases prediction error, confirming it encodes symmetry information.

Figure 3: Comparison of stand dance rear leg joint angles for PPO, PPOeqic, and SKooP.

Sim-to-Sim Transfer

Test Condition	$v_{xy}$ MSE (Isaac)	$v_{xy}$ MSE (Mujoco)	$\omega_z$ MSE (Isaac)	$\omega_z$ MSE (Mujoco)
Friction μ=0.4	0.0341	0.0245	0.2795	0.6017
Friction μ=1.0	0.0271	0.0231	0.1719	0.5048
Mass Offset 0.5kg	0.0273	0.0257	0.1656	0.2904
Push Vel 0.3m/s	0.0320	0.0257	0.1981	0.5325
Armature 0.005	0.0406	0.0227	0.2339	0.4986

Table 3: SKooP sim-to-sim robustness from Isaac Gym to Mujoco. Data from paper Table V.

Linear velocity tracking remains consistent across simulators, while yaw velocity tracking degrades in Mujoco but without failures. This demonstrates SKooP policies are robust to simulator differences and perturbation conditions.

---

Limitations

1. Increased Training Overhead: A single SKooP training iteration takes 34.0% longer than baseline PPO due to concurrent ecDAE training and Koopman prediction computation. While faster convergence partially compensates, overall computational cost remains higher.

2. Symmetry Constraints May Limit Performance: On the push door task, no-symmetry variants (SKooP-NoSym) achieve higher training-task rewards but completely fail to generalize to mirrored tasks. This indicates symmetry constraints trade training-task specialization for generalization.

3. Yaw Tracking Degrades in Sim Transfer: Sim-to-sim evaluation shows yaw velocity tracking significantly degrades in Mujoco (~2-3x), indicating room for improvement in robustness to simulator differences. Additionally, evaluation is conducted only in simulation without real hardware results.

---

Conclusion and Future Directions

SKooP's core contribution is being the first to directly inject Koopman predictions into the policy learning process, combined with morphological symmetries. The equivariant autoencoder's learned linearized latent space provides the critic with smooth future information, while symmetry constraints ensure policy generalization across mirrored tasks. The push door ablation clearly demonstrates the complementary roles of both priors: symmetry constraints are necessary for generalization, while Koopman predictions further improve convergence speed and reward.

Analysis of the learned Koopman model provides additional theoretical insights: ecDAE automatically discovers an invariant eigenvalue $|\lambda| \approx 1$ encoding system symmetry, achieving nearly identical state prediction error on mirrored tasks as on training tasks. Controllability analysis further shows symmetry constraints produce more compact representations by pushing redundant control directions toward zero.

The "equivariant networks + Koopman predictions + privileged observations" paradigm established by SKooP provides a modular framework for incorporating physics priors into RL. Future work can explore more symmetry group types, longer prediction horizons, and real hardware deployment.

"By learning a symmetric linearized latent space concurrently with the policy, we ensure the lifting function is trained on a relevant state space, and by providing the critic with Koopman-predicted latent states as information about the future, we enable it to learn higher-quality motions."