Probe Retrieval

1. Core Idea

Probe Retrieval is a lightweight test-time adaptation method designed to improve the robustness of VLA policies under physical parameter shifts. We do not retrain OpenVLA, nor do we learn an additional dynamics model. Instead, when the robot first makes contact with the target object, we insert a short and standardized physical probe. This probe produces a response signature that reflects the current object-environment dynamics and is then used to guide retrieval.

The core motivation is that visual similarity alone is insufficient under hidden physical changes. Two episodes may look almost identical before contact, but the required post-contact actions can be very different because the object may be heavier, more slippery, harder to push, or more likely to overshoot. Therefore, we first use visual retrieval to find semantically and visually similar candidate episodes, and then use the physical response features produced by the probe to re-rank these candidates. The final action correction is retrieved from successful examples whose physical responses are most similar to the current episode.

The method is lightweight, plug-and-play, and training-free. It only relies on a fixed short probe and a memory bank constructed from historical rollouts.

---

2. Experimental Pipeline

2.1 Data Collection Phase

The goal of the data collection phase is to build a memory bank. Each memory entry contains the visual observation, OpenVLA action, probe response, physical parameter setting, and final task success label.

At the beginning of each episode, OpenVLA executes the task normally. The robot first moves toward the target object. At every simulator step, we monitor MuJoCo contacts. When meaningful contact between the gripper and the target object is detected, the normal OpenVLA rollout is paused and a fixed probe sequence is executed.

The contact detection logic is implemented in:

probe/contact_detector.py

Specifically, the detector reads MuJoCo contact information from:

env.sim.data.contact
env.sim.data.ncon

By default, we do not trigger the probe using arbitrary collisions. Instead, we use filtered contact detection. A contact is considered valid only if:

• at least one geom in the contact pair matches a gripper-related pattern, such as gripper, finger, leftpad, or rightpad;
• the other geom matches the target object pattern of the current task;
• the condition holds for min_contact_steps consecutive steps. In the final setting, this threshold is set to 1.

If valid geom names are unavailable in the simulator, the detector falls back to:

env.sim.data.ncon > 0

The probe action is defined in:

rollout/real_libero.py

The final probe is a fixed 6-step sequence:

1. Close the gripper for 2 control steps.
2. Keep the gripper closed and lift upward with delta_z = 0.015 for 3 control steps.
3. Keep the gripper closed and hold for 1 control step.

During the probe, we continue stepping the same LIBERO environment. Therefore, the probe is not computed offline and does not copy simulator states. It physically changes the current episode state. After the probe finishes, the rollout resumes from the post-probe observation.

The probe stops when one of the following conditions is met:

• the fixed probe sequence is completed;
• the environment returns done=True;
• the task has already succeeded;
• the reward exceeds the success threshold.

Each episode stores at least the following information:

• task name;
• physics setting;
• image observation at contact;
• image embedding;
• OpenVLA action at contact;
• post-probe OpenVLA action chunk;
• whether the probe was triggered;
• probe response features;
• final success label.

The probe response features include:

• mean_force
• max_force
• force_std
• contact_steps
• contact_ratio
• probe_obj_displacement
• end_effector_movement
• probe_motion_ratio
• object_rotation_change_deg
• max_object_step_displacement
• post_probe_object_to_target_distance

The final memory bank used for retrieval is stored at:

artifacts/collect_goal_probe_memory_sweep_4gpu/20260429_015822/merged/

Key files include:

artifacts/collect_goal_probe_memory_sweep_4gpu/20260429_015822/merged/episodes.jsonl
artifacts/collect_goal_probe_memory_sweep_4gpu/20260429_015822/merged/memory_bank/metadata.jsonl
artifacts/collect_goal_probe_memory_sweep_4gpu/20260429_015822/merged/memory_bank/arrays.npz
artifacts/collect_goal_probe_memory_sweep_4gpu/20260429_015822/merged/merge_summary.json

---

2.2 Inference Phase

The inference phase uses the same contact-triggered probe protocol as the data collection phase.

At the beginning of an episode, OpenVLA executes normally. When the robot first makes valid contact with the target object, the system performs retrieval and action correction:

1. Save the current image.
2. Use the current image to retrieve visually similar episodes from the same-task memory bank.
3. Execute the same fixed probe in the current environment.
4. Compute response features from the probe.
5. Re-rank the visually retrieved candidates using response-feature distance.
6. Select the top-5 candidates with the most similar responses.
7. Keep only memory episodes that ended in success.
8. Aggregate the post-probe actions from these successful candidates.
9. Fuse the retrieved action with the current OpenVLA action.
10. Continue the episode from the post-probe observation.

---

2.2.1 Image Retrieval

Image retrieval is implemented in:

retrieval/image_embedder.py
retrieval/image_retrieval.py

The current image embedding is a lightweight implementation that avoids introducing additional large vision models. It concatenates:

• RGB color histogram features;
• low-resolution spatial grid features.

The concatenated embedding is then L2-normalized.

Given a query image embedding and memory-bank image embeddings, we compute cosine similarity:

score(q, x) = cosine(q, x)

Retrieval is restricted to memory entries from the same task. We take the top-20 visually most similar candidates.

---

2.2.2 Probe Response Re-ranking

Response re-ranking is implemented in:

retrieval/response_features.py
retrieval/probe_rerank.py

For each task, we first compute feature-wise mean and standard deviation using probe-triggered memory episodes from the response bank. Both the query and candidate response features are normalized using same-task statistics:

z(x_i) = (x_i - mean_i) / (std_i + eps)

We then compute the weighted L1 distance between the query and each candidate:

d(query, candidate) = sum_i w_i * |z(query_i) - z(candidate_i)|

The final response feature keys used for retrieval are:

probe_start_step
contact_steps
contact_ratio
mean_force
max_force
force_std
probe_obj_displacement
end_effector_movement
probe_motion_ratio
post_probe_object_to_target_distance

where:

probe_motion_ratio = probe_obj_displacement / (end_effector_movement + eps)

Candidates are sorted by response distance in ascending order. We first take the top-5 most response-similar candidates, and then filter them by final success.

---

2.2.3 Action Aggregation

For each successful retrieved candidate, we read its saved post-probe action chunk. If the post-probe action chunk is unavailable, we fall back to the OpenVLA action at the contact step.

The default aggregation method is soft averaging. For candidate j, the weight is:

weight_j = exp(-distance_j / temperature)

The weights are normalized, and the retrieved action is computed as:

action_retrieved = sum_j normalized_weight_j * action_j

Finally, the retrieved action is fused with the current OpenVLA action:

action = alpha * action_retrieved + (1 - alpha) * action_openvla

This allows OpenVLA to preserve its original task execution capability while the retrieved action provides a local correction based on the physical response.

---

3. Experiments and Results

3.0 Environment and Model

The final experiments are conducted on:

LIBERO-goal
openvla/openvla-7b-finetuned-libero-goal

We select five tasks:

push the plate to the front of the stove
put the bowl on the stove
put the bowl on top of the cabinet
put the cream cheese in the bowl
put the wine bottle on top of the cabinet

The test-set physics settings are:

mu = 0.05
mu = 0.5
m = 0.05x
m = 5x
m = 10x

where:

• mu is the absolute MuJoCo friction coefficient, not a multiplier;
• m is a multiplier relative to the original object mass.

Each method is evaluated on:

5 tasks x 5 physics settings x 10 episodes = 250 episodes

---

3.1 Result Paths

Original OpenVLA results on the unmodified LIBERO-goal environment:

artifacts/basemodel_original_environment/bl_ov7bftGoal_libGoal/

Base OpenVLA results on the physics sweep test set:

artifacts/basemodel_sweep_train_test/episodes_base.jsonl
artifacts/basemodel_sweep_train_test/summary_base.json

Full Probe Retrieval results:

artifacts/goal_probe_retrieval_test_eval_4gpu/20260429_161402/

Probe-only ablation:

artifacts/goal_probe_retrieval_test_eval_4gpu_video_local/probe_only/20260502_165701/

Random Retrieve ablation:

artifacts/random_retrieve/20260502_163803/
artifacts/goal_probe_retrieval_test_eval_1gpu_nyu/random_retrieve/20260502_163803/

Image-only Retrieval ablation:

artifacts/image_only_retrieval/20260503_155847/

Presentation videos with visible probe execution:

artifacts/goal_probe_retrieval_test_eval_4gpu_video_local/full_probe_retrieval_visible_probe/
artifacts/goal_probe_retrieval_test_eval_4gpu_video_local/probe_only_visible_probe/

---

3.2 Result Comparison

Overall success rate on the final test sweep:

Method	Successes / 250	Success Rate
Base OpenVLA	124	49.6%
Probe-only	130	52.0%
Image-only Retrieval	110	44.0%
Random Retrieve	118	47.2%
Ours: Probe Retrieval	137	54.8%

Average success rate by test physics setting:

Method	mu=0.05	mu=0.5	m=0.05x	m=5x	m=10x
Base OpenVLA	0.56	0.58	0.46	0.52	0.36
Probe-only	0.54	0.52	0.52	0.56	0.48
Image-only Retrieval	0.40	0.48	0.48	0.44	0.40
Random Retrieve	0.48	0.44	0.50	0.52	0.42
Ours: Probe Retrieval	0.48	0.60	0.54	0.56	0.56

Overall, Probe Retrieval achieves the highest success rate across 250 test episodes. Compared with Base OpenVLA, the results suggest that physical response information provides an additional useful signal for retrieval. Compared with image-only retrieval and random retrieval, the improvement indicates that the gain does not come from retrieval alone, but from response-based re-ranking of physically relevant candidates.

---

4. Code Structure

Core code:

probe/contact_detector.py
diagnostics/probe_goal_physics_sweep.py
diagnostics/collect_goal_probe_memory_sweep.py
diagnostics/eval_goal_probe_retrieval_test.py
diagnostics/baseline_goal_physics_sweep.py
retrieval/memory_bank.py
retrieval/image_embedder.py
retrieval/image_retrieval.py
retrieval/response_bank.py
retrieval/response_features.py
retrieval/probe_rerank.py
rollout/real_libero.py
rollout/video_utils.py

Common scripts:

scripts/run_goal_baseline_4gpu.sh
scripts/run_baseline_goal_physics_sweep_4gpu.sh
scripts/run_probe_goal_physics_sweep_4gpu.sh
scripts/run_goal_probe_memory_collect_4gpu.sh
scripts/run_goal_probe_retrieval_test_eval_4gpu.sh
scripts/resume_goal_probe_retrieval_test_eval_4gpu.sh
scripts/run_visible_probe_full_and_probe_only_4gpu.sh
scripts/submit_goal_probe_retrieval_test_eval_*.sbatch

---

5. Probe Response Features

Feature	Explanation
probe_start_step	The timestep at which contact is detected and the probe starts.
contact_steps	The number of probe steps during which gripper-target contact remains active.
contact_ratio	The fraction of probe steps with active gripper-target contact.
mean_force	The mean force magnitude measured during the probe.
max_force	The maximum force magnitude measured during the probe.
force_std	The standard deviation of force magnitude during the probe.
probe_obj_displacement	The net displacement of the target object from probe start to probe end.
end_effector_movement	The total movement distance of the robot end-effector during the probe.
probe_motion_ratio	The ratio between target-object displacement and end-effector movement during the probe.
post_probe_object_to_target_distance	The distance from the target object to the goal region after the probe.