Reward machine

Overview

A reward machine is a small graph that sits beside the agent's state. Each node is a stage of the task and each edge fires when a proposition becomes true, such as "block 1 is on the table" or "stack has two blocks". The reward depends on which edge fired, so a long task with a sparse final reward gets intermediate signals 2412.10096 §1 2412.10096 §3.2.

The usual setting assumes two things are given: the machine's states and transitions, written by an expert, and the propositions, with a labelling function that maps each environment state to the propositions that hold 2412.10096 §1 2412.10096 §2. The core idea of the inference work is to remove both assumptions for vision-based manipulation. The propositions and the machine are recovered from demonstrations, and the agent then learns with the recovered machine 2412.10096 §1 2412.10096 §4.

The invariant that variants share is the reward structure. Shaping that depends only on the machine's state keeps the optimal policy unchanged, as the paper states for its potential-based term 2412.10096 §4.4.

Card: Reward Machine Inference for Robotic Manipulation (arXiv:2412.10096).

Variants

Expert-written machine with known propositions

The machine and the labelling function are supplied by hand, and the agent is trained on them. The paper uses this setting as its comparison point: a DQN agent with the ground-truth machine and propositions, following an earlier vision-based setup from the literature 2412.10096 §6. The paper does not publish the cost of writing the machine by hand, so the effort saved by the inference method is not quantified there 2412.10096 §6.

Inferred machine from unlabelled demonstrations

Demonstrations are recorded as video frames without labels. The pipeline is:

1. Extract a feature vector per frame with a pretrained ResNet-50, after cropping the frame to the tabletop 2412.10096 §4.2.
2. Cluster all features with DBSCAN. Each cluster is a candidate sub-goal, on the assumption that sub-goals are visited more often than other states 2412.10096 §4.3.
3. Take the frame closest to each cluster mean as its prototype, and make one proposition per prototype. The proposition is true when the feature distance to the prototype is below a threshold κ 2412.10096 §4.3 2412.10096 §4.4.
4. Build the transition table from the sequence of true propositions in each demonstration, keeping the first transition seen for each state and proposition pair 2412.10096 alg. 1 2412.10096 §4.4.

This variant removes the need for hand-written propositions. It adds three per-task knobs, κ, the DBSCAN radius and the minimum cluster size, which the paper tunes for each scenario without stating the procedure 2412.10096 §6 2412.10096 tab. 1. Its results are reported for five simulated tasks only 2412.10096 §6.

When to use which

Implementation sketch

Language-agnostic pseudocode of the inference core, with notes on the variant differences. Thresholds and cluster settings are per task.

# Infer propositions and transitions from unlabelled demonstrations
# Input:  demos: list of lists of feature vectors, one list per demonstration
# Output: prototypes, a transition table keyed by (state, proposition)

def infer_machine(demos, eps, min_points, kappa):
    # Cluster all frames together; DBSCAN labels noise as -1  [arxiv:2412.10096 §4.3]
    feats = [f for demo in demos for f in demo]
    labels = dbscan(feats, eps, min_points)
    # One cluster per sub-goal; the prototype is the frame nearest its mean
    # [arxiv:2412.10096 eq. 3] [arxiv:2412.10096 eq. 4]
    prototypes = {c: nearest_to_mean(feats, labels, c) for c in set(labels) - {-1}}
    # Proposition p_c holds when the feature is within kappa of its prototype
    # [arxiv:2412.10096 eq. 5]
    def true_props(f):
        return {c for c, p in prototypes.items() if dist(f, p) < kappa}
    # Walk each demonstration and record the first transition per pair
    # [arxiv:2412.10096 alg. 1]
    seq = [[true_props(f) for f in demo] for demo in demos]
    return prototypes, infer_transitions(seq, initial_state=0, state_of=lambda c: c)

With a known machine, skip the first block and pass the expert propositions straight into the same shaping and training code. The shaping term is the same in both variants:

# Potential-based shaping on the machine state  [arxiv:2412.10096 eq. 6] [arxiv:2412.10096 eq. 7]
# Potential of a machine state: gamma raised to its edge distance to the goal
def shaped_reward(sparse_r, u, u_next, gamma, d_goal):
    phi = gamma ** d_goal(u)
    phi_next = gamma ** d_goal(u_next)
    return sparse_r + gamma * phi_next - phi

The agent trains one Q-function per machine state, each with its own replay buffer that also holds the demonstrations 2412.10096 §5. Set the discount γ explicitly: the paper does not state its value 2412.10096 §4.4 2412.10096 §6.

Contradictions

• Demonstration source. The evaluation section describes its demonstrations both as unstructured human demonstrations and as output from an algorithmic expert policy, in consecutive sentences 2412.10096 §6. The method is therefore tested only with scripted demonstrations, as far as the text shows.
• Inferred versus ground-truth machine. The paper reports similar learning curves for the two agents on four tasks, and a plateau for the ground-truth agent that is higher on Stack-3 2412.10096 §6 2412.10096 fig. 5. On Place-3, the inferred machine gets a lower placement error, which the paper attributes to two failed runs of the ground-truth agent rather than to a better machine 2412.10096 §6. Read these as author-reported and noisy: ten runs per task, shown only as curves.

Papers

• Reward Machine Inference for Robotic Manipulation: infers a reward machine, its propositions and transitions from unlabelled RGB demonstrations, and trains a DQN with potential-based shaping on simulated block tasks.