TEA: Trajectory Encoding Augmentation for Robust and Transferable Policies in Offline Reinforcement Learning

TL;DR

This paper introduces TEA, a trajectory encoding augmentation method that enhances offline reinforcement learning policies' ability to generalize across unseen environment dynamics by integrating latent environment representations.

Contribution

The paper proposes TEA, a novel approach that extends state space with environment dynamics encodings, improving policy transferability in offline RL.

Findings

TEA improves policy transfer to new environments with unseen dynamics.

Incorporating environment encodings enhances generalization over state-only methods.

TEA captures environment-specific features critical for robust policy performance.

Abstract

In this paper, we investigate offline reinforcement learning (RL) with the goal of training a single robust policy that generalizes effectively across environments with unseen dynamics. We propose a novel approach, Trajectory Encoding Augmentation (TEA), which extends the state space by integrating latent representations of environmental dynamics obtained from sequence encoders, such as AutoEncoders. Our findings show that incorporating these encodings with TEA improves the transferability of a single policy to novel environments with new dynamics, surpassing methods that rely solely on unmodified states. These results indicate that TEA captures critical, environment-specific characteristics, enabling RL agents to generalize effectively across dynamic conditions.