Isaac-GR00T/external_dependencies/SimplerEnv/eval_config.py

import os

from ml_collections import ConfigDict

# BridgeDatav2 metadata (computed from long time ago, not the new dataloader)

ACT_MEAN = [
    1.9296819e-04,
    1.3667766e-04,
    -1.4583133e-04,
    -1.8390431e-04,
    -3.0808983e-04,
    2.7425270e-04,
    5.9716219e-01,
]

ACT_STD = [
    0.00912848,
    0.0127196,
    0.01229497,
    0.02606696,
    0.02875283,
    0.07807977,
    0.48710242,
]

ACT_MIN = [
    -0.0437546,
    -0.052831028,
    -0.035931006,
    -0.14489305,
    -0.15591072,
    -0.26039174,
    -0.780331,
]  # 0.1% quantile

ACT_MAX = [
    0.04158026,
    0.05223833,
    0.05382493,
    0.15559858,
    0.142592,
    0.25956747,
    0.79311615,
]  # 99.9% quantile


REPO_DIR = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "..",
    "..",
)


jaxrl_gc_policy_kwargs = dict(
    encoder="resnetv1-34",
    encoder_kwargs=dict(
        pooling_method="avg",
        add_spatial_coordinates=False,
        act="swish",
    ),
    obs_image_size=256,
    seed=42,
    policy_class="gc_bc",
    checkpoint_path=os.path.expanduser(
        "~/SimplerEnv/checkpoint_75000"
    ),
    agent_kwargs=dict(
        network_kwargs=dict(
            hidden_dims=(256, 256, 256),
            dropout_rate=0.1,
        ),
        policy_kwargs=dict(
            tanh_squash_distribution=False,
            std_parameterization="fixed",
            fixed_std=[1, 1, 1, 1, 1, 1, 0.1],
        ),
        early_goal_concat=True,
        shared_goal_encoder=True,
        use_proprio=False,
        learning_rate=3e-4,
        warmup_steps=2000,
        decay_steps=int(2e6),
        # freeze_encoder=True,
    ),
    ACT_MEAN=ACT_MEAN,
    ACT_STD=ACT_STD,
)

soar_policy_kwargs = dict(
    **jaxrl_gc_policy_kwargs,
    susie_kwargs=dict(
        diffusion_checkpoint="kvablack/susie",
        diffusion_wandb="kvablack/dlimp-diffusion/9n9ped8m",
        diffusion_num_steps=50,
        prompt_w=7.5,
        context_w=4.0,
        diffusion_pretrained_path="lodestones/stable-diffusion-v1-5-flax",
        image_size=256,
    ),
)


octo_policy_kwargs = dict()


openvla_policy_kwargs = dict()
finetuned_openvla_open_drawer_kwargs = dict(
    **openvla_policy_kwargs,
    lora_adapter_dir=os.path.expanduser(
        "~/checkpoints/auto-eval-openvla-drawer/adapter_checkpoints/openvla-7b+expert_demos+b16+lr-0.0005+lora-r64+dropout-0.0/"
    ),
    dataset_stats_path=os.path.expanduser(
        "~/checkpoints/auto-eval-openvla-drawer/checkpoints/openvla-7b+expert_demos+b16+lr-0.0005+lora-r64+dropout-0.0/dataset_statistics.json"
    ),
)
finetuned_openvla_blue_sink_kwargs = dict(
    **openvla_policy_kwargs,
    lora_adapter_dir=os.path.expanduser(
        "~/checkpoints/auto-eval-openvla-blue-sink/adapter_checkpoints/openvla-7b+bridge_orig+b8+lr-0.0005+lora-r32+dropout-0.0+q-4bit--image_aug"
    ),
    dataset_stats_path=os.path.expanduser(
        "~/checkpoints/auto-eval-openvla-blue-sink/checkpoints/openvla-7b+bridge_orig+b8+lr-0.0005+lora-r32+dropout-0.0+q-4bit--image_aug/dataset_statistics.json"
    ),
)

scripted_close_door_policy_kwargs = dict(
    policy_save_path=os.path.join(
        REPO_DIR, "scripted_policies", "close the drawer.pkl"
    ),
    reset_language_cond="close the drawer",
)

scripted_out_of_drawer_handle_policy_kwargs = dict(
    policy_save_path=os.path.join(
        REPO_DIR, "scripted_policies", "out of drawer handle.pkl"
    ),
    recovery_steps=30,
)


def get_config(config_string):
    possible_structures = {
        "open_drawer": ConfigDict(
            dict(
                text_cond="open the drawer",
                success_detector_type="paligemma",
                success_detector_kwargs=dict(
                    vlm_type="paligemma",
                    vlm_config={
                        "model_id": os.path.expanduser(
                            "~/checkpoints/auto-eval-paligemma/paligemma-checkpoint-660-drawer-thick-handle-collect-2"
                        ),
                        "device": "cuda:0",
                        "quantize": True,
                    },
                    vlm_question="is the drawer open? answer yes or no",
                    ground_truth_answer_eval_task="yes",
                    ground_truth_answer_reset_task="no",
                ),
                # eval_policy_type="octo",
                # eval_policy_kwargs=octo_policy_kwargs,
                eval_policy_type="openvla",
                eval_policy_kwargs=openvla_policy_kwargs,
                # eval_policy_type="jaxrl_gc_policy",
                # eval_policy_kwargs=jaxrl_gc_policy_kwargs,
                # eval_policy_type="soar",
                # eval_policy_kwargs=soar_policy_kwargs,
                reset_policy_type="scripted",
                reset_policy_kwargs=scripted_close_door_policy_kwargs,
                recovery_policy_type="scripted",
                recovery_policy_kwargs=scripted_out_of_drawer_handle_policy_kwargs,
                workspace_bounds=dict(
                    x=[0.12, float("inf")],  # edge of table
                    y=[-float("inf"), float("inf")],
                    z=[-float("inf"), float("inf")],
                ),
                # x is towards front wall, y is towards left wall, z is up
                failure_conditions=[
                    {
                        "x": lambda x: x >= 0.43,
                        "y": lambda y: True,
                        "z": lambda z: z <= 0.03,
                    },  # robot pushing the micromove and falling
                    {
                        "x": lambda x: True,
                        "y": lambda y: True,
                        "z": lambda z: z <= 0,
                    },  # robot falling on the table somewhere
                    {
                        "x": lambda x: x >= 0.382,
                        "y": lambda y: y >= 0.01,
                        "z": lambda z: z <= 0.07,
                    },  # robot arm stuck behind drawer handle so it's hard to get back to reset
                ],
                stuck_conditions=[
                    {
                        "x": lambda x: 0.27 <= x <= 0.3,
                        "y": lambda y: -0.05 <= y <= 0.05,
                        "z": lambda z: 0.02 <= z <= 0.063,
                    },  # handle in drawer handle
                ],
            )
        ),
        "close_drawer": ConfigDict(
            dict(
                text_cond="close the drawer",
                success_detector_type="paligemma",
                success_detector_kwargs=dict(
                    vlm_type="paligemma",
                    vlm_config={
                        "model_id": os.path.expanduser(
                            "~/checkpoints/auto-eval-paligemma/paligemma-checkpoint-660-drawer-thick-handle-collect-2"
                        ),
                        "device": "cuda:0",
                        "quantize": True,
                    },
                    vlm_question="is the drawer open? answer yes or no",
                    ground_truth_answer_eval_task="no",
                    ground_truth_answer_reset_task="yes",
                ),
                # eval_policy_type="octo",
                # eval_policy_kwargs=octo_policy_kwargs,
                eval_policy_type="openvla_client",
                eval_policy_kwargs=openvla_policy_kwargs,
                # eval_policy_type="soar",
                # eval_policy_kwargs=soar_policy_kwargs,
                reset_policy_type="openvla_client",
                reset_policy_kwargs=openvla_policy_kwargs,
                # eval_policy_type="diffusion_policy_client",
                # reset_policy_type="openvla_client",
                # reset_policy_kwargs=openvla_reset_policy_kwargs,
                workspace_bounds=dict(
                    x=[-float("inf"), float("inf")],
                    y=[-float("inf"), float("inf")],
                    z=[-float("inf"), float("inf")],
                ),
            )
        ),
        "put_eggplant_in_sink": ConfigDict(
            dict(
                text_cond="put the eggplant into the blue sink",
                success_detector_type="paligemma",
                success_detector_kwargs=dict(
                    vlm_type="paligemma",
                    vlm_config={
                        "model_id": os.path.expanduser(
                            "~/checkpoints/auto-eval-paligemma/paligemma-checkpoint-540"
                        ),
                        "device": "cuda:0",
                        "quantize": True,
                    },
                    vlm_question="is the eggplant in the sink or in the basket? answer sink or basket or invalid",
                    ground_truth_answer_eval_task="sink",
                    ground_truth_answer_reset_task="basket",
                ),
                eval_policy_type="openvla",
                eval_policy_kwargs=openvla_policy_kwargs,
                # eval_policy_type="octo",
                # eval_policy_kwargs=octo_policy_kwargs,
                # reset_policy_type="soar",
                # reset_policy_kwargs=soar_policy_kwargs,
                reset_policy_type="scripted",
                reset_policy_kwargs=scripted_close_door_policy_kwargs,  # TODO: change
                workspace_bounds=dict(
                    x=[-float("inf"), float("inf")],
                    y=[-float("inf"), float("inf")],
                    z=[-float("inf"), float("inf")],
                ),
                failure_conditions=[],
                stuck_conditions=[],
            )
        ),
    }

    return possible_structures[config_string]