Enhancing generalization in reinforcement learning for robotic control through deep policy similarities

Abstract

This work introduces Deep Policy Similarities (DeePS), a learning-based bisimulation approach designed to enhance generalization in reinforcement learning for robotic control. Traditional reward-based bisimulation metrics often fail to enable effective policy transfer, particularly in environments with inconsistent reward structures. DeePS overcomes these limitations by leveraging policy similarities and approximating the policy similarity metric using a forward dynamics model. This approach facilitates more efficient and effective policy transfer across semantically equivalent environments. Through a series of control experiments, DeePS is shown to significantly outperform standard reinforcement learning methods and reward-based bisimulation approaches. In the noisy cartpole environment with randomized rewards, DeePS achieved 53.5% and 79.0% higher test rewards compared to Soft Actor-Critic and Deep Bisimulation for Control (DBC), respectively. Additionally, in a zero-shot evaluation on the Minigrid simple crossing environment, DeePS outperformed existing approaches, with test rewards 79.9% and 99.5% higher than DBC and RAPID, respectively. These results demonstrate that DeePS significantly enhances the ability of reinforcement learning models to generalize to diverse, unseen environments. This makes DeePS a valuable advancement for reinforcement learning, particularly in robotic applications where adaptability and robustness are critical.

Introduction

Reinforcement learning is a powerful framework for developing autonomous systems capable of making sequential decisions to maximize cumulative rewards. In the realm of robotic control, reinforcement learning enables robots to learn complex behaviors through interactions with their environment, eliminating the need for pre-programmed instructions. This capability is critical for tasks such as navigation, manipulation, and coordination in complex environments with undefined solutions.

A significant breakthrough in reinforcement learning for robotics was the development of Deep Q-Network, ¹ which integrated deep learning with traditional Q-learning. ² This innovative approach enabled reinforcement learning to manage high-dimensional state spaces, making it feasible to apply reinforcement learning techniques to more sophisticated robotic tasks. For instance, reinforcement learning was successfully applied to train robots for object manipulation, locomotion, and autonomous navigation.^3–6

Despite these advances, applying reinforcement learning to real-world robotic problems remains challenging, particularly in achieving generalization—the ability to transfer learned behaviors from training environments to unseen, semantically equivalent environments. Real-world environments are inherently unpredictable, posing a significant challenge for reinforcement learning models that are often optimized for static, well-defined environments. For instance, a robot trained with reinforcement learning to navigate a specific indoor environment may perform well in training, but struggle in new settings with different layouts or goals. This difficulty in adapting to changes highlights the fragility of current reinforcement learning approaches as a result of overfitting to training conditions leading to poor performances in novel situations.^7,8

Traditional reinforcement learning benchmarks such as MuJoCo ⁹ or the Arcade Learning Environment ¹⁰ evaluate agents in controlled and repeatable settings. While useful for assessing specific agent capabilities, these benchmarks fail to capture the complexity and variability of dynamic real-world environments. To this end, more recent benchmarks, such as those presented in Nichol et al., ¹¹ which introduce a train–test split reveal that standard deep reinforcement learning methods perform poorly under varied conditions. This discrepancy in performance between tightly controlled experiments and practical deployment further underscores the need for reinforcement learning models that are robust and can generalize effectively across diverse conditions.

To address this challenge, studies have explored various strategies to improve the generalizability of reinforcement learning in robotics. Techniques such as data augmentation^12–14 and domain randomization^15,16 introduce variability during training, helping to develop more robust models. Regularization and dropout^17,18 were also applied to neural models to reduce overfitting to training environments. In parallel, procedurally generated environments^7,19 offer a means to create diverse training conditions, alleviating concerns with overfitting. These environments continuously generate new and varied scenarios, guiding the agent to learn more adaptable policies. However, these existing methods largely fail to leverage the temporal dependency between sequential decisions when considering generalization for reinforcement learning—a characteristic that is unique to reinforcement learning.

More recent approaches have adapted the temporal properties of reinforcement learning to learn general policies with bisimulation techniques.^20–25 Bisimulation methods, in general, seek to encode state representations such that states leading to similar behaviors are represented similarly. This helps in learning policies that are more general and adaptable by recognizing and responding to behavioral similarities in different states. Notably, bisimulation approaches are flexible and can be integrated with other existing enhancement techniques such as the data augmentation or algorithmic modifications described previously.

Enhancing the generalization capabilities of reinforcement learning is crucial for deploying it in real-world robotics, making it a key focus of ongoing research. By addressing these generalization challenges, reinforcement learning can maximize its potential, enabling trained robots to perform a wide range of tasks autonomously and effectively in diverse settings.

The contribution of this work is as follows: (i) we introduce a novel learning-based bisimulation approach that leverages policy-driven actions to enhance the generalization capabilities of reinforcement learning models efficiently and (ii) propose the use of an additional nonlinear projection layer on the encoder representations for improved performance. (iii) We also demonstrate that our proposed method improves generalization across diverse environments that are semantically equivalent, through a series of control-related experiments.

These advancements represent a critical step toward developing generalized reinforcement learning policies that are not only effective in controlled scenarios, but also capable of handling the complexities and unpredictability of real-world applications. Our findings suggest that integrating bisimulation with policy-driven strategies using a learning-based approach can bridge the gap between theoretical research and practical deployment, paving the way for more reliable and versatile robotic systems.

Markov Decision Processes

Markov decision processes (MDPs) ²⁶ provide a mathematical framework for modeling decision-making in situations where outcomes are influenced by both the action of the agent and the probabilistic elements of the environment. MDPs are widely used in various fields of work, including robotics, transportation systems, ²⁷ and artificial intelligence, particularly in the context of reinforcement learning.

An MDP is defined by a tuple $(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma )$ where:

$\mathcal{S}$ is a finite set of states representing all possible situations the agent might encounter.

$\mathcal{A}$ is a finite set of actions available to the agent at each state.

$\mathcal{P}:\mathcal{S}\times \mathcal{A}$ is the state transition probability function, where $\mathcal{P}({\mathbf{s}}^{\prime }|\mathbf{s},\mathbf{a})$ represents the probability of transitioning to state ${\mathbf{s}}^{\prime }$ from state $\mathbf{s}$ after taking action $\mathbf{a}$ .

$\mathcal{R}:\mathcal{S}\times \mathcal{A}$ is the reward function, where $\mathcal{R}(\mathbf{s},\mathbf{a})$ gives the immediate reward received after taking action $\mathbf{a}$ in state $\mathbf{s}$ .

$\gamma \in [0,1]$ is the discount factor, which quantifies the importance of future rewards compared to immediate rewards.

The objective in reinforcement learning is to find an optimal policy $\pi *$ that maximizes the expected cumulative reward for the agent over time and the policy $\pi :\mathcal{S}\to \mathcal{A}$ determines the action that the agent should take in each state.

MDPs assume the Markov property, where the future states depend only on the current state and action, simplifying decision-making. While this assumption facilitates modeling and optimization within individual environments, generalization across environments requires understanding and identification of similarities between the states for effective policy transfer. Solving an MDP focuses on maximizing returns in a single environment^28–30 and generalization requires the policy to recognize and adapt to similarities across multiple MDPs.

Although deep reinforcement learning techniques have shown remarkable success in developing effective policies for individual MDPs, $\mathcal{M}$ , they face difficulties when agents are required to operate across a collection of environments, represented as $\rho =\{{\mathcal{M}}_1,\dots ,{\mathcal{M}}_n\}$ , after being trained on only a limited subset of the collection akin to zero-shot learning. Agents trained on this limited subset often struggle to adapt to new but semantically related scenarios. This highlights the critical challenge of generalization in reinforcement learning: ensuring policies perform robustly in unseen environments with shared behavioral structures.

MDPs form the backbone of reinforcement learning algorithms, where an agent learns to make decisions by interacting with the environment. Through trial and error, the agent seeks to discover the optimal policy that yields the highest rewards and MDPs provide the theoretical foundation for understanding how agents can learn effective strategies in these uncertain environments. It offers a robust framework for modeling complex decision-making problems essential for developing intelligent robotic systems capable of operating autonomously.

Bisimulation in Reinforcement Learning

In MDPs, bisimulation provides a systematic approach to measure the similarity between states, where bisimilar states are expected to exhibit similar behavior. This approach is especially useful for tasks such as model reduction or state aggregation, ³¹ which aim to simplify a complex MDP while maintaining key behavioral characteristics. A significant advancement for applying bisimulation in MDPs is the introduction of the bisimulation metric, ²⁰ which provides a quantitative measurement of the similarity between states, focusing particularly on their reward structures.

Definition 1 (Bisimulation metrics ²⁰ )

Let $\mathcal{F}:\mathbb{M}\to \mathbb{M}$ be defined by

$$\begin{array}{rl}\mathcal{F}(d)({\mathbf{s}}_i,{\mathbf{s}}_j)=& \underset{\mathbf{a}\in \mathcal{A}}{max}(|\mathcal{R}({\mathbf{s}}_i,\mathbf{a})-\mathcal{R}({\mathbf{s}}_j,\mathbf{a})|\\ & +c{W}_1(d)(\mathcal{P}(\cdot |{\mathbf{s}}_i,\mathbf{a}),\mathcal{P}(\cdot |{\mathbf{s}}_j,\mathbf{a})))\end{array}$$

(1)

where the operator $\mathcal{F}$ has a unique fixed-point $d_{\sim }$ in the metric space $\mathbb{M}$ representing the bisimulation metric.

This bisimulation metric provides a nuanced method for evaluating similarities between states in MDPs, comprising two main components: the absolute difference in immediate rewards from specific actions and the variation in future state transitions, assessed using the 1-Wasserstein distance, ³² denoted as $W_1$ , which is adjusted by a discount factor $c$ , similar to the $\gamma$ in MDPs. Derived from optimal transport theory, the 1-Wasserstein distance quantifies the similarity between two distributions by determining the minimum cost to transform one distribution into the other. Essentially, a lower cost indicates a higher similarity between the distributions, implying that the corresponding actions lead to similar subsequent states and thus comparable future values.

The bisimulation metric evaluates the differences between states ${\mathbf{s}}_i$ and ${\mathbf{s}}_j$ by examining the maximum disparity across all actions. While effective, the use of the $max$ operator can lead to overly pessimistic estimates, limiting policy transfer, ³³ and increasing computation costs, particularly in environments with continuous action spaces. These challenges make the bisimulation metric impractical for complex applications.

The on-policy bisimulation metric ²² was introduced to address these limitations by focusing on a specific policy $\pi$ and evaluating only actions it selects. By disregarding actions not taken under the chosen policy, the metric reduces unnecessary pessimism and computational complexity, enhancing the efficiency and practicality of bisimulation in reinforcement learning, especially in environments with large action spaces.

Definition 2 ( $\pi$ -bisimulation metric ²² )

Let $\mathcal{F}:\mathbb{M}\to \mathbb{M}$ be defined by

$$\begin{array}{rl}\mathcal{F}(d)({\mathbf{s}}_i,{\mathbf{s}}_j)=& |{\mathcal{R}}^{\pi }({\mathbf{s}}_i)-{\mathcal{R}}^{\pi }({\mathbf{s}}_j)|\\ & +c{W}_1(d)({\mathcal{P}}^{\pi }(\cdot |{\mathbf{s}}_i),{\mathcal{P}}^{\pi }(\cdot |{\mathbf{s}}_j))\end{array}$$

(2)

where the operator $\mathcal{F}$ has a unique fixed-point $d_{\pi }$ in the metric space $\mathbb{M}$ representing the $\pi$ -bisimulation metric.

However, relying solely on reward similarities may not be sufficient for effective policy transfer between environments which is essential for the broader adoption of these metrics in reinforcement learning. The policy similarity metric (PSM) represents a novel approach that differs from traditional reward-based bisimulation metrics by focusing on the alignment of actions derived from policy $\pi$ when evaluating state similarity. ²³ This method offers a new perspective and is particularly effective for policy transfer with its emphasis on the consistency of policy actions.

Definition 3 (Policy similarity metric ²³ )

Let $\mathcal{F}:\mathbb{M}\to \mathbb{M}$ be defined by

$$\begin{array}{rl}\mathcal{F}(d)({\mathbf{s}}_i,{\mathbf{s}}_j)=& \text{DIST}(\pi ({\mathbf{s}}_i),\pi ({\mathbf{s}}_j))\\ & +c{W}_1(d)({\mathcal{P}}^{\pi }(\cdot |{\mathbf{s}}_i),{\mathcal{P}}^{\pi }(\cdot |{\mathbf{s}}_j))\end{array}$$

(3)

where the operator $\mathcal{F}$ has a unique fixed-point $d*$ in metric space $\mathbb{M}$ representing the PSM.

In the PSM, emphasis is placed on the difference in agent actions measured using a pseudo-metric DIST. For discrete action spaces, DIST uses the total variation distance as its measure. In continuous domains, it corresponds to the $L_1$ distance between the means of the policy actions. We illustrate the differences between PSM and $\pi$ -bisimulation metric within the context of MDPs in Figure 1.

Figure 1.

Consider MDPs with states $\mathbf{x}$ and $\mathbf{y}$ . The optimal action at each state is represented by the colored path, with the value of the reward indicated by the color. (a) The distance between states ${\mathbf{x}}_1$ and ${\mathbf{y}}_1$ measured using the $\pi$ -bisimulation is non-zero, while the distance, according to PSM, is zero. (b) The $\pi$ -bisimulation indicates a distance of zero between states ${\mathbf{x}}_4$ and ${\mathbf{y}}_4$ , while the PSM measures a non-zero distance. Under $\pi$ -bisimulation, the policy is expected to behave differently when the states are deemed bisimilar, which is both anomalous and inconsistent. MDP: Markov decision process; PSM: policy similarity metric.

As shown in Figure 1, the action-based PSM provides a more effective comparison metric for generalizing learned policies. Unlike $\pi$ -bisimulation, which is heavily dependent on reward structures, the PSM focuses on the similarity of actions, emphasizing behavioral similarities based on agent reactions rather than environmental state values. This makes PSM more suitable for generalizing between states without imposing excessive restrictions on the reward similarities and the distinction is crucial because reward-based $\pi$ -bisimulation can introduce significant challenges in policy transfer, particularly when dealing with states that, despite having similar values, necessitate divergent actions to achieve optimal outcomes.

By prioritizing action-based comparisons, PSM facilitates more robust policy generalization and enhances the agent’s ability to adapt to new and varied environments. This capability is essential for real-world robotic applications where consistency and adaptability of the learned policies are critical for success.

Learning-based bisimulation metrics

Learning-based bisimulation approaches have been extensively explored^{22,24,25,34,35} in recent works, marking a significant progress in addressing the generalization challenges posed by environments with extensive or continuous state spaces. These advanced methods utilize deep learning to approximate the bisimulation metric, facilitating its scalable application in complex settings and contributing to the development of more robust and generalized reinforcement learning approaches that can be effectively applied across diverse environments.

The Deep Bisimulation for Control (DBC) ³⁴ is a notable example that approximates the $\pi$ -bisimulation metric by modeling environmental dynamics. It estimates the transition probabilities as a Gaussian distribution, parameterized by mean $\mu$ and covariance $\mathrm{\Sigma }$ and represents the next latent state probabilistically as ${\mathbf{h}}^{\prime }\sim \mathcal{N}(\mu ,\mathrm{\Sigma })$ . This allows the convenient substitution of the computationally intensive 1-Wasserstein distance with the 2-Wasserstein distance,^31,36 which has a closed-form solution given as:

$$\begin{array}{rl}{W}_2(\mathcal{N}({\mu }_i,{\mathrm{\Sigma }}_i),\mathcal{N}({\mu }_j,{\mathrm{\Sigma }}_j){)}^2=& \Vert {\mu }_i-{\mu }_j{\Vert }_2^2\\ & +\text{Tr}({\mathrm{\Sigma }}_i+{\mathrm{\Sigma }}_j-2({\mathrm{\Sigma }}_i^{1/2}{\mathrm{\Sigma }}_j{\mathrm{\Sigma }}_i^{1/2}{)}^{1/2})\end{array}$$

(4)

The $\pi$ -bisimulation metric can then be estimated as:

$$\begin{array}{rl}{\hat{d}}_{\pi }({\mathbf{s}}_i,{\mathbf{s}}_j)=& |{\mathcal{R}}^{\pi }({\mathbf{s}}_i)-{\mathcal{R}}^{\pi }({\mathbf{s}}_j)|\\ & +c{W}_2({\hat{\mathcal{P}}}^{\pi }(\cdot |\varphi ({\mathbf{s}}_i)),{\hat{\mathcal{P}}}^{\pi }(\cdot |\varphi ({\mathbf{s}}_j)))\end{array}$$

(5)

where $\hat{\mathcal{P}}$ is the dynamics model, and $\varphi$ is the encoder.

DBC directly enforces bisimilarity properties on the encoder representations by applying an auxiliary loss designed to aggregate bisimilar states (i.e. in this case, states with very similar values). This is achieved by minimizing the mean squared error:

$${\mathcal{L}}_{\text{DBC}}={\mathbb{E}}_{({\mathbf{s}}_i,{\mathbf{s}}_j)\sim \mathcal{D}}{[|\varphi ({\mathbf{s}}_i)-\varphi ({\mathbf{s}}_j)|-{\hat{d}}_{\pi }({\mathbf{s}}_i,{\mathbf{s}}_j)]}^2$$

(6)

Gradient stopping is used for the $W_2$ term to ensure that the ${\mathcal{L}}_{\text{DBC}}$ loss affects only the encoder, promoting the alignment of the latent space with the estimated $\pi$ -bisimulation metric.

However, the use of an approximate dynamics model to estimate the $\pi$ -bisimulation metric can result in significant issues, such as the runaway expansion of the latent space. ³¹ To prevent this unconstrained expansion, existing work proposed to normalize the latent space within a theoretically predefined closed ball ${\mathbb{B}}_c$ to constrain the search for the unique metric within a bounded space.

To the best of our knowledge, most learning-based bisimulation research continues to focus on the reward-based $\pi$ -bisimulation metric, despite its inefficiencies for policy transfer (see Figure 1). This limitation motivates our proposal of a learning-based approach that leverages the action-based PSM to address these shortcomings.

Deep policy similarities

This work proposes adapting the DBC methodology to learn generalized policies using action-based similarities, specifically the PSM. We call our proposed methodology the Deep Policy Similarities (DeePS), where a neural network $\hat{\mathcal{P}}$ is employed to model the environmental dynamics and estimate long-term differences in order to approximate the PSM.

Definition 4 (Policy similarity metric using approximate forward dynamics)

Let $\mathcal{F}:\mathbb{M}\to \mathbb{M}$ be defined by

$$\begin{array}{rl}\mathcal{F}(d)({\mathbf{s}}_i,{\mathbf{s}}_j)=& \text{DIST}(\pi ({\mathbf{s}}_i),\pi ({\mathbf{s}}_j))\\ & +c{W}_1(d)({\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_i),{\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_j))\end{array}$$

(7)

where the operator $\mathcal{F}$ has a unique fixed-point $\hat{d}*$ in metric space $\mathbb{M}$ representing the approximate PSM.

By representing the transition probabilities as a Gaussian distribution, $\mathcal{N}(\mu ,\mathrm{\Sigma })$ , it is possible to further write the approximate PSM as:

$$\begin{array}{rl}\hat{d}*=& \text{DIST}(\pi ({\mathbf{s}}_i),\pi ({\mathbf{s}}_j))+c{W}_2({\hat{\mathcal{P}}}^{\pi }(\cdot |\varphi ({\mathbf{s}}_i)),{\hat{\mathcal{P}}}^{\pi }(\cdot |\varphi ({\mathbf{s}}_j)))\end{array}$$

(8)

We present a theoretical bound for convergence guarantees on the approximate PSM, similar to the convergence guarantees established for the approximate $\pi$ -bisimulation in Kemertas and Aumentado-Armstrong. ³¹

Lemma 1 (Bounds on Wasserstein distances ³⁷ )

For any two distributions $\mu$ , $\lambda$ over a space $X$ , for all $p\ge 1$ ,

$$W_1(\mu ,\lambda )\le W_p(\mu ,\lambda )\le \text{diam}{(X)}^{(p-1)/p}{W}_1(\mu ,\lambda {)}^{1/p}$$

(9)

where $\text{diam}$ is the diameter of the space.

Theorem 1 (Boundedness condition for convergence)

Assume $\mathcal{S}$ is compact. If the support of an approximate forward dynamics model $\hat{\mathcal{P}}$ (i.e., ${\mathcal{S}}^{\prime }=\text{supp}(\hat{\mathcal{P}})$ ) is a closed subset of $\mathcal{S}$ , then there exists a unique approximate PSM $\hat{d}*$ as in equation 7. The bounds of this metric, dependent on the nature of the action space, are given as:

$$\text{diam}(\mathcal{S};\hat{d}*)\le \{\begin{array}{cc}\frac{1}{1-c}& \text{discrete }\\ \frac{1}{1-c}(A_{max}-A_{min})& \text{continuous }\end{array}$$

(10)

Proof.

The existence proof is virtually identical to the proof presented in Kemertas and Aumentado-Armstrong. ³¹ For continuous MDPs, the proof of the existence of the unique metric depends on the state space $\mathcal{S}$ being compact. ²¹ Therefore, if the support of the approximate forward dynamics model is a closed subset of the compact $\mathcal{S}$ , then the same proof can be used to show the existence of a unique metric.

$$\begin{array}{rl}\mathcal{F}(d)-\mathcal{F}(d^{\prime })& =c(W_1(d){\hat{\mathcal{P}}}_{i,j}^{\pi }-W_1(d^{\prime }){\hat{\mathcal{P}}}_{i,j}^{\pi })\\ & =c(W_1(d-d^{\prime }+d^{\prime }){\hat{\mathcal{P}}}_{i,j}^{\pi }-W_1(d^{\prime }){\hat{\mathcal{P}}}_{i,j}^{\pi })\\ & \le c(W_1(\Vert d-d^{\prime }{\Vert }_{\mathrm{\infty }}+d^{\prime }){\hat{\mathcal{P}}}_{i,j}^{\pi }-W_1(d^{\prime }){\hat{\mathcal{P}}}_{i,j}^{\pi })\\ & \le c(\Vert d-d^{\prime }{\Vert }_{\mathrm{\infty }}+W_1(d^{\prime }){\hat{\mathcal{P}}}_{i,j}^{\pi }-W_1(d^{\prime }){\hat{\mathcal{P}}}_{i,j}^{\pi })\\ & =c\Vert d-d^{\prime }{\Vert }_{\mathrm{\infty }},\mathrm{\forall }({\mathbf{s}}_i,{\mathbf{s}}_j)\in \mathcal{S}\times \mathcal{S}\end{array}$$

where ${\hat{\mathcal{P}}}_{i,j}^{\pi }=({\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_i),{\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_j))$ . This implies that $\mathcal{F}$ is a $c$ -contraction and a unique metric $\hat{d}*$ exists according to the Banach fixed-point theorem. To prove that the distance is bounded, according to Lemma 1,

$$\begin{array}{rl}\text{supp}(\hat{\mathcal{P}})\subseteq & \mathcal{S}⟹\underset{{\mathbf{s}}_i,{\mathbf{s}}_j\in \mathcal{S}\times \mathcal{S}}{sup}{W}_p(\hat{d}*)({\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_i),{\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_j))\\ & \le \text{diam}(\mathcal{S};\hat{d}*),\mathrm{\forall }p\ge 1\end{array}$$

Then,

$$\begin{array}{rl}\hat{d}*({\mathbf{s}}_i,{\mathbf{s}}_j)=& \text{DIST}(\pi ({\mathbf{s}}_i),\pi ({\mathbf{s}}_j))+c{W}_p(\hat{d}*)({\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_i),{\hat{\mathcal{P}}}^{\pi }(\cdot |{\mathbf{s}}_j))\end{array}$$

implies that

$$\hat{d}*({\mathbf{s}}_i,{\mathbf{s}}_j)\le \text{DIST}(\pi ({\mathbf{s}}_i),\pi ({\mathbf{s}}_j))+c\cdot \text{diam}(\mathcal{S};\hat{d}*)$$

and

$$\begin{array}{rl}\text{diam}(\mathcal{S};\hat{d}*)& \le \text{DIST}(\pi ({\mathbf{s}}_i),\pi ({\mathbf{s}}_j))+c\cdot \text{diam}(\mathcal{S};\hat{d}*)\\ & \le \frac{1}{1-c}\text{DIST}(\pi ({\mathbf{s}}_i),\pi ({\mathbf{s}}_j)),\mathrm{\forall }({\mathbf{s}}_i,{\mathbf{s}}_j)\in \mathcal{S}\times \mathcal{S}\end{array}$$

For discrete action spaces, $\text{DIST}$ is the total variation distance bounded between 0 and 1. For continuous spaces, $\text{DIST}$ is the $L_1$ between the means of the policy. Therefore,

$$\text{diam}(\mathcal{S};\hat{d}*)\le \{\begin{array}{cc}\frac{1}{1-c}& \text{discrete}\\ \frac{1}{1-c}(A_{max}-A_{min})& \text{continuous}\end{array}$$

where $A_{max}$ and $A_{min}$ are the upper and lower bounds on the action, respectively.

Instead of normalizing the latent states within a closed ball ${\mathbb{B}}_c$ , we notice that employing a nonlinear projection effectively prevents runaway expansion in the latent space. The projection layer disrupts the positive feedback loop within the DBC architecture, where encoder representations estimate transition probabilities that, in turn, constrain these representations, thus stabilizing the latent space. Moreover, the projection eases the constraints on the encoder. Direct enforcement of bisimilarity properties onto the encoder representations is too restrictive since the aggregated representations of two bisimilar observations will always produce the same value estimates. This conflicts with our expectation that states with different values but similar action sequences should be aggregated. We instead enforce the bisimilarity properties on a nonlinear projection such that bisimilar states need not have identical encoder representations, as long as their projections are similar. Therefore, the auxiliary loss function is given as:

$${\mathcal{L}}_{\text{DeePS}}={\mathbb{E}}_{({\mathbf{s}}_i,{\mathbf{s}}_j)\sim \mathcal{D}}{[|z_{\theta }({\mathbf{h}}_i)-z_{\theta }({\mathbf{h}}_j)|-\hat{d}*({\mathbf{s}}_i,{\mathbf{s}}_j)]}^2$$

(11)

where $\mathbf{h}=\varphi (\mathbf{s})$ is the encoder representation, and $z_{\theta }$ is the nonlinear projection. We present our DeePS architecture in Figure 2, where the policy network and dynamics model are trained with standard reinforcement learning algorithms (e.g. PPO ³⁸ , SAC ³⁹ ) and maximum likelihood estimation, respectively. The overall training process is demonstrated in Algorithm 1.

Figure 2.

The proposed DeePS architecture bears a strong resemblance to the DBC architecture, with the key distinction being the omission of the reward network and the additional projection layer. The encoder processes the input observation, extracting a latent representation $\mathbf{h}=\varphi (\mathbf{s})$ . This latent representation is the basis for the policy network $\pi$ and also the projection layer $z_{\theta }$ . The dynamics model $\hat{\mathcal{P}}$ takes a concatenation of the latent representation and policy action, $[\mathbf{h},\pi (\mathbf{h})]$ , to predict the next latent state in the form of a Gaussian distribution, ${\mathbf{h}}^{\prime }\sim \mathcal{N}(\mu ,\mathrm{\Sigma })$ . The bisimilarity properties are then enforced on the projection $z_{\theta }(\mathbf{h})$ instead of the latent representations. DeePS: Deep Policy Similarities; DBC: Deep Bisimulation for Control.

Experiments

Noisy sparse cartpole environment

We study the effectiveness of our proposed approach using a modified version of the traditional cartpole environment. This variant ³¹ offers a more challenging adaptation of the well-known OpenAI Gym environment ⁴⁰ and is specifically designed to test the robustness and adaptability of reinforcement learning models for control under increasingly complex conditions. The noisy sparse cartpole environment introduces three key modifications:

Sparse reward structure: Unlike the original environment, which provides rewards for keeping the pole within a wide angular range at each timestep, this adapted version restricts the rewards to a narrow angular window of $[-{\theta }_{\text{rew}},{\theta }_{\text{rew}}]$ with ${\theta }_{\text{rew}}$ drastically reduced to 1% of the original threshold.

Noisy observations: The environment complicates state observations by concatenating the state vector with $N_m\cdot \text{dim}(\mathcal{S})$ dimensions of noise, sampled from an isotropic Gaussian distribution, mimicking uncertainty in observations. This added noise challenges the agent’s ability to generalize and correctly aggregate noisy states into consistent representations.

Continuous action space: The action space, $\mathcal{A}\in [-1,1]$ , is redefined from discrete to continuous, demanding more nuanced control strategies. This shift requires the agent to make precise adjustments to the force applied to the cart, significantly increasing the complexity of the control problem and evaluating the agent’s ability to handle fine-grained action selection.

These adaptations not only increase the complexity of the task but also necessitate advanced techniques to accurately capture latent state representations for developing effective policies. By evaluating our approach in this challenging environment, which mimics the complexities of a robotic control problem, we aim to demonstrate its robustness and ability to learn generalized policies, highlighting its potential for real-world robotic applications.

We assess the performance of our proposed approach DeePS alongside its reward-based counterpart DBC, both of which employ the Soft Actor-Critic (SAC) ³⁹ as the standard reinforcement learning model. This analysis aims to underscore the strengths and potential weakness of our method in comparison to the state-of-the-art for generalization in reinforcement learning which leverages on the $\pi$ -bisimulation metric. For DBC, $L_2$ normalization is applied to the latent representations to ensure convergence. In contrast, DeePS uses a nonlinear projection to prevent runaway expansion in the latent state space.

In models that utilize intrinsic rewards, the forward model error in the latent space is employed.^31,41–43 This intrinsic reward is mathematically defined as:

$${\tilde{r}}_{I,t}={\eta }_r\Vert \hat{\varphi }({\mathbf{s}}_t,{\mathbf{a}}_t)-\varphi ({\mathbf{s}}_{t+1}){\Vert }_2^2/(2n)$$

(12)

where $\hat{\varphi }({\mathbf{s}}_t,{\mathbf{a}}_t)={\mathbb{E}}_{\hat{\varphi }({\mathbf{s}}_{t+1})\sim \hat{\mathcal{P}}(\cdot |{\mathbf{s}}_t,{\mathbf{a}}_t)}[\hat{\varphi }({\mathbf{s}}_{t+1})]$ denotes the expected latent state representation predicted by the approximate forward dynamics model, and ${\eta }_r>0$ is the weight coefficient. To ensure the intrinsic reward remains within practical bounds, it is capped at a maximum value of $R_{max}$ , leading to a final intrinsic reward of $r_{I,t}=min(R_{max},{\tilde{r}}_{I,t})$ . This upper bound on the intrinsic reward ensures it contributes positively to the learning process without overwhelming the true reward signal.

Network architecture and hyperparameters

The neural network architecture used in this experiment closely follows the design outlined in Kemertas and Aumentado-Armstrong. ³¹

The encoder network comprises a four-layer multilayer perceptron (MLP) with LeakyReLU activation functions applied after each layer except the final one. The output dimension of the encoder is set to 50. Separate encoders are used for the actor and critic networks, with the weights of their fully connected layers shared to ensure consistency in the learned representations. However, the shared weights are only updated by the critic, and gradients from the actor are stopped before propagating through the encoder.

The critic network employs double Q-learning, with each Q-function represented by a three-layer MLP featuring ReLU activation functions after each layer, except the final one. Similarly, the actor network, which is responsible for generating the mean and covariance for a diagonal Gaussian distribution representing the policy, follows the same structure. Both the actor and critic networks have a hidden layer dimension of 256. Additionally, the weight matrices of the fully connected layers in both networks are orthogonally initialized, and all biases are set to zero. The target critic network undergoes a soft update with a factor of $\tau =0.005$ every two iterations, while the actor network is also updated every two iterations.

When utilizing a predictive dynamics model, an additional network approximates a Gaussian distribution representing the transition probabilities. For simplicity, deterministic transitions were assumed in this environment. The forward dynamics model is a two-layer MLP with a hidden layer dimension of 512, with layer normalization applied after the first fully connected layer, followed by a ReLU activation function to introduce nonlinearity.

In our proposed approach, which incorporates a nonlinear projection, the projection layer consists of a single fully connected layer with an output dimension of 32, followed by a ReLU activation function to introduce nonlinearity.

All other hyperparameters used in the experiment are summarized in Table 1.

Table 1.

Hyperparameters for noisy sparse cartpole environment.

Hyperparameter	Value
Batch size	512
Discount factor ( $\gamma$ )	0.99
Discount factor ( $c$ )	0.99
Initialization steps	1000
Learning rate	0.001
Maximum intrinsic reward ( $R_{max}$ )	0.1
Replay buffer size	50,000
Scaling factor ( ${\eta }_r$ )	2.0
Training steps	50,000

Results and discussion

Our initial analysis concentrates on the boundedness of the state space. We demonstrate the feasibility of DeePS in the noise-free setting ( $N_m=0$ ), comparing the performance with and without the projection layer. We also evaluate the performance of a variant that normalizes the latent representation rather than using a nonlinear projection. Figure 3 shows that the norm of the representations for the model without projection grows rapidly and uncontrollably, ultimately leading to the model’s failure in learning an effective policy. Furthermore, we observe that the average norm of the representations, when using the nonlinear projection, remains well within the theoretical limits, unlike the normalization approach where the norm grows to the maximum allowed by the bounds. ³¹ This observation further validates the effectiveness of our solution in preventing unconstrained expansion of the latent space.

Figure 3.

The solid lines represent the average episode rewards, while the dashed lines indicate the $L_2$ norm values of the encoder representations. The shaded regions depict the 95% confidence intervals.

The performance of the different models (SAC, DBC, DBC w/ ir, DeePS, and DeePS w/ ir) are illustrated in Figure 4, and the average episode rewards are summarized in Table 2, with noise intensity incrementally raised from $N_m=0$ to $N_m=2$ . Overall, approaches utilizing intrinsic rewards perform the best as expected given the sparsity of the rewards. Among the different approaches, the action-based DeePS w/ ir showed the best performance.

Figure 4.

Performance comparison of SAC, DBC, and DeePS (ours) in the noisy sparse cartpole environment. ‘ir’ indicates the use of intrinsic rewards. Each plot represents the average episode reward over 10 episodes, averaged across 20 runs, with the 95% confidence interval indicated. SAC: Soft Actor-Critic; DBC: Deep Bisimulation for Control; DeePS: Deep Policy Similarities.

Table 2.

The average episode reward over 10 episodes for the noisy sparse cartpole environment, averaged across 20 runs, is presented with the standard deviation in brackets. ‘ir’ indicates the use of intrinsic rewards. The maximum attainable episode reward is 200.

	$N_m=0$	$N_m=1$	$N_m=2$
SAC	200.00 (0.00)	117.87 (64.87)	10.25 (2.80)
DBC	165.78 (58.96)	178.74 (52.73)	93.73 (98.62)
DBC w/ ir	197.75 (10.08)	197.15 (12.75)	194.22 (17.42)
DeePS	200.00 (0.00)	157.03 (46.74)	95.22 (93.59)
DeePS w/ ir	200.00 (0.00)	200.00 (0.00)	199.07 (3.29)
DeePS (norm)	200.00 (0.00)	54.10 (78.12)	16.38 (43.27)
DeePS (norm) w/ ir	200.00 (0.00)	200.00 (0.00)	199.85 (0.65)
DeePS (w/o projection)	34.20 (16.45)	38.54 (12.72)	36.02 (8.86)
DeePS (w/o projection) w/ ir	35.98 (17.73)	41.90 (40.64)	38.63 (4.81)

Bold values are the best performing algorithm for each scenario.

SAC: Soft Actor-Critic; DBC: Deep Bisimulation for Control; DeePS: Deep Policy Similarities.

The SAC, a representative of standard reinforcement learning methods, managed to achieve optimal performance in the noise-free setting despite the sparse rewards. However, its performance declined significantly as the noise intensity increased, eventually failing at $N_m=2$ . This underscores the limitation of standard reinforcement learning approaches in generalizing noisy observations as consistent latent representations.

Interestingly, the DBC approaches, both with and without intrinsic rewards, failed to achieve optimal performance even in the noise-free setting, indicating instability. This instability may stem from inaccurate state clustering where states with similar values but different optimal actions are aggregated incorrectly, as discussed previously. The declining performance of DBC, along with the high standard deviations, further supports the notion that state aggregation using $\pi$ -bisimulation is insufficient for effective generalization.

Although DeePS also experienced reduced performance as noise intensity increased (as seen in Table 2), the evaluation curves suggest that, unlike DBC which appeared to plateau, training was incomplete for DeePS, hinting that extended training could improve performance. Our approach showed delayed improvements during training, likely due to the difficulty of policy learning in sparse reward settings early in the training process, and state aggregation based on these random or poor policies hurt the training process. This points to a potential weakness in our approach in the reliance on the quality of the baseline policy. This is evident from the improved performance of DeePS w/ ir as seen in Figure 4 when intrinsic rewards were used to counteract reward sparsity. Despite this weakness, DeePS still outperforms the SAC, indicating that the model likely benefits from state aggregation based on policy-driven behaviors.

We conducted additional ablation studies to assess the impact of nonlinear projection on performance. As shown in Table 2, the model suffers catastrophic failure when no projection or normalization constraint is applied to the latent space. When intrinsic rewards are included, models utilizing nonlinear projection and normalization exhibit comparable performance. However, in the absence of intrinsic rewards, the normalized variant performs significantly worse, highlighting the robustness of the nonlinear projection approach.

In this experiment, our approach based on an approximation of the policy-driven PSM demonstrates significant promise in addressing the generalization challenge more effectively than standard reinforcement learning algorithms or the reward-based DBC. The superior performance of our approach validates the feasibility of using approximate forward dynamics to estimate the PSM, which is then used to improve generalization for reinforcement learning.

Noisy cartpole with randomized rewards

To study the impact of reward structures on the generalization ability of our action-based approach, we introduce slight variations to the noisy cartpole environment. We simulate a generalization problem where semantically equivalent environments have different reward structures by integrating randomized rewards. Specifically, the agent receives a reward that uniformly varies between 50% and 150% of the original reward for each successful timestep. This setup disregards reward sparsity to allow for greater variations in state values.

The randomization of reward signals introduces stochasticity, creating a more challenging and realistic environment for evaluating generalization. By testing the models in these varied environments, we can assess their ability to learn robust policies that generalize across different reward structures—essential for real-world applications where reward functions may be inconsistent. This allows us to compare the flexibility and adaptability of our proposed policy-driven DeePS approach to standard reinforcement learning algorithms and DBC.

Network architecture and hyperparameters

In general, we made use of the same network architecture and hyperparameters as with the previous experiment. However, to better manage the increased complexity introduced by randomized rewards, we made several adjustments to both the network architecture and training process. The hidden layer dimensions for both the actor and critic networks were reduced from 256 to 128 in an attempt to reduce the model’s capacity, preventing it from memorizing and overfitting to specific reward structures. In addition, we decay the learning rate by 0.99 every 1000 environment steps and increase the training batch size from 512 to 1024 to stabilize the training process. These adjustments help ensure that the models can handle the increased complexity and stochasticity of the noisy cartpole environment with randomized rewards, facilitating a thorough evaluation of the robustness and adaptability of our proposed method in a setting that closely simulates real-world scenarios.

Results and discussion

The performance of different models (SAC, DBC, and DeePS) in the noisy cartpole environment with randomized rewards is illustrated in the provided plots (Figure 5) and the table of results (Table 3). The results demonstrate the superiority of our proposed DeePS approach against the SAC and DBC across different noise levels ranging from $N_m=0$ to $N_m=20$ .

Figure 5.

Performance comparison of SAC, DBC, and DeePS (ours) in the noisy cartpole environment with randomized rewards. Each plot represents the average episode reward over 10 episodes, averaged across 20 runs, with the 95% confidence interval indicated. SAC: Soft Actor-Critic; DBC: Deep Bisimulation for Control; DeePS: Deep Policy Similarities.

Table 3.

The average episode reward over 10 episodes for the noisy cartpole environment with randomized rewards, averaged across 20 runs, is presented with the standard deviation in brackets. The maximum attainable episode reward is 200.

	$N_m=0$	$N_m=10$	$N_m=20$
SAC	190.29 (23.10)	186.73 (17.41)	115.94 (43.45)
DBC	199.73 (1.19)	170.68 (17.69)	99.41 (49.82)
DeePS	199.00 (3.53)	195.31 (8.56)	177.99 (17.21)

Bold values are the best performing algorithm for each scenario.

SAC: Soft Actor-Critic; DBC: Deep Bisimulation for Control; DeePS: Deep Policy Similarities.

In the noise-free setting ( $N_m=0$ ), both DeePS and DBC achieved near-optimal performance, with DBC slightly outperforming DeePS. However, as noise level increases from $N_m=0$ to $N_m=20$ , DeePS maintained strong performance, while both DBC and SAC showed clear signs of deterioration. At the highest noise level ( $N_m=20$ ), DeePS outperformed SAC and DBC by approximately 53.5% and 79.0%, respectively, in terms of episode rewards.

The evaluation curves in Figure 5 further support these findings, showing that DeePS consistently delivers higher and more stable performance as noise increases. The decline in SAC’s performance underscores the limitation of standard reinforcement learning in effectively generalizing across MDPs with variable reward structures. While DBC marginally outperformed DeePS in the absence of noise, its performance deteriorated more rapidly as noise levels increased, eventually performing worse than SAC at $N_m=20$ . DBC also exhibited greater instability, reflected in its larger standard deviation, particularly in high-noise environments. These results suggest that reward-based state aggregation methods not only negatively affect performance in environments with inconsistent reward structures, but also introduce instability.

Overall, these experiments highlight the effectiveness of DeePS in generalizing across continuous control tasks affected by noise and inconsistent rewards. Our method consistently outperforms both standard reinforcement learning algorithms and reward-based bisimulation approaches, demonstrating its robustness in handling complex and variable environments. By focusing on action-based similarities, DeePS enables more reliable policy transfer and adaptability, making it a valuable tool for real-world robotic applications. This success underscores the potential of our approach to enhance the development of more resilient and versatile reinforcement learning agents for actual robotic systems.

Crossing environment

We further validate our proposed approach in a simple environment that mimics robotic navigation problems in indoor settings. For this purpose, we utilize the simple crossing environment (SimpleCrossingS9N1-v0) implemented in the procedurally generated Minigrid environment. ¹⁹ Each environment consists of a $7\times 7$ grid world, where the agent starts in the top left corner of the room and must reach the green goal square in the opposite corner. To do so, the agent must navigate through a single crossing point in the wall that spans either horizontally or vertically across the room. At each step, the agent receives a partial observation of shape $(7,7,3)$ , describing the current environment in its forward direction. The agent’s action space consists of seven actions: turn left, turn right, move forward, and four other redundant actions not used in this environment. A discounted reward of $1-0.9(\mathtt{s}\mathtt{t}\mathtt{e}\mathtt{p}/\mathtt{m}\mathtt{a}\mathtt{x}\mathtt{\_}\mathtt{s}\mathtt{t}\mathtt{e}\mathtt{p})$ is given when the agent successfully reaches the goal, where $\mathtt{s}\mathtt{t}\mathtt{e}\mathtt{p}$ is the number of steps taken to reach the goal and $\mathtt{m}\mathtt{a}\mathtt{x}\mathtt{\_}\mathtt{s}\mathtt{t}\mathtt{e}\mathtt{p}$ is the maximum allowable steps, set at 324. If the agent fails to reach the goal within $\mathtt{m}\mathtt{a}\mathtt{x}\mathtt{\_}\mathtt{s}\mathtt{t}\mathtt{e}\mathtt{p}$ , it receives a reward of 0. This environment challenges the agent to locate the crossing in the wall and move toward the goal as quickly as possible to maximize rewards.

We study generalization by training the agent on 10 different layouts (training set) and testing it on three other semantically equivalent layouts (testing set) as shown in Figures 6 and 7, respectively. The layouts in the training and testing sets do not overlap, requiring the agent to adapt to the test environments after training for 1,500,000 timesteps in the training set. This presents a very challenging task known as zero-shot generalization.

Figure 6.

Environment layouts used to train the agent in the crossing environment. The agent should move through the crossing in the wall towards the goal located at the bottom right corner as quickly as possible.

Figure 7.

Environment layouts used to test the agents in the crossing environment. The agent should move through the crossing in the wall towards the goal located at the bottom right corner as quickly as possible.

We evaluate the performance of our proposed approach alongside standard reinforcement learning and also DBC. Additionally, we consider the use of the RAPID algorithm, ⁴⁴ designed to enhance reinforcement learning in procedurally generated environments. While the RAPID algorithm is not specifically designed for generalization, it is state-of-the-art for improving agent performance in procedurally generated environments. The RAPID algorithm uses a ranking buffer to gather high-quality state-action transitions for imitation learning. These transitions are ranked based on the weighted sum of their extrinsic rewards and their local and global exploration scores, defined as

$$S=w_0{S}_{\text{ext}}+w_1{S}_{\text{local}}+w_2{S}_{\text{global}}$$

(13)

where $S_{\text{ext}}$ is the extrinsic reward, $S_{\text{local}}$ is the local exploration score measuring episode diversity, $S_{\text{global}}$ is the global count-based exploration score with hyperparameters $w_0$ , $w_1$ and $w_2$ , respectively. All models are trained with Proximal Policy Optimization (PPO). ³⁸

As noted in previous experiments, DeePS is dependent on the quality of the baseline policy during training. Therefore, to ensure effective generalization, we first train the DeePS model with an additional RAPID loss for 1,000,000 timesteps to establish a strong baseline policy (for the training domain) before applying state aggregation for the remaining 500,000 timesteps.

Network architecture and hyperparameters

The encoder network consists of a two-layer MLP with Tanh activation function applied after the initial layer. The output dimension of the encoder is set to 64. Both the actor and the critic are single fully connected layers. For simplicity, transitions are assumed to be deterministic, and the forward dynamics model is a two-layer MLP with hidden dimensions of 512. Layer normalization is applied after the initial layer, followed by a ReLU activation function to introduce nonlinearity. The projection layer is a single fully connected layer with a dimension of 32, followed by a ReLU activation for nonlinearity. Gradient clipping with a norm value of 1.0 is applied for the PPO loss. The hyperparameters used are summarized in Table 4.

Table 4.

Hyperparameters for crossing environment.

Hyperparameter	Value
Learning rate	0.0001
Training steps	1,500,000
Clip	0.2
Discount ( $\gamma$ )	0.99
Epochs	4
Entropy coefficient	0.01
Generalized advantage estimation ( $\lambda$ )	0.95
Horizon	128
Minibatch	4
Batch size	256
Buffer size	10,000
Extrinsic weight ( $w_0$ )	1.0
Local exploration weight ( $w_1$ )	0.1
Global exploration weight ( $w_2$ )	0.001

Results and discussion

The performance of the different models (PPO, RAPID, DBC, and DeePS) is shown in Figure 8, and the summary of the results provided in Table 5.

Figure 8.

Comparison of training (top) and testing (bottom) performance for PPO, RAPID, DBC, and DeePS (ours) in the crossing environment. The training plots represent the average episode rewards over 100 episodes, and the testing plots represent the average episode rewards over 30 episodes. Both are averaged across five runs, with 95% confidence intervals included. PPO: Proximal Policy Optimization; DBC: Deep Bisimulation for Control; DeePS: Deep Policy Similarities.

Table 5.

The average episode reward over 100 episodes (train) and 30 episodes (test) for the crossing environment, averaged across five runs, is presented with the standard deviation in brackets.

	train	test
PPO	0.488 (0.454)	0.224 (0.365)
RAPID	0.864 (0.082)	0.202 (0.255)
DBC	0.929 (0.014)	0.224 (0.069)
DeePS	0.895 (0.091)	0.403 (0.104)

Bold values are the best performing algorithm for each scenario.

PPO: Proximal Policy Optimization; DBC: Deep Bisimulation for Control; DeePS: Deep Policy Similarities.

We observed that PPO, as a representative of standard reinforcement learning models, performed poorly on both the training and testing sets. Although extended training could potentially enhance PPO’s performance on the training set, its testing performance is unlikely to improve due to its tendency to overfit to the training conditions.

The other models demonstrated strong training performance, as shown in Table 5, with DBC achieving the highest average training reward. Surprisingly, DBC performed well in the test environments during the early stages of training. However, its near-optimal training performance did not translate into effective generalization as the training progressed. This is evident in its test performance, which remained comparable to PPO despite the superior training results. The inefficiency of DBC is further illustrated in the evaluation curve in Figure 8, where its test performance deteriorates sharply after plateauing in training, underscoring its susceptibility to overfitting.

Similarly, while RAPID achieved better training performance than the baseline PPO, this improvement did not generalize. In fact, RAPID’s test performance was observed to be worse than PPO and clear signs of overfitting were apparent in its evaluation curve as seen in Figure 8. This further highlights the impracticality of using RAPID for generalization purposes.

In contrast to the other models, DeePS demonstrated substantial improvements in generalization, achieving approximately 99.5% higher test rewards than RAPID and 79.9% higher test rewards than DBC (and PPO). Despite achieving slightly lower training performance compared to DBC, DeePS’s performance on the test environments was the best among all models, with a test reward of 0.403, nearly doubling that of PPO and DBC. This result highlights the ability of DeePS to strike a balance between strong training performance and robust generalization.

As seen in the evaluation curves, DeePS effectively leverages its state aggregation mechanism to enhance both training and testing performance. Unlike RAPID and DBC, DeePS avoids overfitting, enabling it to perform consistently across unseen environments. By focusing on action-based similarities, DeePS offers a reliable method for policy transfer and adaptation, ensuring that agents remain effective across diverse environments.

Although DeePS’s test performance remains below optimal levels, it represents a significant step forward in reinforcement learning methodologies. By improving generalization without compromising training performance, DeePS demonstrates its potential to enhance the robustness and adaptability of reinforcement learning policies, particularly in challenging and varied environments.

Limitations

While the proposed DeePS method is effective in enhancing the generalization of reinforcement learning policies, its reliance on a quality baseline and certain assumptions about environmental transitions limit its overall effectiveness. Specifically, the assumption of a Gaussian transition probability may not hold across all environments, and the requirement of accurate modeling of the transition dynamics may limit the reliability of the metric estimates. Nonetheless, a significant amount of existing works have shown the possibility of accurately modeling environmental transitions.^45–48 which supports the use of our proposed methodology.

Conclusion

In conclusion, this work introduced a learning-based bisimulation approach that leverages policy similarities. Our experiments demonstrated the effectiveness and robustness of the proposed methodology, showing that DeePS outperforms the state-of-the-art reward-based bisimulation methods. Furthermore, the approach manages to preserve the integrity of the training process while substantially improving the model’s ability to generalize to new, unseen environments with semantically equivalent dynamics. Even in environments with different reward structures, DeePS was able to generalize effectively. This makes it particularly valuable for real-world robotic applications, where adaptability and robustness are essential.