3D-AGOS-Model / README.md
Devon018's picture
Normalize model card formatting
767b13c verified
|
Raw History Blame Contribute Delete
6.34 kB
---
tags:
- robotics
- diffusion-policy
- depth
- tactile
- agos
- isaac-gym
- isaac-lab
---
# 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](https://github.com/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:
```bash
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.
```python
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.