Incorporating Task Progress Knowledge for Subgoal Generation in Robotic Manipulation through Image Edits

In plain words

A robot that follows a sentence such as "open the drawer" has to do several small steps in the right order. Some steps are easy to miss, and some are easy to repeat. This paper has the robot imagine the picture it should reach next, then move its gripper until the camera view looks like that picture, then imagine the next one.

The imagined pictures come from an image generator that is trained on short demonstration clips. The paper picks the key moments from each clip by looking at how fast the scene changes relative to the final frame. A small memory component keeps track of the moments already reached, so the next picture follows on from the earlier ones. The robot checks its own progress by comparing the current view with the imagined picture. When the two look similar enough, or when too many steps have passed, it imagines the next picture.

The authors compare this with a method that imagines pictures at fixed time intervals. Their claim is that picking moments by task progress gives better guidance, especially when the robot starts from different poses.

TL;DR

• TaKSIE trains a latent-diffusion-style image generator to produce the next keyframe image from the current view, the instruction and a GRU progress embedding §3.4.1 §3.3.2.
• Training subgoals are selected from demonstrations by a smoothed progress curve, not at fixed intervals §3.3.1.
• At run time, a time-contrastive progress check triggers regeneration §3.3.3 §4.2.3.
• Author-reported CALVIN LH-MTLC average length 3.18, above SuSIE (2.80) and LCD (2.88) tab. 2.
• The released code differs from the paper in several settings, including the feature used by the progress encoder and the sampler. Read the Paper vs code section before reimplementing.

Problem

• Fixed-interval subgoal generation (as in SuSIE) produces too many subgoals for a short interval, or misses key steps for a long one §1 §3.3.1.
• Without knowledge of task progress, generated subgoals can be inconsistent when several solutions exist, so the robot switches between modes §3.3.2.
• Language-conditioned policies that use an LLM or an API need text descriptions or action APIs; the paper targets visual goals instead §2.
• Assumptions made by the method: a single static RGB camera as observation §4.2.1; demonstrations that include a language annotation and a time-contrastive visual representation that tracks task progress §3.3; a separately trained low-level policy learned from the same demonstrations §3.5.
• Scope in the paper: fully observable, mostly atomic, sequential tasks; mobile manipulation and partial observation are future work §5.

Method

Problem formulation

• The subgoal generator G_θ produces the next visual subgoal f_gen from the initial observation, the instruction l*, and the progress so far s_0 … s_t §3.1.
• Training minimises a distance between generated subgoals and ground-truth subgoals τ_t sampled from demonstrations eq. 1.

\min_\theta \; \mathbb{E}_{\tau_t \sim \mathcal{T}_{0:T}}\big[\,\mathrm{dist}\big(\mathcal{G}_\theta(f_{gen} \mid s_0,\dots,s_t,\,l^*),\ \tau_t\big)\big] eq. 1

• A separate low-level policy π_φ(a | s_t, f_gen) outputs actions that reach the current subgoal §3.1 §3.4.2.

Framework loop

• The system is two-stage: (1) subgoal generation, (2) subgoal-conditioned policy rollout §3.2 fig. 2.
• The policy is rolled out until the progress evaluator says the subgoal is reached or has not been reached for too long; then a new subgoal is generated from the updated observation §3.2.

Ground-truth subgoal selection (training data)

Applied to each demonstration §3.3.1:

1. Compute, for every frame, the L2 distance between its progress embedding and the embedding of the final frame s_T.
2. Smooth this distance curve with LOWESS.
3. Compute the slope I_i of each frame and normalise the slopes to [0, 1].
4. Select frame i if the slope pattern around it changes sharply (the paper states this with two thresholds, δ1 and δ2).
5. Enforce a minimum gap of d frames between selected subgoals.

# Ground-truth subgoal selection, one demonstration  [§3.3.1] [§4.2.2]
# Input:  progress embeddings (T+1, d) of the frames of one demonstration
# Output: sorted list of frame indices, the ground-truth subgoals

def select_subgoals(embeddings, delta_1, delta_2, min_gap):
    # Distance of every frame to the final frame  [§3.3.1]
    final = embeddings[-1]
    distance = [norm(e - final) for e in embeddings]
    # Smooth the distance curve with LOWESS  [§3.3.1] [§4.2.2]
    distance = lowess(distance, fraction=1 / 6)
    # Rescale to [0, 1]  [§3.3.1]
    distance = min_max_scale(distance)
    # Slope of each frame  [§3.3.1]
    slope = [distance[i + 1] - distance[i] for i in range(T)]
    subgoals = []
    for i in range(1, T - 1):
        # Keep frame i when the slope pattern around it matches  [§3.3.1]
        if slope[i - 1] > delta_1 and slope[i + 1] < delta_2:
            # Enforce the minimum gap d between subgoals  [§3.3.1] [§4.2.2]
            if not subgoals or i - subgoals[-1] >= min_gap:
                subgoals.append(i)
    return subgoals

• The paper states δ1 = 0, δ2 = 0 and d = 7 for the main experiments §4.2.2. The smoothing factor is 0.167 §4.2.2.
• The paper's example: on CALVIN, 2.74 subgoals per task including the final goal, average interval 16.4; on the real robot, 1.6 subgoals, interval 23 §4.2.2.

Progress encoder (memory)

• A GRU reads the features of the frames chosen so far and outputs a progress embedding h_t §3.3.2.
• The feature input is one of CLIP, R3M or LIV, chosen once; the paper reports R3M as the best §4.2.2 §4.3.5 tab. 2.
• The embedding is fused with the language embedding and the diffusion timestep, and enters every level of the U-Net §3.4.1.

Progress evaluator (when to switch subgoal)

• Compare the current frame with the current generated subgoal using cosine similarity of a fine-tuned LIV embedding §3.3.3 §4.2.2.
• Move to the next subgoal when the similarity reaches δ = 0.96, or when the current subgoal has been pursued for more than λ = 20 rollout steps §3.3.3 §4.2.3.

Subgoal generator

• A 2D image-conditioned diffusion model. The paper accepts ControlNet or InstructPix2Pix; both are reported §3.4.1 app. D fig. 11.
• Conditions: the current image s_t (keeps the scene), the instruction embedding from a pretrained text encoder (the paper names CLIP's text encoder), and the progress embedding h_t §3.4.1.
• Classifier-free guidance uses two signals: text guidance with a negative prompt, and image guidance toward the current image §3.4.1. The negative prompt is the antonym of the instruction (for example, "close the drawer" for "open the drawer") §3.4.1.

Subgoal training loss

• The paper writes the training loss as a squared error between the generated and ground-truth subgoal images eq. 2.

\mathcal{L}_{gen}(f_{gen}, f_{gt}) = \sum \big(f_{gen} - f_{gt}\big)^2 eq. 2

• The released training script uses the standard diffusion noise-prediction loss on the subgoal latent, not a direct image-space loss (see Paper vs code) train_taksie.py.

Subgoal-conditioned policy (D-GCBC)

• The policy is π_φ(a | s_t, f_gen) and outputs the 6-DoF end-effector change and a gripper open/close action §3.4.2.
• It is a diffusion-based goal-conditioned behavior-cloning policy. It stacks the current and goal images, encodes them with a ResNet-50, and conditions a 3-layer 256-unit MLP §4.2.2.
• The policy predicts four action sequences and averages them dimension-wise §4.2.2.
• Training samples a sub-trajectory with a window between k_min and k_max; the final state is repeated k_δ times §3.5 §4.2.2. The paper's values: k_min = 8, k_max = 20 on CALVIN; k_min = 4, k_max = 20 on the real robot §4.2.2.

Training pipeline

1. Pretrain or reuse a text-to-image latent diffusion model §4.2.2.
2. Fine-tune the unconditional diffusion model on all frames of the trajectories (CALVIN: about 64 frames per trajectory; real robot: about 32) for 100k steps, batch 256 §4.2.2.
3. Initialise the conditioned generator from the fine-tuned weights, then co-train the ControlNet branch with the U-Net branch (or train InstructPix2Pix fully) for 300k steps, batch 256 §4.2.2.
4. Train the policy separately for 400k steps, batch 256 §4.2.2.

Inference (rollout)

# Adaptive subgoal rollout  [§3.2] [§3.3.3] [§4.2.3]
# Constants (paper): delta = 0.96 (cosine similarity), lam = 20 steps

def rollout(env, policy, generator, evaluator, instruction):
    # Generate the first subgoal from the start observation  [§3.2]
    obs = env.reset()
    history = [obs]
    subgoal = generator.generate(obs, instruction, history)
    steps_on_subgoal = 0
    while not env.done():
        # One low-level action toward the current subgoal  [§3.2] [§3.4.2]
        action = policy(obs, subgoal)
        obs = env.step(action)
        steps_on_subgoal += 1
        # Similarity of the current frame to the subgoal  [§3.3.3]
        similarity = cosine(evaluator(obs), evaluator(subgoal))
        # Move on after reaching the subgoal or after lam steps  [§4.2.3]
        if similarity >= 0.96 or steps_on_subgoal > 20:
            # Record the progress point, then make the next subgoal  [§3.3.2]
            history.append(obs)
            subgoal = generator.generate(obs, instruction, history)
            steps_on_subgoal = 0
    return env.success()

• The paper does not state a maximum episode length for the rollout §4.2.3. A commented-out line in the source mentions 360 steps from other work; it is not used in the text.
• Sampler: DDIM with 50 steps; text and image guidance scales 2.5 §4.2.3.

Hyperparameters

Values marked "not stated" are absent from the paper. Code values come from the released repository at code_commit.

Data settings from §4.1: CALVIN task D with 34 tasks and more than 5,000 training trajectories, average length about 64; evaluation on 1,000 five-task chains. The real robot is a Franka Research 3 with 28 tasks.

Results

Author-reported. Setup: CALVIN benchmark, task D, the 1,000-chain LH-MTLC evaluation; success rate in %, mean over 3 runs with standard deviation in parentheses tab. 2. Compared with HULC, LCD and SuSIE; HULC and LCD also use a gripper camera, TaKSIE uses only the static view §4.2.1.

Simulation (CALVIN, task D, LH-MTLC) tab. 2

The paper's text calls TaKSIE's relative gain larger at longer chains §4.3.2.

Goal-image type (CALVIN validation set, about 30 rollouts per task) tab. 1

Real robot (50 demonstrations per task, 5 trials per task) tab. 3

Training-set size (CALVIN) tab. 4

Unseen scenarios tab. 5

Slope thresholds (CALVIN validation set) tab. 6

The paper concludes the number of subgoals does not predict success; informativeness matters app. E. SuSIE selects 3 subgoals on average in this comparison app. E.

Other author claims: the real-robot run uses about 1.4k trajectories recorded over 7 hours, versus 60k in SuSIE §4.4. The paper contrasts its visual-goal approach with LLM-based planners, which need text descriptions and action APIs §2.

Paper vs code

Repository: https://github.com/Shua-Kang/TaKSIE at code_commit a822c05b844e93c54a1254354f41cf7d9fb5816c. It is linked as "Code" from the paper's project page https://live-robotics-uva.github.io/TaKSIE/. The GitHub owner "Shua-Kang" is not verifiably one of the two authors, but the project page links it as the code, and the README cites the same project page and a Hugging Face account "ShuaKang" for the checkpoints README.md README.md. The repo has no LICENSE file; several files carry Apache-2.0 headers copied from Hugging Face diffusers code, which does not license the rest. No verbatim code is quoted. The paper's own tex is CC BY 4.0 as recorded by arXiv.

Discrepancies found by reading the code:

1. Feature for the progress encoder. The paper's main row uses R3M tab. 2 and says one feature serves both subgoal selection and the GRU input §4.2.2. The training command caches CLIP features for the GRU train_taksie.bash and the inference config sets feature: clip cfg.yaml; the wrapper builds the CLIP path taksie_wrapper.py. Selection uses R3M select_keyframe.py README.md.
2. Distance reference for selection. The paper measures distance to the final frame s_T §3.3.1. The script measures distance from the current anchor frame to each later frame in the same segment, recursively select_keyframe.py, which does not match the paper's formula.
3. Threshold defaults. The paper uses δ1 = δ2 = 0 §4.2.2. The script defaults are -0.001 and 0.001 select_keyframe.py.
4. Threshold logic. The paper says I_{i-1} > δ1 and I_{i+1} < δ2 §3.3.1. The script tests the current slope against the first threshold with "less than or equal" and the previous slope against the second with "greater than or equal" select_keyframe.py. The roles and comparisons differ; I did not verify which one reproduces the paper's subgoal counts (paper: 2.74 on CALVIN).
5. Inference sampler. The paper says DDIM with 50 steps §4.2.3. The wrapper replaces the scheduler with UniPC multistep taksie_wrapper.py.
6. Guidance formula. The paper describes two guidance terms §3.4.1. The pipeline combines three branches: uncond + w_text·(text − image) + w_image·(image − uncond) pipeline_controlnet.py. This is an InstructPix2Pix-style combination that the paper does not write out.
7. Training loss. The paper writes an image-space squared error eq. 2. The script uses the standard noise-prediction MSE on the latent train_taksie.py.
8. ControlNet initialisation. The paper copies the fine-tuned diffusion weights into the ControlNet branch §4.2.2. The script calls from_unet(..., load_weights_from_unet=False) when no ControlNet path is given train_taksie.py. I read this as no weight copy; I did not run it.
9. Base weights. The paper cites a latent diffusion model §4.2.2. The training script's base model is stable-diffusion-v1-5 train_taksie.bash.
10. Training length and batch. The paper gives 300k steps at batch 256 §4.2.2. The script uses batch 32 with gradient accumulation 4 per device, 10000 epochs, and no max-steps cap train_taksie.bash. The total step count therefore depends on the dataset and the number of devices.
11. Action horizon. The paper predicts four action sequences §4.2.2. The training config uses 4 gcbc_train_config.py, the inference config sets 5 cfg.yaml. The wrapper uses 5 only to build the example input shape for checkpoint restore taksie_wrapper.py, and keeps 4 predictions in its buffer taksie_wrapper.py taksie_wrapper.py.
12. Policy training. The paper uses 400k steps and sub-trajectory window 8–20 on CALVIN §4.2.2. The config uses 2e6 steps and goal_delta [0, 24] gcbc_train_config.py gcbc_train_config.py. The policy training code is not in this repository; the README refers to an external JAX codebase README.md.
13. Negative prompts. The paper generates the negative prompt by swapping antonyms §3.4.1. The repository ships per-task YAML lists new_playtable_validation_negative.yaml; the code that produced them is not included.
14. Progress evaluator fine-tune. The paper fine-tunes LIV for 10k steps §4.2.2. The fine-tuned checkpoint is on Hugging Face README.md; the fine-tuning code is not included.
15. Hard-coded save path. On every regeneration after the first, the wrapper writes images to a fixed scratch directory taksie_wrapper.py. Outside that machine the write fails because the directory does not exist.
16. Unused settings. if_regenerate is read but never used taksie_wrapper.py. generation_evaluator is defined but never called taksie_wrapper.py.
17. Seeds. The paper reports 3 runs tab. 2 without naming seeds. The wrapper fixes the policy's sampling key to 42 taksie_wrapper.py taksie_wrapper.py, so repeated runs with that code sample actions identically.
18. Real-robot code. The repository contains only the CALVIN path. The real-robot data, the 28 tasks and the Franka setup are not in it.

Paper claims with no code in this repository: the 100k-step unconditional fine-tune, the progress-evaluator fine-tune, the real-robot results, the LOWESS-based selection statistics in §4.2.2, and the 3-run protocol.

Pitfalls

• Pick the feature deliberately. The paper's best row uses R3M; the released inference and training configs use CLIP for the GRU cfg.yaml. An implementation that follows the paper's table needs R3M for the GRU, and it should say so.
• Threshold sign conventions. δ1 and δ2 are both 0 in the paper but ±0.001 in the script. The script's comparisons differ from the paper's inequalities (Paper vs code, items 3–4). Subgoal count and success depend on these values tab. 6.
• Sensitivity. Table 6 shows success from 76.17 to 86.57 across small threshold changes, and subgoal count does not track success app. E.
• Progress threshold units. The paper states cosine similarity 0.96. The config stores 0.040, a distance. Use 1 − similarity, and confirm that LIV's sim is cosine; that function is in the LIV package, not this repository, and I did not verify it.
• Compute per subgoal. Each generation runs a diffusion sampler with 50 steps. The pipeline stacks three batch copies per step pipeline_controlnet.py. The paper gives no latency or throughput.
• Regeneration rule. The paper switches after λ = 20 steps or at similarity 0.96 §4.2.3. The code switches when either holds, but its step-count branch also requires more than min_per_frame steps (1) taksie_wrapper.py cfg.yaml. The paper does not mention that minimum.
• Gripper binarisation. The wrapper maps the gripper output to ±1 by sign taksie_wrapper.py. The paper does not say this.
• Action ensembling. The wrapper averages the valid buffered predictions per step taksie_wrapper.py. The paper says "dimension-wise mean" §4.2.2 without the buffer details.
• Image sizes differ by module. Generator: 256 px train_taksie.bash; policy: 200 px taksie_wrapper.py; gripper input: 84 px taksie_wrapper.py.
• Unconditional pretraining drives quality. The paper says the fine-tune on all frames is key to realistic images §4.2.2. This step is outside the repository.
• Failure modes the authors report. Wrong subgoals when the language encoder misses spatial terms (for example, push left vs push right) §4.7; blurry or distorted images, handled by regeneration §4.7; a good subgoal the policy fails to reach §4.7.
• Evaluation dependencies. CALVIN, the JAX GCBC runtime and LIV must be installed, plus the Hugging Face checkpoints README.md.

Lineage

• Builds on fixed-interval diffusion subgoals as in SuSIE, which this paper retrains from the original repository §4.2.1 §2.
• Uses a time-contrastive progress representation from R3M §3.3, LIV §3.3.3 and CLIP §4.2.2. Only LIV is cited in the bibliography with an arXiv id, 2306.00958 [raw/src/2410.11013/egbib.bib].
• Uses ControlNet (arXiv:2302.05543) and InstructPix2Pix as the conditioned generator §3.4.1 [raw/src/2410.11013/egbib.bib].
• Uses D-GCBC from BridgeData V2 (arXiv:2308.12952) as the policy family §4.2.1 [raw/src/2410.11013/egbib.bib].
• Compared against: HULC, LCD, SuSIE, D-LCBC §4.2.1.
• Superseded by: not known from the content of this wiki.