Bringup Pi_0 Pytorch model

[ashokkumarkannan1]

Ticket

Link to Github Issue - tenstorrent/tt-xla#2979

Problem description

to bringup Pi_0 model

What's changed

added Pi_0 pytorch model

Checklist

• New/Existing tests provide coverage for changes

Logs

• xla_inference_test_base.log
• xla_inference_test_libero.log

[ppadjinTT]

Can you explain why do we need to override the forward method in this way?

[ashokkumarkannan1]

• The select_action method is responsible for inference (as described in source). The Policy.forward method in LeRobot is only used during training to compute the loss and is not part of the inference path.

• Since our use case targets inference, we needed to override forward to incorporate the select_action logic, ensuring the correct execution flow during model inference. (as already mentioned in the forward method docstring)

• Additionally, the steps inside preprocess_for_sampling, which are normally invoked within select_action, have been separated and applied directly during the load_inputs stage. This is because those steps are specifically responsible for preprocessing a batch before action sampling.

[ppadjinTT]

can you avoid global module overriding function like this. You can make a function like get_custom_pi0_policy that returns this object instead of it being global

[ashokkumarkannan1]

• As suggested, I’ve added a get_custom_pi0_policy factory function. This function loads the original PI0Policy, overrides the forward and preprocess_for_sampling methods on that instance, and returns the modified policy object.

• This avoids any global module-level overrides and keeps the changes scoped to the specific instance being used.

[ppadjinTT]

due to the global PI0Policy object I mentioned earlier, it is not clear whether you are using the original PI0Policy object, or is it getting overwritten by the one from .src.model? Please change to from .src.model import get_custom_pi0_policy and then I assume this from lerobot.policies.pi0 import PI0Policy import will be unnessesary

[ashokkumarkannan1]

• I’ve updated the load_model function to instantiate the model using get_custom_pi0_policy, which internally handles importing the original PI0Policy and applying the required overrides.

• As a result, the direct import from lerobot.policies.pi0 import PI0Policy has been removed from this file to avoid any confusion about which policy object is being used.

[ppadjinTT]

Is this the standard way to do this? I did a quick google search and I found this pip dependency for PI0:

'For lerobot 0.4.0, if you want to install pi tag, you will have to do: pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git".'

[ashokkumarkannan1]

• this pip install "lerobot[pi]@git+[https://github.com/huggingface/lerobot.git installs LeRobot’s default dependencies include versions of torch(<2.8) and other related ones that conflict with our target XLA environment.
• Because of that, we install LeRobot itself with --no-deps, then install only the required dependencies manually (as listed in our requirements.txt), ensuring compatibility with the current environment.

[ppadjinTT]

Ok tnx, makes sense

[ppadjinTT]

Thanks!