[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!