Devon018/3D-AGOS-Model

3D-AGOS insertion policies

Five trained diffusion policies for Franka plug/socket insertion, with the original checkpoints and the sensing calibration needed by the AGOS variants. Code: Devon018/Insertion-RL. See the code repository's docs/RELEASE.md for runtime setup, scan generation and source lineage.

Models

Bundle directory	Inputs	Gripper	Formal training run
depth_agos_tactile	wrist depth + AGOS + tactile	GelSight-Mini	tactile_agos_pcsim4/depth_tactile_cnn_unet_agos/20260927_010716
depth_agos_gelsight	wrist depth + AGOS	GelSight-Mini	tactile_agos_pcsim4_notactile/depth_cnn_unet_agos/20260929_075123
depth_tactile	wrist depth + tactile	GelSight-Mini	tactile_base_v5_minmax/depth_tactile_cnn_unet/20260908_000723
depth_agos_stock	wrist depth + AGOS	stock Franka	stock_agos_pcsim4_noforce/depth_cnn_unet_agos/20261001_003908
depth_agos_wrench	wrist depth + AGOS + 6-D wrist wrench	stock Franka	wrench_stock_agos_pcsim4_mdb012/depth_wrench_cnn_unet_agos/20260930_135438

These are the baseline policies; no RECAP-improved checkpoint is included. All five are the original checkpoints/last.pt at global step 145000, with training.ema.enabled=true. The repository loader chooses the EMA state used by the benchmark. best_fixed_noise.pt is a different checkpoint and is not substituted.

Each model directory contains checkpoints/last.pt, resolved_config.yaml, normalization_stats.json, normalization_stats.npz, summary.json and git_state.json. Checkpoints preserve the complete original schema, including the raw model, EMA, optimizer state, config and embedded normalization statistics. manifest.json records source runs, training commits, SHA-256, sizes, input shapes and CPU loading checks. Training paths in configs identify source datasets; those paths are not required for inference.

Inputs and architecture

All models use one observation step and predict a chunk of ten 9-D actions: EEF-frame translation plus rot6d. Depth is [B,1,1,480,640]; AGOS image is [B,1,2,80,80]. Tactile sensors are ordered [left,right], with depth, normal and two shear channels on a 20x25 grid. Wrench is [B,1,6] (force + moment). Shapes, normalization and history alignment are defined by each checkpoint. Use the saved normalizer and observation-history adapter before prediction.

Depth and AGOS use separate ResNet-18 encoders. Tactile uses a shared small CNN with sensor embeddings; wrench uses an MLP. Features concatenate and project to a 512-D global condition. A 1-D action UNet receives that condition through FiLM; DDIM uses 20 inference steps, epsilon prediction and eta 0. The AGOS image has plug and socket occupancy channels; the readout vector is produced by the estimator but these released policies encode its image branch.

The wrench model consumes the joint-space wrist_wrench_observer in the stock Franka environment. Its moment components use a 0.012 N m deadband before normalization; force components are retained. This setting is in the checkpoint's saved data config and is applied by the existing loader/live capture helpers.

Download and load

Activate the code repository's Python environment. Choose MODEL_ROOT on a volume with at least 4 GB available, and run from the cloned code repository root:

hf download Devon018/3D-AGOS-Model --local-dir "$MODEL_ROOT"
mkdir -p tmp/model_weights outputs/agos_plug_scans/sim_mv
cp "$MODEL_ROOT/pretrained/resnet18-f37072fd.pth" tmp/model_weights/
cp "$MODEL_ROOT/calibration/plug_scans/"*.npz outputs/agos_plug_scans/sim_mv/

The common pretrained ResNet file satisfies the stored initialization path when constructing the model; the trained checkpoint then strictly replaces its weights. Calibration scans are only needed for models that use AGOS.

import torch
from policy.frozen_base import load_frozen_multimodal_base

base = load_frozen_multimodal_base(
    run_dir="/path/to/model_root/depth_agos_tactile",
    checkpoint="checkpoints/last.pt",
    device=torch.device("cpu"),
)
print(base.describe())

For Isaac Gym evaluation, use scripts/eval_visuo_tactile_policy.sh with the selected bundle directory and asset ID. Set GRIPPER=stock for depth_agos_stock and depth_agos_wrench; the other three use GRIPPER=gelsight. Simulator evaluation needs the TacSL Isaac Gym runtime and consistent camera GPU routing. The existing Isaac Lab evaluator supports the GelSight robot; merging that port does not add a stock-gripper or wrist-wrench Isaac Lab evaluator.

Plug-scan calibration

calibration/plug_scans/ contains 54 existing 700xxx multiview scans. AGOS v4 (plug_scan=sim) expects them at outputs/agos_plug_scans/sim_mv/asset_<id>.npz in the code checkout. Each file contains surface points_plug [N,3] in metres and unit normals_plug [N,3] in the plug frame.

They are acquired by scripts/scan_plug_multiview.sh / isaacgymenvs/scan_plug_multiview.py: one bottom view and eight azimuths at each of three elevations (25 views), a 640x480 depth camera, 60-degree FOV, 8 cm distance, and 0.3 mm voxel subsampling. Simulated segmentation identifies plug pixels, and the known simulated plug pose registers them to the plug frame. OBJ surface sampling only audits coverage/error. Online point-to-plane ICP tracks the scanned plug against wrist depth, including in-hand slip/rotation. Real deployment requires measured depth, segmentation and object-frame registration for its own calibration.

The code release also bundles 900xxx connector geometry. These checkpoints and scans target the existing 700xxx assets; no 900xxx calibration or new connector policy-quality result is supplied.

Validation and scope

Preparation verified the five checkpoint schemas and normalizer checksums, byte-identical copies, strict EMA loading, finite synthetic CPU predictions with shape [1,10,9], and deterministic state-reload agreement. These are loading checks, not new closed-loop success measurements. No new GPU rollout or sim-to-real result was produced for this release. Training summaries preserve the original metrics.