lerobot

/

pi0

like 82

Follow

LeRobot 510

After making some adjustments to the config.json file from the original download in order to work with the SimperEnv Google Robot environment such that now it looks like this:

{
    "type": "pi0",
    "n_obs_steps": 1,
    "normalization_mapping": {
        "VISUAL": "IDENTITY",
        "STATE": "MEAN_STD",
        "ACTION": "MEAN_STD"
    },
    "input_features": {
        "observation.image.top": {
          "shape": [
            3,
            224,
            224
          ],
          "type": "VISUAL"
        },
        "observation.state": {
          "shape": [
            7
          ],
          "type": "STATE"
        }
    },
    "output_features": {
      "action": {
        "shape": [7],
        "type": "ACTION"
    }
    },
    "chunk_size": 50,
    "n_action_steps": 50,
    "max_state_dim": 32,
    "max_action_dim": 32,
    "resize_imgs_with_padding": [
        224,
        224
    ],
    "empty_cameras": 0,
    "adapt_to_pi_aloha": false,
    "use_delta_joint_actions_aloha": false,
    "tokenizer_max_length": 48,
    "proj_width": 1024,
    "num_steps": 10,
    "use_cache": true,
    "attention_implementation": "eager",
    "freeze_vision_encoder": true,
    "train_expert_only": false,
    "train_state_proj": true,
    "optimizer_lr": 2.5e-05,
    "optimizer_betas": [
        0.9,
        0.95
    ],
    "optimizer_eps": 1e-08,
    "optimizer_weight_decay": 1e-10,
    "scheduler_warmup_steps": 1000,
    "scheduler_decay_steps": 30000,
    "scheduler_decay_lr": 2.5e-06
}

I previously tried not performing any operations over the robot's outputs

"normalization_mapping": {
        "VISUAL": "IDENTITY",
        "STATE": "IDENTITY",
        "ACTION": "IDENTITY"
}

I am directly using the policy with a batch of actions

action = pi_policy.select_action(in_batch)

With in_batch being:

{'observation.image.top': tensor([[[[0.2157, 0.2157, 0.2196,  ..., 0.3647, 0.3412, 0.3255],
           [0.2118, 0.2118, 0.2118,  ..., 0.3373, 0.3098, 0.2941],
           [0.2118, 0.2118, 0.2118,  ..., 0.3294, 0.3020, 0.2902],
           ...,
           [0.4157, 0.4157, 0.4196,  ..., 0.6863, 0.6824, 0.6824],
           [0.4157, 0.4157, 0.4196,  ..., 0.6824, 0.6784, 0.6745],
           [0.4157, 0.4157, 0.4196,  ..., 0.6745, 0.6706, 0.6667]],
 
          [[0.1373, 0.1373, 0.1412,  ..., 0.2627, 0.2392, 0.2235],
           [0.1333, 0.1333, 0.1333,  ..., 0.2353, 0.2078, 0.1922],
           [0.1333, 0.1333, 0.1333,  ..., 0.2235, 0.2000, 0.1882],
           ...,
           [0.2902, 0.2902, 0.2941,  ..., 0.5412, 0.5373, 0.5373],
           [0.2902, 0.2902, 0.2941,  ..., 0.5373, 0.5333, 0.5294],
           [0.2902, 0.2902, 0.2941,  ..., 0.5294, 0.5255, 0.5216]],
 
          [[0.0941, 0.0941, 0.0980,  ..., 0.2118, 0.1882, 0.1725],
           [0.0902, 0.0902, 0.0902,  ..., 0.1882, 0.1569, 0.1412],
           [0.0902, 0.0902, 0.0902,  ..., 0.1765, 0.1490, 0.1373],
           ...,
           [0.1843, 0.1843, 0.1882,  ..., 0.4118, 0.4078, 0.4078],
           [0.1843, 0.1843, 0.1882,  ..., 0.4078, 0.4039, 0.4000],
           [0.1843, 0.1843, 0.1882,  ..., 0.4000, 0.3961, 0.3922]]]],
        device='cuda:0'),
 'observation.state': tensor([[ 9.3738e-01,  8.5981e-01,  3.8762e-01,  6.4738e-01,  1.6412e+00,
           1.1104e+00, -4.7183e-01, -9.9201e-05, -9.9773e-05, -2.8596e-03,
           7.8514e-01]], device='cuda:0'),
 'task': ['pick coke can']}

And directly executing the output action on the environment... But this resulted in some funky results:

So therefore I need the mean and std for both the state and the action to be rescaled, but I'm guessing these values are domain-specific and have to be computed by my own, right? If so, there is not a proper way to load these values into the model (as the normalizer and unnormalizer are created on the __init__() method).

This means that the pretrained weights that are loaded should feature such values, as per the error message
AssertionError: meanis infinity. You should either initialize withstats as an argument, or use a pretrained model.

Could the weights be updated to feature this values? Otherwise, could I get some insights as to how to properly scale these values?

Thank you very much in advance