[Pi0] The dataset_stats are not provided on the model weights

[dfloreaa]

Starting off, this is kind of a duplicate of #694

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

[DelinQu]

This is truly amazing! Could you share how you managed to load the simpler environment into LeRobot? I would greatly appreciate it if you could share the code as well. Thank you so much for your help!

[lixinyao11]

Since Dataset statistics to be used for normalization. If not passed here, it is expected that they will be passed with a call to load_state_dict before the policy is used., maybe we can manually load the statistics into model by calling load_state_dict one more time with a state dict containing only the normalizer and the unnormalizer? (i.e. load the model and statistics separately).

btw, do you have any idea of how to compute the mean and std for simpler env?