Action-driven contrastive representation for reinforcement learning

Abstract

In reinforcement learning, reward-driven feature learning directly from high-dimensional images faces two challenges: sample-efficiency for solving control tasks and generalization to unseen observations. In prior works, these issues have been addressed through learning representation from pixel inputs. However, their representation faced the limitations of being vulnerable to the high diversity inherent in environments or not taking the characteristics for solving control tasks. To attenuate these phenomena, we propose the novel contrastive representation method, Action-Driven Auxiliary Task (ADAT), which forces a representation to concentrate on essential features for deciding actions and ignore control-irrelevant details. In the augmented state-action dictionary of ADAT, the agent learns representation to maximize agreement between observations sharing the same actions. The proposed method significantly outperforms model-free and model-based algorithms in the Atari and OpenAI ProcGen, widely used benchmarks for sample-efficiency and generalization.

Fig 1. Reconstruction of auto-encoder based representation.

This is the Bigfish of ProcGen environment, which has various wallpaper patterns in the background. Under the auxiliary task of reconstructing original inputs, left images are pixel-inputs and right are reconstructions of their representations. It shows a problem that if there are various control-irrelevant details in wallpaper pixels, the encoder first learns abstraction about irrelevant details for gameplay instead of the essential pixels (i.e., fish in this example). Our proposal was motivated by this issue.

Fig 2. The overall framework of ADAT.

State-action pairs are sampled from a history of interactions. The dictionary is built with states by action types, labeling keys based on the query’s action. Through augmentation ‘Random-Translate’ [22], query is encoded for reinforcement learning and query-key pairs are encoded for contrastive learning [20]. Only the query encoder learns from contrastive loss, and the key encoder is trained as Momentum Update [19].

Fig 3. The sampling-efficacy following by two methods.

By extracting twice the amount and sorting them, all elements are used equally for learning. Unbiased Sampling prevents unwise early records from being leveraged much more than rich labels.

Fig 4. The frame-stacked inputs.

For Atari agents, four sequentially frame-stacked observations are translated from 84x84 to 92x92 pixels with zero pads. Likewise, for ProcGen agents, two frame-stacked observations are translated from 64x64 to 72x72 pixels [22].

Fig 5. Comparison of generalization capability.

Bigfish and Plunder Games in OpenAI Progen. 200-levels trained agents were evaluated on 100,000 levels. On-policy ADAT leveraged the samples rolled out by PPO policy. A difference in both frameworks is contrastive representation learning through ADAT. This figure reports the average score and standard deviation over five random seeds.

Fig 6. Saliency map of Bigfish and Plunder Games in OpenAI Progen.

From left, original rendering, saliency map of the PPO, and the PPO with ADAT in order.