Add direct access to action chunks (#1020)

* fix: sharing predicted chunk with user

* [pre-commit.ci] pre-commit autoupdate (#1011)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Revert "[pre-commit.ci] pre-commit autoupdate" (#1025)

* fix(ci): Pin draccus (<0.10.0) and torch (<2.7) to fix pipeline (#1022)

Co-authored-by: imstevenpmwork <steven.palma@huggingface.co>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>

* fix(ci): Pin `torchcodec` (==0.2.1) to fix pipeline temporarly (#1030)

* Update tutorial (#1021)

Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>

* Add description motor order SO-101 leader (#1051)

* feat(encoding): switching to PyAV for ffmpeg related tasks (#983)

* feat(docs): Add new docs build process (#1046)

Co-authored-by: Mishig Davaadorj <dmishig@gmail.com>
Co-authored-by: Steven Palma <steven.palma@huggingface.co>

* Docs: adapt text + fix video code (#1064)

* Fix typos (#1070)

* docs: minor corrections and clean-up (#1089)

* Update 10_use_so100.md; use diff syntax (#944)

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>

* Update 12_use_so101.md (#1081)

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>

* bug fix for #1071 When --display_data=true, Failed running control_robot. (#1073)

* Add editable -e for feetech install command (#1133)

* Fix: emptying action queue between resets (#1117)

* fix: typos and grammar (#1148)

* Update README.md (#1160)

* Update README.md (#1163)

* [Fix]  Unpin torch beyond 2.6.0 & torchcodec beyond 0.2.1  (#1127)

* (hotfix): nightly CI by clipping pymunk version below 7.0.0 (#1182)

* [pre-commit.ci] pre-commit autoupdate (#1048)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Simon Alibert <simon.alibert@huggingface.co>

* Add SmolVLA (#1175)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: fracapuano <francesco.capuano@huggingface.co>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Co-authored-by: Dana Aubakirova <118912928+danaaubakirova@users.noreply.github.com>
Co-authored-by: Remi <remi.cadene@huggingface.co>

* Fix SmolVLA loss not sent to wandb (#1198)

* Hardware API redesign (#777)

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Steven Palma <steven.palma@huggingface.co>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: Pepijn <pepijn@huggingface.co>

* fix(smolvla): update record.py, fix populate_queues and remove unused dependencies (#1208)

* replaced OBS_ROBOT with OBS_STATE constant (#1211)

* Fix test_teleoperate (#1216)

* Fix LeKiwi example (#1217)

* Fix smolVLA dependencies (#1218)

* fix(pyserial): adding pyserial dependency to global ones (#1219)

* Update SmolVLA README.md (#1228)

* Fix unable to set camera width/height to non-default (#1225)

* Update tutorial link (#1250)

* update KochFollower.get_observation() so it returns same observation structure as SO101 (#1248)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* [pre-commit.ci] pre-commit autoupdate (#1185)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>

* Proposal for fix for enter_pressed on Windows (#1230)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>

* fix: update pi0 dependency version constraint (#1247)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Match motor names with ids lekiwi (#1261)

* fix issues: checkpoints keys mismatch and 'task' tokenisation in smolvla (#1256)

Co-authored-by: danaaubakirova <d.aubakirova@alumni.edu.kz>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Co-authored-by: Simon Alibert <simon.alibert@huggingface.co>

* fix(docs): update realsense documentation (#1268)

* Use HF Papers (#1120)

* Skip normalization parameters in load_smolvla (#1274)

* fix(record): no teleop needed when running with policy (#1284)

* Port HIL SERL (#644)

Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: Eugene Mironov <helper2424@gmail.com>
Co-authored-by: s1lent4gnt <kmeftah.khalil@gmail.com>
Co-authored-by: Ke Wang <superwk1017@gmail.com>
Co-authored-by: Yoel Chornton <yoel.chornton@gmail.com>
Co-authored-by: imstevenpmwork <steven.palma@huggingface.co>
Co-authored-by: Simon Alibert <simon.alibert@huggingface.co>

* fix(docs): SmolVLA fine-tuning getting started (#1201)

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: danaaubakirova <d.aubakirova@alumni.edu.kz>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Co-authored-by: Francesco Capuano <francesco_capuano@aol.com>
Co-authored-by: Steven Palma <steven.palma@huggingface.co>

* chore(teleop): print calibration path saved (#1286)

* chore(dependencies): add gamepad support with pygame and hidapi (#1287)

* Robot integration tutorial (#1285)

* fix(docs): update send_feedback docstrings

* Add sim tutorial, fix lekiwi motor config, add notebook links (#1275)

Co-authored-by: AdilZouitine <adilzouitinegm@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: s1lent4gnt <kmeftah.khalil@gmail.com>
Co-authored-by: Michel Aractingi <michel.aractingi@gmail.com>
Co-authored-by: Eugene Mironov <helper2424@gmail.com>
Co-authored-by: imstevenpmwork <steven.palma@huggingface.co>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

* Fixes on robot integration tutorial (#1290)

* Add keyboard teleop device to control the end effector robot  (#1289)

* Improve type hints (#1293)

* fix(record): no teleop arg in reset environment (#1294)

* `learner.py` import so101_leader instead of so100 (#1295)

Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>

* Fixing `PI0` Policy (#1297)

* `gym_manipulator.py` Remove None value action_intervention of BaseLeaderTeleoperator (#1299)

* (chore): incorrect resume parameter in recording documentation (#1301)

* Update lekiwi.mdx  (#1229)

* bump `pi0` and `hil` transformers version (#1298)

* docs: fix imitation learning robots docs command (#1308)

* fix(benchmarks): remove .numpy() from frame in benchmark script (#1354)

* add smolvla to the supported policies to run tests (:

* add: chunk-level access for the policy

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add: smolvla in availables

* remove: smolvla from library supported policies

* fix: change env for training, xarm is broken as of now

* add: predict_action_chunk to all supported policies

* fix: add robot type constants

* add: predict action chunk in base policy class

* restore original Makefile

* fix: minor

* fix: dict keys come from lerobot/constants

* fix: improve act encapsulation, properly supporting temporal ensembling

* fix: smolvla action chunking

* fix: very minor, but very annoying

* fix: minor

* fix minor naming

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com>

* fix: refactoring inference for single actions and chunks into different components

* fix: minor

* fix: temporal ensembling

* fix: moving populate queues out of modular component for batch preparation

* fix: minor for CI

* fix: smovla debug

* fix: reward classifier, maybe the last policy lacking?

---------

Signed-off-by: Francesco Capuano <74058581+fracapuano@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: imstevenpmwork <steven.palma@huggingface.co>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Caroline Pascal <caroline8.pascal@gmail.com>
Co-authored-by: Mishig Davaadorj <dmishig@gmail.com>
Co-authored-by: omahs <73983677+omahs@users.noreply.github.com>
Co-authored-by: CharlesCNorton <135471798+CharlesCNorton@users.noreply.github.com>
Co-authored-by: masato-ka <jp6uzv@gmail.com>
Co-authored-by: Ragnar <rodiondenmark@gmail.com>
Co-authored-by: mshukor <mustafa.shukor97@gmail.com>
Co-authored-by: Simon Alibert <simon.alibert@huggingface.co>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Co-authored-by: Dana Aubakirova <118912928+danaaubakirova@users.noreply.github.com>
Co-authored-by: Remi <remi.cadene@huggingface.co>
Co-authored-by: Ben Zhang <5977478+ben-z@users.noreply.github.com>
Co-authored-by: Pepijn <pepijn@huggingface.co>
Co-authored-by: Dhruva <51377003+utterwqlnut@users.noreply.github.com>
Co-authored-by: Daisuke Sato <tiryoh@gmail.com>
Co-authored-by: Sarunas Kalade <sarunas.kalade@amd.com>
Co-authored-by: koenvanwijk <koenvanwijk@users.noreply.github.com>
Co-authored-by: Yushun Xiang <73413365+YushunXiang@users.noreply.github.com>
Co-authored-by: danaaubakirova <d.aubakirova@alumni.edu.kz>
Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>
Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: Eugene Mironov <helper2424@gmail.com>
Co-authored-by: s1lent4gnt <kmeftah.khalil@gmail.com>
Co-authored-by: Ke Wang <superwk1017@gmail.com>
Co-authored-by: Yoel Chornton <yoel.chornton@gmail.com>
Co-authored-by: Michel Aractingi <michel.aractingi@gmail.com>
Co-authored-by: tidely <43219534+tidely@users.noreply.github.com>
Co-authored-by: David <17435126+DavidLMS@users.noreply.github.com>

lerobot/common/policies/diffusion/modeling_diffusion.py

@@ -33,7 +33,7 @@ from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 from torch import Tensor, nn
 
-from lerobot.common.constants import OBS_ENV_STATE, OBS_STATE
+from lerobot.common.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.common.policies.normalize import Normalize, Unnormalize
 from lerobot.common.policies.pretrained import PreTrainedPolicy
@@ -99,6 +99,18 @@ class DiffusionPolicy(PreTrainedPolicy):
         if self.config.env_state_feature:
             self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps)
 
+    @torch.no_grad
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """Predict a chunk of actions given environment observations."""
+        # stack n latest observations from the queue
+        batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
+        actions = self.diffusion.generate_actions(batch)
+
+        # TODO(rcadene): make above methods return output dictionary?
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
+        return actions
+
     @torch.no_grad
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select a single action given environment observations.
@@ -124,23 +136,15 @@ class DiffusionPolicy(PreTrainedPolicy):
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = torch.stack(
-                [batch[key] for key in self.config.image_features], dim=-4
-            )
+            batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
         # Note: It's important that this happens after stacking the images into a single key.
         self._queues = populate_queues(self._queues, batch)
 
-        if len(self._queues["action"]) == 0:
-            # stack n latest observations from the queue
-            batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
-            actions = self.diffusion.generate_actions(batch)
+        if len(self._queues[ACTION]) == 0:
+            actions = self.predict_action_chunk(batch)
+            self._queues[ACTION].extend(actions.transpose(0, 1))
 
-            # TODO(rcadene): make above methods return output dictionary?
-            actions = self.unnormalize_outputs({"action": actions})["action"]
-
-            self._queues["action"].extend(actions.transpose(0, 1))
-
-        action = self._queues["action"].popleft()
+        action = self._queues[ACTION].popleft()
         return action
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:
@@ -148,9 +152,7 @@ class DiffusionPolicy(PreTrainedPolicy):
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = torch.stack(
-                [batch[key] for key in self.config.image_features], dim=-4
-            )
+            batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
         batch = self.normalize_targets(batch)
         loss = self.diffusion.compute_loss(batch)
         # no output_dict so returning None