Reward Machine Inference for Robotic Manipulation

TL;DR

The paper tackles the challenge of long-horizon robotic manipulation with sparse rewards by learning Reward Machines directly from visual demonstrations. It proposes a four-step RM inference pipeline that discovers sub-goals via clustering of visual features and constructs an RM with a labeling function $L$, enabling potential-based reward shaping to provide dense guidance for RL. An actor-critic-like Deep Q-Network is then trained per RM state using two FCNs for pick and place actions, demonstrating that the inferred RM closely matches a ground-truth RM and supports learning an optimal policy with placement accuracy comparable to using true object positions. Overall, the approach reduces the need for manually designed high-level task representations and yields interpretable RM structures that improve sample efficiency and policy quality in vision-based robotic manipulation.

Abstract

Learning from Demonstrations (LfD) and Reinforcement Learning (RL) have enabled robot agents to accomplish complex tasks. Reward Machines (RMs) enhance RL's capability to train policies over extended time horizons by structuring high-level task information. In this work, we introduce a novel LfD approach for learning RMs directly from visual demonstrations of robotic manipulation tasks. Unlike previous methods, our approach requires no predefined propositions or prior knowledge of the underlying sparse reward signals. Instead, it jointly learns the RM structure and identifies key high-level events that drive transitions between RM states. We validate our method on vision-based manipulation tasks, showing that the inferred RM accurately captures task structure and enables an RL agent to effectively learn an optimal policy.

Figures (5)

• Figure 1: Overview of the proposed method applied to the task of building a predefined pyramid. (1) Visual demonstrations are captured, and feature embeddings are extracted using a pre-trained model $\phi$. (2) Sub-goals are inferred by clustering the feature vectors obtained from the demonstrations. (3) An RM is constructed, capturing the valid temporal ordering and transitions between the inferred sub-goals. The $0$-prototype corresponds with the initial RM state.
• Figure 2: Overview of our DQN formulation. Each action consists of a picking operation at pixel coordinates $(u_{\textrm{pick}},v_{\textrm{pick}})$ followed by a placing operation at $(u_{\textrm{place}}, v_{\textrm{place}})$. The Q-function $Q_{\theta}(s,a)$ is modeled using two FCN's: $\psi_{\textrm{pick}}$ and $\psi_{\textrm{place}}$ which generate pixel-wise Q-value maps for their respective actions.
• Figure 3: Inferred RM for the Place-3 task (left) and the corresponding state prototypes (right). We omitted the self-loop transitions for clarity.
• Figure 4: Inferred RM for the Stack-3 task (left) and the corresponding state prototypes (right). We omitted the self-loop transitions for clarity.
• Figure 5: Total reward obtained during one episode (top) and average distance between each object and its ground goal location, i.e. placement error (bottom) during training. Results are averaged over 10 runs, the shaded region represents the standard deviation.