TLDR: A new research paper introduces Supervised Contrastive Imitation Learning (SCIL), a method that improves how AI agents learn to play video games from visual inputs. By structuring the agent’s internal representations based on action similarity, SCIL enables faster learning, better generalization, and enhanced performance in both 2D Atari and complex 3D games like Astro Bot and Returnal. The approach discretizes continuous actions and uses action labels to guide the learning process, avoiding problematic data augmentations.
Training artificial intelligence (AI) agents to play video games, especially from raw visual input, presents significant challenges. Traditional methods often struggle with high-dimensional visual data, leading to longer training times and a risk of overfitting. A new research paper introduces a novel approach called Supervised Contrastive Imitation Learning (SCIL) to address these issues by focusing on learning more effective state representations for game agents.
Imitation Learning (IL) is a technique where AI agents learn by observing demonstrations, much like a human learning by watching an expert. However, when agents are limited to visual inputs (like pixels on a screen) rather than detailed internal game states, it becomes harder for them to understand the crucial cause-effect relationships between what they see and the actions they should take. The goal is to create a ‘latent representation’ – a simplified, meaningful summary of the visual input – that highlights action-relevant factors, such as “the player jumps whenever an obstacle appears ahead.”
The core idea behind SCIL is to structure this latent space so that observations leading to similar actions are grouped closely together, while observations leading to different actions are kept separate. This is achieved by integrating a Supervised Contrastive (SupCon) loss function into the IL training process. Unlike some self-supervised learning methods that rely on artificial data augmentations (like rotating or shifting images), SCIL uses the actual action labels to define what constitutes a ‘positive pair’ (observations associated with similar actions) and a ‘negative pair’ (observations associated with different actions). This is crucial in video games where precise spatial information is vital and geometric augmentations could distort critical cues.
A key innovation in SCIL is its ability to handle continuous action spaces, which are common in many video games (e.g., joystick movements, aiming angles). Since the original SupCon framework is designed for discrete classification tasks, the researchers developed a method to discretize continuous action dimensions into a set of ‘bins’. This allows them to treat nearby continuous values as equivalent, effectively assigning them to the same ‘class’ for the purpose of computing the SupCon loss. These discretized values are then combined using a positional encoding scheme to create a unique categorical label for each combination of actions, enabling the contrastive learning process to work effectively.
The researchers also addressed a practical challenge: the possibility of a mini-batch of training data not containing any positive pairs for a given sample, which could lead to errors. They implemented safeguards to handle these edge cases, recommending the use of sufficiently large mini-batches to mitigate this issue.
Experiments were conducted on a variety of games, including the 3D titles Astro Bot and Returnal, and several 2D Atari games like Ms. Pac-Man, Montezuma’s Revenge, and Space Invaders. The results were consistently positive. For the Atari games, SCIL showed significant percentage improvements in scores, with Ms. Pac-Man seeing an average increase of 33.47% and Montezuma’s Revenge 26.53%. In 3D games, agents trained with SCIL demonstrated higher success rates in Astro Bot’s challenging levels and inflicted more damage on bosses in Returnal, finishing the first phase of a boss fight in 37.5% of trials compared to the baseline’s 7.5%.
Beyond improved performance, SCIL also led to faster learning convergence and better generalization, meaning the agents could adapt more effectively to previously unseen game states. The method proved to be architecture-agnostic, working well with different underlying model structures. The findings strongly support the central hypothesis: by enforcing a structure in the latent space where observations corresponding to similar actions are represented by similar embeddings, IL agents can learn more efficiently and generalize better.
Also Read:
- Information Theory Unlocks Deeper Understanding and Diagnosis of Reinforcement Learning Agents
- Rethinking Actor-Critic Design for Discrete Action Reinforcement Learning
This work, detailed further in the paper Learning Representations in Video Game Agents with Supervised Contrastive Imitation Learning, suggests that the insights gained could extend beyond video game applications, potentially benefiting imitation learning agents in other complex domains such as robotics.
