Research map / Multi-agent systems and games / Human-AI and goal-conditioned

Data Augmentation for Instruction Following Policies via Trajectory Segmentation

AAAI 2025 · Multi-agent systems and games › Human-AI and goal-conditioned · approach: Embodied agent learning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Play Segmentation (PS), a probabilistic segmentation model trained only on single annotated instruction segments, to extract labelled segments from unannotated play trajectories and augment instruction-following policy training, compared against adapted UnLoc and TriDet video segmentation methods.

Problem

Instruction-following agents in robotics and games lack instruction-trajectory data; large unannotated play trajectories exist but extracting correctly labelled instruction segments from them is hard, since random segments often cut across instructions.

Approach

A conditional probabilistic segmentation model over play trajectories, trained with negative segments created by translating and elongating/shortening annotated segments, and decoded with a dynamic-programming (Viterbi-like) maximum-likelihood segmentation, so that it generalises from short labelled segments to long play trajectories.

Key results

• On BabyAI, PS-augmented data reaches 0.702 task accuracy versus 0.544 for the 10% starting split and 0.645 for 25% GT, exceeding the 25% GT level, while TriDet drops to 0.261 and random relabelled segments to 0.053.
• PS achieves segmentation precision 0.826, recall 0.627, and label accuracy 0.516 on BabyAI, compared with 0.283 precision for UnLoc and 0.338 for TriDet.
• On CALVIN, PS outperforms the best video segmentation model and reaches a higher average sequence length (1.477) than the 50% GT subset (1.411).

Measured effects reported in the paper

Surprising findings

• Randomly sampled segments severely degrade policy performance on BabyAI (0.053 with random relabelling), and on CALVIN they leave performance unchanged.
• Adapted video segmentation methods (UnLoc, TriDet) have little or negative effect on BabyAI policy performance because extracted segments are too short or inaccurate.

Limitations

• PS segmentation has cubic complexity in trajectory length, so full play trajectories must be segmented in windows.
• Extracted labelled segments are imbalanced across tasks on CALVIN, with some tasks underrepresented, which explains the remaining gap to ground-truth relabelling.
• Without ground-truth annotations for CALVIN play trajectories, segmentation quality cannot be directly measured.

Experimental setup

Datasets: BabyAI GoTo play trajectories (generated) CALVIN

Benchmarks: BabyAI GoTo CALVIN

Metrics: task accuracy (percentage of tasks solved within 25 timestep average sequence length (instructions completed before first segmentation precision segmentation recall F1 label accuracy

Baselines: UnLoc TriDet Random segments GT relabel GT 50%/25%/10% subsets

Builds on: UnLoc TriDet MCIL I3D Viterbi algorithm for hidden semi-Markov models

Closest papers by meaning

• Unsupervised Video Object Segmentation for Deep Reinforcement Learning (NeurIPS 2018 · similarity 0.88)
• Unsupervised Hierarchical Skill Discovery (ICML 2026 · similarity 0.87)
• CompILE: Compositional Imitation Learning and Execution (ICML 2019 · similarity 0.87)
• Towards Automatic Learning of Procedures From Web Instructional Videos (AAAI 2018 · similarity 0.87)
• GO-DICE: Goal-Conditioned Option-Aware Offline Imitation Learning via Stationary Distribution Correction Estimation (AAAI 2024 · similarity 0.86)
• CCIL: Continuity-Based Data Augmentation for Corrective Imitation Learning (ICLR 2024 · similarity 0.85)