TraKDis: A Transformer-based Knowledge Distillation Approach for Visual Reinforcement Learning with Application to Cloth Manipulation

TL;DR

TraKDis tackles cloth manipulation with visual reinforcement learning by bridging vision and privileged state information through a two-stage Transformer-based knowledge distillation workflow. A privileged Decision Transformer is first trained with full cloth-state data, then a vision-based student is learned via a pre-trained state-estimation encoder and weight initialization to imitate the teacher, enabling robust, temporally aware control from RGB inputs. The approach achieves superior performance over SOTA baselines in SoftGym tasks, demonstrates robustness to estimation noise, and proves feasibility on a real UR5 robot. This work advances practical vision-based cloth manipulation by leveraging temporal transformers and cross-domain knowledge transfer to overcome observation gaps and dynamic cloth behavior.

Abstract

Approaching robotic cloth manipulation using reinforcement learning based on visual feedback is appealing as robot perception and control can be learned simultaneously. However, major challenges result due to the intricate dynamics of cloth and the high dimensionality of the corresponding states, what shadows the practicality of the idea. To tackle these issues, we propose TraKDis, a novel Transformer-based Knowledge Distillation approach that decomposes the visual reinforcement learning problem into two distinct stages. In the first stage, a privileged agent is trained, which possesses complete knowledge of the cloth state information. This privileged agent acts as a teacher, providing valuable guidance and training signals for subsequent stages. The second stage involves a knowledge distillation procedure, where the knowledge acquired by the privileged agent is transferred to a vision-based agent by leveraging pre-trained state estimation and weight initialization. TraKDis demonstrates better performance when compared to state-of-the-art RL techniques, showing a higher performance of 21.9%, 13.8%, and 8.3% in cloth folding tasks in simulation. Furthermore, to validate robustness, we evaluate the agent in a noisy environment; the results indicate its ability to handle and adapt to environmental uncertainties effectively. Real robot experiments are also conducted to showcase the efficiency of our method in real-world scenarios.

Figures (9)

• Figure 1: Pipeline of TraKDis for training visual RL agents for cloth manipulation. A privilege agent is trained to imitate the expert policy with privileged state information (cloth particle locations) to obtain a robust performance. Then, a student agent that can only access partial observation (RGB images) is trained to imitate the privileged agent. By leveraging a CNN encoder and weight copy, the student policy can realize an enhanced performance with limited observation information.
• Figure 2: Model Training of the privileged agent. We first conduct offline training using the expert data collected from the human-designed heuristic algorithm. An online fine-tuning is then applied for the trained model to obtain a better performance
• Figure 3: CNN Model: We design a CNN encoder to estimate the state information from image observation. Image augmentation is applied to improve the robustness of state estimation.
• Figure 4: The figure demonstrates the overview of our knowledge distillation procedure. We run the teacher and student policy simultaneously for knowledge distillation. The teacher policy has access to all the state dynamics of the cloth and thus performs better. The student policy, which receives image inputs, is trained to imitate the actions of the teacher policy. By using a pre-trained CNN, the image can be estimated to encoded input $s_{t}$. Since the student agent and privileged agent have the same architecture, we initialize the weight of the student policy by copying the weights from the teacher's policy, which aids the knowledge distillation process. The parameters of CNN and teacher policy are frozen during training. Only the student policy is updated via the imitation loss.
• Figure 5: Snapshots of the proposed method. We evaluate our proposed methods by adopting three cloth manipulation tasks. Only image observation is used for the agent to generate trajectory. More demonstration can be found in the multi-media resource