CtrlFormer: Learning Transferable State Representation for Visual Control via Transformer
Transformer has achieved great successes in learning vision and language
representation, which is general across various downstream tasks. In visual
control, learning transferable state representation that can transfer between
different control tasks is important to reduce the training sample size.
However, porting Transformer to sample-efficient visual control remains a
challenging and unsolved problem. To this end, we propose a novel Control
Transformer (CtrlFormer), possessing many appealing benefits that prior arts do
not have. Firstly, CtrlFormer jointly learns self-attention mechanisms between
visual tokens and policy tokens among different control tasks, where multitask
representation can be learned and transferred without catastrophic forgetting.
Secondly, we carefully design a contrastive reinforcement learning paradigm to
train CtrlFormer, enabling it to achieve high sample efficiency, which is
important in control problems. For example, in the DMControl benchmark, unlike
recent advanced methods that failed by producing a zero score in the "Cartpole"
task after transfer learning with 100k samples, CtrlFormer can achieve a
state-of-the-art score with only 100k samples while maintaining the performance
of previous tasks. The code and models are released in our project homepage.