Scene Graphs for Robotic Imitation Learning

Introduction

Imitation learning enables robots to acquire complex skills from demonstrations, but performance often degrades when tasks span large environments or long time horizons. A core challenge is partial observability: onboard sensors rarely capture the full scene at once.

We propose using task-relevant scene graphs as an explicit, structured memory for imitation learning. Each object is represented as a node with visual embeddings, 2D bounding boxes when visible, and persistent 3D centroids when out of view. This compact representation lets a diffusion policy reason over spatial and temporal context accumulated throughout an episode.

Experiments on simulated mobile manipulation with Boston Dynamics Spot and real-world tabletop manipulation with a Franka arm show substantial gains over RGB and object-centric baselines, especially on long-horizon tasks that require remembering objects outside the current field of view.

Method

Our pipeline builds and updates a task-driven scene graph online, then conditions a transformer-based diffusion policy on the resulting node sequence.

Task-driven scene graph pipeline and policy architecture — We maintain a dynamic, task-relevant scene graph and use it to condition a transformer-based diffusion policy. Visible objects update their 2D boxes; out-of-view objects retain 3D centroids for memory-efficient reasoning.

Task-relevant object enumeration

Given a natural-language task description, an LLM extracts a list of relevant object nouns (e.g., microwave, cup for “heat up tea”).

Scene graph construction & maintenance

Grounding DINO detects task-relevant objects in the first frame. Each node stores a DINO-v2 embedding, 2D bounding box, and 3D centroid from depth backprojection. XMem tracking preserves identity over time; new detections are merged when they do not overlap existing nodes.

Scene graph-conditioned diffusion policy

Nodes are encoded with MLPs over geometry and processed by a transformer encoder together with a scene-level root token. The resulting representation conditions a diffusion policy (DDPM training, DDIM inference) that outputs short-horizon action sequences.

Partially observed tabletop setup — Real-world evaluation uses only a wrist-mounted ZED Mini camera, forcing policies to rely on memory accumulated over the trajectory.

Demo

Rollout videos from simulated Spot experiments and real-world Franka tabletop manipulation.

Simulated Mobile Manipulation (Spot)

Heat-up Tea

Heat-up Tea Long

Real-World Tabletop Manipulation (Franka)

Pineapple-Bowl

Feed-Animals

Feed-Giraffe

Clean-Table