Skip to content

Creating a S2 policy

s2-Diffusion is a method using diffusion policies with semantically segmented and depth input. This way, the diffusion policy can generalize better on the physical appearance of an object. For semantic segmentation, it uses GroundedSAM-2, and for the depth image it uses AnyDepth.

Policies trained on scooping rice evaluated on scooping cereal

An extension for s2-Diffusion already exists in the incar_baselines package. This example walks through how to use the extension, but you can also look at the code as an example for writing your own data processing steps.

Approach

The approach is to use three processing steps:

  1. A process step to apply GroundedSAM-2 on our video feature, and create a new feature video_segmented which contains the segmented image.
  2. A process step to apply AnyDepth on our video feature, and create a new feature video_depth which contains the depth image.
  3. A process step that combines the first channel of the video_segmented feature, and the first channel of the video_depth feature and puts it into a new feature called video_s2.

These process steps will hook into the DATASET hook, so that these features are available prior to training. They will also hook into the OBSERVATION hook so that the features are created from the live incomming data for inference.

The only VISUAL-type feature the policy is trained on, is the video_s2 feature.

Prerequisites

In the incar environment, install the s2-extras. This will allow us to get the extra dependencies needed (notably anydepth and grounded-sam dependencies)

# Remember to source the incar environment
pip install incar[s2]

Train Config

A train config for a s2-Diffusion policy looks as follows: 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
{
    "workspace": "<PATH_TO_WORKSPACE>",
    "model_save_name": "<MODEL_SAVE_NAME>",
    "preprocessed_dataset_name": "<DS_NAME>",
    "policy": {
        "type": "diffusion",
        "dt": 0.1,
        "preprocessing": {
            "raw_dataset_name": "<DS_NAME>",
            "validation_ratio": 0.15,
            "steps": [
                {
                    "type": "sample_dt",
                    "dt": 0.1
                },
                {
                    "type": "filter_by_buttons",
                    "button_feature_names": ["left.buttons"],
                    "button_names": ["gripValue"],
                },
                {
                    "type": "depth_anything",
                    "hooks": ["DATASET", "OBSERVATION"],
                    "source_feature_name": "viper_left_ee",
                    "target_feature_name": "viper_left_ee_depth"
                },
                {
                    "type": "grounded_SAM",
                    "hooks": ["DATASET", "OBSERVATION"],
                    "source_feature_name": "viper_left_ee",
                    "target_feature_name": "viper_left_ee_segmented",
                    "prompts": ["<PROMPT_A>", "<PROMPT_B>"],
                },
                {
                    "type": "combine_single_channels",
                    "hooks": ["DATASET", "OBSERVATION"],
                    "first_feature": "viper_left_ee_segmented",
                    "second_feature": "viper_left_ee_depth",
                    "target_feature_name": "viper_left_ee_s2"
                },
                {
                    "type": "downsample_video",
                    "features": ["viper_left_ee_s2"],
                    "new_size": [240, 320]
                }
            ]
        },
        "dependent_features": {
            "viper_left_ee": {
                "type": "VISUAL",
                "shape": [3, 240, 320]
            }
        },
        "input_features": {
            "viper_left_ee_s2": {
                "type": "VISUAL",
                "shape": [3, 240, 320]
            },
            "left.arm.ee.pose": {
                "type": "STATE",
                "shape": [7]
            },
            "left.gripper.joints.position": {
                "type": "STATE",
                "shape": [1]
            }
        },
        "output_features": {
            "left.commands.arm.ee.velocity": {
                "type": "ACTION",
                "shape": [6]
            },
            "left.commands.gripper.openclose": {
                "type": "ACTION",
                "shape": [1]
            }
        }
    },
    "steps": 100000,
    "save_freq": 20000,
    "log_freq": 1000,
    "validate_freq": 1000,
    "num_workers": 4
}

Note that this is pretty similar to any other policy to train. The only notable difference is the inclusion of the grounded_SAM, depth_anything and combine_single_channels processing steps.

Running inference

Run the policy like any other policy by selecting and loading it in the skills panel.