CGGT: Concept-Guided Grounded Training for Generalizable and Interpretable Reinforcement Learning in Robot Control

Abstract

Visuomotor reinforcement learning (RL) has shown a strong capability in enabling robotic agents to perceive their environments and generate actions to accomplish downstream manipulation tasks. However, existing approaches often do shortcut learning, resulting in poor performance under outof-distribution (OOD) conditions, while also inducing incomprehensible interpretations. To address these issues, we propose Concept-Guided Grounded Training (CGGT), a framework that leverages task-specific concepts from the environment to improve visuomotor learning, not only in terms of success rates but also in terms of generalization and interpretability. Specifically, our CGGT framework employs a single stage training scheme that integrates visual concept alignment at the feature extraction and spatial concept alignment at the decisionmaking network, along with the reward optimization. This approach enables robots to focus on action-relevant features and subsequently generate actions for task completion. Through our evaluations on eleven robotsuite tasks with multiple OOD scenarios in the RL-Vigen benchmark, CGGT outperforms state-of-the-art baselines in terms of generalization while also enabling failure inspection through concept visualizations. Finally, we also deploy our proposed method on the Franka Emika robot in four different tasks with real-world objects to test its feasibility in the real-world domain.