From ff271e8b512d2164e40c47fe730d3de092489418 Mon Sep 17 00:00:00 2001
From: Pepijn <138571049+pkooij@users.noreply.github.com>
Date: Tue, 23 Dec 2025 23:58:34 +0100
Subject: [PATCH] pi fixes for dependencies (#2706)

* pi fixes for dependencies

* add walls sarm conflict

* also add conflicts for pi

* fix(ci): use --extra all instead of --all-extras + --no-extra

---------

Co-authored-by: Steven Palma <steven.palma@huggingface.co>
---
 .github/workflows/full_tests.yml         |  2 +-
 .github/workflows/unbound_deps_tests.yml |  2 +-
 pyproject.toml                           | 39 +++++++++++++++++++++++-
 src/lerobot/policies/pi0/modeling_pi0.py |  7 +++--
 4 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/full_tests.yml b/.github/workflows/full_tests.yml
index ad222b04f..7962d4a3a 100644
--- a/.github/workflows/full_tests.yml
+++ b/.github/workflows/full_tests.yml
@@ -85,7 +85,7 @@ jobs:
           python-version: ${{ env.PYTHON_VERSION }}
 
       - name: Install lerobot with all extras
-        run: uv sync --all-extras --no-extra groot --no-extra wallx # TODO(Steven): Make flash-attn optional
+        run: uv sync --extra all # TODO(Steven): Make flash-attn optional
 
       - name: Run pytest (all extras)
         run: uv run pytest tests -vv --maxfail=10
diff --git a/.github/workflows/unbound_deps_tests.yml b/.github/workflows/unbound_deps_tests.yml
index 95562d0dd..bf93dcda6 100644
--- a/.github/workflows/unbound_deps_tests.yml
+++ b/.github/workflows/unbound_deps_tests.yml
@@ -78,7 +78,7 @@ jobs:
           echo "Dependencies unbound:" && cat pyproject.toml
 
       - name: Install lerobot with all extras
-        run: uv sync --all-extras --no-extra groot --no-extra wallx # TODO(Steven): Make flash-attn optional
+        run: uv sync --extra all # TODO(Steven): Make flash-attn optional
 
       - name: Run pytest (all extras)
         run: uv run pytest tests -vv
diff --git a/pyproject.toml b/pyproject.toml
index 48e071d32..61b802bc5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -168,7 +168,7 @@ all = [
     "lerobot[kinematics]",
     "lerobot[intelrealsense]",
     # "lerobot[wallx]",
-    "lerobot[pi]",
+    # "lerobot[pi]", TODO(Pepijn): Update pi to transformers v5
     "lerobot[smolvla]",
     # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn
     "lerobot[xvla]",
@@ -405,6 +405,10 @@ conflicts = [
         { extra = "wallx" },
         { extra = "xvla" },
     ],
+    [
+        { extra = "wallx" },
+        { extra = "sarm" },
+    ],
     [
         { extra = "wallx" },
         { extra = "hilserl" },
@@ -417,4 +421,37 @@ conflicts = [
         { extra = "wallx" },
         { extra = "all" },
     ],
+    # pi uses custom branch which conflicts with transformers-dep
+    [
+        { extra = "pi" },
+        { extra = "transformers-dep" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "smolvla" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "groot" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "xvla" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "sarm" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "hilserl" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "libero" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "all" },
+    ],
 ]
diff --git a/src/lerobot/policies/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py
index 0d9c77e00..e4970dcf1 100644
--- a/src/lerobot/policies/pi0/modeling_pi0.py
+++ b/src/lerobot/policies/pi0/modeling_pi0.py
@@ -93,10 +93,11 @@ def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedd
 
 
 def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
-    alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device)
-    beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device)
+    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
+    alpha_t = torch.tensor(alpha, dtype=torch.float32)
+    beta_t = torch.tensor(beta, dtype=torch.float32)
     dist = torch.distributions.Beta(alpha_t, beta_t)
-    return dist.sample((bsize,))
+    return dist.sample((bsize,)).to(device)
 
 
 def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (exact copy)