Add files using upload-large-folder tool

examples/aloha_real/Dockerfile

@@ -0,0 +1,70 @@
+# Dockerfile for the Aloha real environment.
+
+# Build the container:
+# docker build . -t aloha_real -f examples/aloha_real/Dockerfile
+
+# Run the container:
+# docker run --rm -it --network=host -v /dev:/dev -v .:/app --privileged aloha_real /bin/bash
+
+FROM ros:noetic-robot@sha256:0e12e4db836e78c74c4b04c6d16f185d9a18d2b13cf5580747efa075eb6dc6e0
+SHELL ["/bin/bash", "-c"]
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    cmake \
+    curl \
+    libffi-dev \
+    python3-rosdep \
+    python3-rosinstall \
+    python3-rosinstall-generator \
+    whiptail \
+    git \
+    wget \
+    openssh-client \
+    ros-noetic-cv-bridge \
+    ros-noetic-usb-cam \
+    ros-noetic-realsense2-camera \
+    keyboard-configuration
+
+WORKDIR /root
+RUN curl 'https://raw.githubusercontent.com/Interbotix/interbotix_ros_manipulators/main/interbotix_ros_xsarms/install/amd64/xsarm_amd64_install.sh' > xsarm_amd64_install.sh
+RUN chmod +x xsarm_amd64_install.sh
+RUN export TZ='America/Los_Angeles' && ./xsarm_amd64_install.sh -d noetic -n
+
+COPY ./third_party/aloha /root/interbotix_ws/src/aloha
+RUN cd /root/interbotix_ws && source /opt/ros/noetic/setup.sh && source /root/interbotix_ws/devel/setup.sh && catkin_make
+
+# Install python 3.10 because this ROS image comes with 3.8
+RUN mkdir /python && \
+    cd /python && \
+    wget https://www.python.org/ftp/python/3.10.14/Python-3.10.14.tgz && \
+    tar -zxvf Python-3.10.14.tgz && \
+    cd Python-3.10.14 && \
+    ls -lhR && \
+    ./configure --enable-optimizations && \
+    make install && \
+    echo 'alias python3="/usr/local/bin/python3.10"' >> ~/.bashrc && \
+    echo 'alias python="/usr/local/bin/python3.10"' >> ~/.bashrc && \
+    cd ~ && rm -rf /python && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY --from=ghcr.io/astral-sh/uv:0.5.6 /uv /bin/uv
+ENV UV_HTTP_TIMEOUT=120
+ENV UV_LINK_MODE=copy
+COPY ./examples/aloha_real/requirements.txt /tmp/requirements.txt
+COPY ./packages/openpi-client/pyproject.toml /tmp/openpi-client/pyproject.toml
+RUN uv pip sync --python 3.10 --system /tmp/requirements.txt /tmp/openpi-client/pyproject.toml
+
+ENV PYTHONPATH=/app:/app/src:/app/packages/openpi-client/src:/root/interbotix_ws/src/aloha/aloha_scripts:/root/interbotix_ws/src/aloha
+WORKDIR /app
+
+# Create an entrypoint script to run the setup commands, followed by the command passed in.
+RUN cat <<'EOF' > /usr/local/bin/entrypoint.sh
+#!/bin/bash
+source /opt/ros/noetic/setup.sh && source /root/interbotix_ws/devel/setup.sh && "$@"
+EOF
+RUN chmod +x /usr/local/bin/entrypoint.sh
+
+ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
+CMD ["python3", "/app/examples/aloha_real/main.py"]

examples/aloha_real/README.md

@@ -0,0 +1,126 @@
+# Run Aloha (Real Robot)
+
+This example demonstrates how to run with a real robot using an [ALOHA setup](https://github.com/tonyzhaozh/aloha). See [here](../../docs/remote_inference.md) for instructions on how to load checkpoints and run inference. We list the relevant checkpoint paths for each provided fine-tuned model below.
+
+## Prerequisites
+
+This repo uses a fork of the ALOHA repo, with very minor modifications to use Realsense cameras.
+
+1. Follow the [hardware installation instructions](https://github.com/tonyzhaozh/aloha?tab=readme-ov-file#hardware-installation) in the ALOHA repo.
+1. Modify the `third_party/aloha/aloha_scripts/realsense_publisher.py` file to use serial numbers for your cameras.
+
+## With Docker
+
+```bash
+export SERVER_ARGS="--env ALOHA --default_prompt='take the toast out of the toaster'"
+docker compose -f examples/aloha_real/compose.yml up --build
+```
+
+## Without Docker
+
+Terminal window 1:
+
+```bash
+# Create virtual environment
+uv venv --python 3.10 examples/aloha_real/.venv
+source examples/aloha_real/.venv/bin/activate
+uv pip sync examples/aloha_real/requirements.txt
+uv pip install -e packages/openpi-client
+
+# Run the robot
+python examples/aloha_real/main.py
+```
+
+Terminal window 2:
+
+```bash
+roslaunch --wait aloha ros_nodes.launch
+```
+
+Terminal window 3:
+
+```bash
+uv run scripts/serve_policy.py --env ALOHA --default_prompt='take the toast out of the toaster'
+```
+
+## **ALOHA Checkpoint Guide**
+
+
+The `pi0_base` model can be used in zero shot for a simple task on the ALOHA platform, and we additionally provide two example fine-tuned checkpoints, “fold the towel” and “open the tupperware and put the food on the plate,” which can perform more advanced tasks on the ALOHA.
+
+While we’ve found the policies to work in unseen conditions across multiple ALOHA stations, we provide some pointers here on how best to set up scenes to maximize the chance of policy success. We cover the prompts to use for the policies, objects we’ve seen it work well on, and well-represented initial state distributions. Running these policies in zero shot is still a very experimental feature, and there is no guarantee that they will work on your robot. The recommended way to use `pi0_base` is by finetuning with data from the target robot.
+
+
+---
+
+### **Toast Task**  
+
+This task involves the robot taking two pieces of toast out of a toaster and placing them on a plate.  
+
+- **Checkpoint path**: `s3://openpi-assets/checkpoints/pi0_base`
+- **Prompt**: "take the toast out of the toaster"
+- **Objects needed**: Two pieces of toast, a plate, and a standard toaster.  
+- **Object Distribution**:  
+  - Works on both real toast and rubber fake toast  
+  - Compatible with standard 2-slice toasters  
+  - Works with plates of varying colors  
+
+### **Scene Setup Guidelines**
+<img width="500" alt="Screenshot 2025-01-31 at 10 06 02 PM" src="https://github.com/user-attachments/assets/3d043d95-9d1c-4dda-9991-e63cae61e02e" />
+
+- The toaster should be positioned in the top-left quadrant of the workspace.  
+- Both pieces of toast should start inside the toaster, with at least 1 cm of bread sticking out from the top.  
+- The plate should be placed roughly in the lower-center of the workspace.  
+- Works with both natural and synthetic lighting, but avoid making the scene too dark (e.g., don't place the setup inside an enclosed space or under a curtain).  
+
+
+### **Towel Task**  
+
+This task involves folding a small towel (e.g., roughly the size of a hand towel) into eighths.
+
+- **Checkpoint path**: `s3://openpi-assets/checkpoints/pi0_aloha_towel`
+- **Prompt**: "fold the towel"  
+- **Object Distribution**:  
+  - Works on towels of varying solid colors 
+  - Performance is worse on heavily textured or striped towels 
+
+### **Scene Setup Guidelines**  
+<img width="500" alt="Screenshot 2025-01-31 at 10 01 15 PM" src="https://github.com/user-attachments/assets/9410090c-467d-4a9c-ac76-96e5b4d00943" />
+
+- The towel should be flattened and roughly centered on the table.  
+- Choose a towel that does not blend in with the table surface.  
+
+
+### **Tupperware Task**  
+
+This task involves opening a tupperware filled with food and pouring the contents onto a plate.  
+
+- **Checkpoint path**: `s3://openpi-assets/checkpoints/pi0_aloha_tupperware`
+- **Prompt**: "open the tupperware and put the food on the plate"
+- **Objects needed**: Tupperware, food (or food-like items), and a plate.  
+- **Object Distribution**:  
+  - Works on various types of fake food (e.g., fake chicken nuggets, fries, and fried chicken).  
+  - Compatible with tupperware of different lid colors and shapes, with best performance on square tupperware with a corner flap (see images below).  
+  - The policy has seen plates of varying solid colors.  
+
+### **Scene Setup Guidelines** 
+<img width="500" alt="Screenshot 2025-01-31 at 10 02 27 PM" src="https://github.com/user-attachments/assets/60fc1de0-2d64-4076-b903-f427e5e9d1bf" />
+
+- Best performance observed when both the tupperware and plate are roughly centered in the workspace.  
+- Positioning:  
+  - Tupperware should be on the left.  
+  - Plate should be on the right or bottom.  
+  - The tupperware flap should point toward the plate.  
+
+## Training on your own Aloha dataset
+
+1. Convert the dataset to the LeRobot dataset v2.0 format. 
+    
+    We provide a script [convert_aloha_data_to_lerobot.py](./convert_aloha_data_to_lerobot.py) that converts the dataset to the LeRobot dataset v2.0 format. As an example we have converted the `aloha_pen_uncap_diverse_raw` dataset from the [BiPlay repo](https://huggingface.co/datasets/oier-mees/BiPlay/tree/main/aloha_pen_uncap_diverse_raw) and uploaded it to the HuggingFace Hub as [physical-intelligence/aloha_pen_uncap_diverse](https://huggingface.co/datasets/physical-intelligence/aloha_pen_uncap_diverse). 
+
+
+2. Define a training config that uses the custom dataset. 
+
+    We provide the [pi0_aloha_pen_uncap config](../../src/openpi/training/config.py) as an example. You should refer to the root [README](../../README.md) for how to run training with the new config.
+ 
+IMPORTANT: Our base checkpoint includes normalization stats from various common robot configurations. When fine-tuning a base checkpoint with a custom dataset from one of these configurations, we recommend using the corresponding normalization stats provided in the base checkpoint. In the example, this is done by specifying the trossen asset_id and a path to the pretrained checkpoint’s asset directory within the AssetsConfig.

packages/openpi-client/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "openpi-client"
+version = "0.1.0"
+requires-python = ">=3.7"
+dependencies = [
+    "dm-tree>=0.1.8",
+    "msgpack>=1.0.5",
+    "numpy>=1.21.6",
+    "pillow>=9.0.0",
+    "tree>=0.2.4",
+    "websockets>=11.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.uv]
+dev-dependencies = [
+    "pytest>=8.3.4",
+]
+
+[tool.ruff]
+line-length = 120
+target-version = "py37"

packages/openpi-client/src/openpi_client/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

packages/openpi-client/src/openpi_client/action_chunk_broker.py

@@ -0,0 +1,45 @@
+from typing import Dict
+
+import numpy as np
+import tree
+from typing_extensions import override
+
+from openpi_client import base_policy as _base_policy
+
+
+class ActionChunkBroker(_base_policy.BasePolicy):
+    """Wraps a policy to return action chunks one-at-a-time.
+
+    Assumes that the first dimension of all action fields is the chunk size.
+
+    A new inference call to the inner policy is only made when the current
+    list of chunks is exhausted.
+    """
+
+    def __init__(self, policy: _base_policy.BasePolicy, action_horizon: int):
+        self._policy = policy
+
+        self._action_horizon = action_horizon
+        self._cur_step: int = 0
+
+        self._last_results: Dict[str, np.ndarray] | None = None
+
+    @override
+    def infer(self, obs: Dict) -> Dict:  # noqa: UP006
+        if self._last_results is None:
+            self._last_results = self._policy.infer(obs)
+            self._cur_step = 0
+
+        results = tree.map_structure(lambda x: x[self._cur_step, ...], self._last_results)
+        self._cur_step += 1
+
+        if self._cur_step >= self._action_horizon:
+            self._last_results = None
+
+        return results
+
+    @override
+    def reset(self) -> None:
+        self._policy.reset()
+        self._last_results = None
+        self._cur_step = 0

packages/openpi-client/src/openpi_client/base_policy.py

@@ -0,0 +1,12 @@
+import abc
+from typing import Dict
+
+
+class BasePolicy(abc.ABC):
+    @abc.abstractmethod
+    def infer(self, obs: Dict) -> Dict:
+        """Infer actions from observations."""
+
+    def reset(self) -> None:
+        """Reset the policy to its initial state."""
+        pass

packages/openpi-client/src/openpi_client/image_tools.py

@@ -0,0 +1,58 @@
+import numpy as np
+from PIL import Image
+
+
+def convert_to_uint8(img: np.ndarray) -> np.ndarray:
+    """Converts an image to uint8 if it is a float image.
+
+    This is important for reducing the size of the image when sending it over the network.
+    """
+    if np.issubdtype(img.dtype, np.floating):
+        img = (255 * img).astype(np.uint8)
+    return img
+
+
+def resize_with_pad(images: np.ndarray, height: int, width: int, method=Image.BILINEAR) -> np.ndarray:
+    """Replicates tf.image.resize_with_pad for multiple images using PIL. Resizes a batch of images to a target height.
+
+    Args:
+        images: A batch of images in [..., height, width, channel] format.
+        height: The target height of the image.
+        width: The target width of the image.
+        method: The interpolation method to use. Default is bilinear.
+
+    Returns:
+        The resized images in [..., height, width, channel].
+    """
+    # If the images are already the correct size, return them as is.
+    if images.shape[-3:-1] == (height, width):
+        return images
+
+    original_shape = images.shape
+
+    images = images.reshape(-1, *original_shape[-3:])
+    resized = np.stack([_resize_with_pad_pil(Image.fromarray(im), height, width, method=method) for im in images])
+    return resized.reshape(*original_shape[:-3], *resized.shape[-3:])
+
+
+def _resize_with_pad_pil(image: Image.Image, height: int, width: int, method: int) -> Image.Image:
+    """Replicates tf.image.resize_with_pad for one image using PIL. Resizes an image to a target height and
+    width without distortion by padding with zeros.
+
+    Unlike the jax version, note that PIL uses [width, height, channel] ordering instead of [batch, h, w, c].
+    """
+    cur_width, cur_height = image.size
+    if cur_width == width and cur_height == height:
+        return image  # No need to resize if the image is already the correct size.
+
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+    resized_image = image.resize((resized_width, resized_height), resample=method)
+
+    zero_image = Image.new(resized_image.mode, (width, height), 0)
+    pad_height = max(0, int((height - resized_height) / 2))
+    pad_width = max(0, int((width - resized_width) / 2))
+    zero_image.paste(resized_image, (pad_width, pad_height))
+    assert zero_image.size == (width, height)
+    return zero_image

packages/openpi-client/src/openpi_client/image_tools_test.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+import openpi_client.image_tools as image_tools
+
+
+def test_resize_with_pad_shapes():
+    # Test case 1: Resize image with larger dimensions
+    images = np.zeros((2, 10, 10, 3), dtype=np.uint8)  # Input images of shape (batch_size, height, width, channels)
+    height = 20
+    width = 20
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (2, height, width, 3)
+    assert np.all(resized_images == 0)
+
+    # Test case 2: Resize image with smaller dimensions
+    images = np.zeros((3, 30, 30, 3), dtype=np.uint8)
+    height = 15
+    width = 15
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (3, height, width, 3)
+    assert np.all(resized_images == 0)
+
+    # Test case 3: Resize image with the same dimensions
+    images = np.zeros((1, 50, 50, 3), dtype=np.uint8)
+    height = 50
+    width = 50
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (1, height, width, 3)
+    assert np.all(resized_images == 0)
+
+    # Test case 3: Resize image with odd-numbered padding
+    images = np.zeros((1, 256, 320, 3), dtype=np.uint8)
+    height = 60
+    width = 80
+    resized_images = image_tools.resize_with_pad(images, height, width)
+    assert resized_images.shape == (1, height, width, 3)
+    assert np.all(resized_images == 0)

packages/openpi-client/src/openpi_client/msgpack_numpy.py

@@ -0,0 +1,57 @@
+"""Adds NumPy array support to msgpack.
+
+msgpack is good for (de)serializing data over a network for multiple reasons:
+- msgpack is secure (as opposed to pickle/dill/etc which allow for arbitrary code execution)
+- msgpack is widely used and has good cross-language support
+- msgpack does not require a schema (as opposed to protobuf/flatbuffers/etc) which is convenient in dynamically typed
+    languages like Python and JavaScript
+- msgpack is fast and efficient (as opposed to readable formats like JSON/YAML/etc); I found that msgpack was ~4x faster
+    than pickle for serializing large arrays using the below strategy
+
+The code below is adapted from https://github.com/lebedov/msgpack-numpy. The reason not to use that library directly is
+that it falls back to pickle for object arrays.
+"""
+
+import functools
+
+import msgpack
+import numpy as np
+
+
+def pack_array(obj):
+    if (isinstance(obj, (np.ndarray, np.generic))) and obj.dtype.kind in ("V", "O", "c"):
+        raise ValueError(f"Unsupported dtype: {obj.dtype}")
+
+    if isinstance(obj, np.ndarray):
+        return {
+            b"__ndarray__": True,
+            b"data": obj.tobytes(),
+            b"dtype": obj.dtype.str,
+            b"shape": obj.shape,
+        }
+
+    if isinstance(obj, np.generic):
+        return {
+            b"__npgeneric__": True,
+            b"data": obj.item(),
+            b"dtype": obj.dtype.str,
+        }
+
+    return obj
+
+
+def unpack_array(obj):
+    if b"__ndarray__" in obj:
+        return np.ndarray(buffer=obj[b"data"], dtype=np.dtype(obj[b"dtype"]), shape=obj[b"shape"])
+
+    if b"__npgeneric__" in obj:
+        return np.dtype(obj[b"dtype"]).type(obj[b"data"])
+
+    return obj
+
+
+Packer = functools.partial(msgpack.Packer, default=pack_array)
+packb = functools.partial(msgpack.packb, default=pack_array)
+
+Unpacker = functools.partial(msgpack.Unpacker, object_hook=unpack_array)
+unpackb = functools.partial(msgpack.unpackb, object_hook=unpack_array)

packages/openpi-client/src/openpi_client/runtime/agent.py

@@ -0,0 +1,17 @@
+import abc
+
+
+class Agent(abc.ABC):
+    """An Agent is the thing with agency, i.e. the entity that makes decisions.
+
+    Agents receive observations about the state of the world, and return actions
+    to take in response.
+    """
+
+    @abc.abstractmethod
+    def get_action(self, observation: dict) -> dict:
+        """Query the agent for the next action."""
+
+    @abc.abstractmethod
+    def reset(self) -> None:
+        """Reset the agent to its initial state."""

packages/openpi-client/src/openpi_client/runtime/environment.py

@@ -0,0 +1,32 @@
+import abc
+
+
+class Environment(abc.ABC):
+    """An Environment represents the robot and the environment it inhabits.
+
+    The primary contract of environments is that they can be queried for observations
+    about their state, and have actions applied to them to change that state.
+    """
+
+    @abc.abstractmethod
+    def reset(self) -> None:
+        """Reset the environment to its initial state.
+
+        This will be called once before starting each episode.
+        """
+
+    @abc.abstractmethod
+    def is_episode_complete(self) -> bool:
+        """Allow the environment to signal that the episode is complete.
+
+        This will be called after each step. It should return `True` if the episode is
+        complete (either successfully or unsuccessfully), and `False` otherwise.
+        """
+
+    @abc.abstractmethod
+    def get_observation(self) -> dict:
+        """Query the environment for the current state."""
+
+    @abc.abstractmethod
+    def apply_action(self, action: dict) -> None:
+        """Take an action in the environment."""

packages/openpi-client/src/openpi_client/runtime/runtime.py

@@ -0,0 +1,92 @@
+import logging
+import threading
+import time
+
+from openpi_client.runtime import agent as _agent
+from openpi_client.runtime import environment as _environment
+from openpi_client.runtime import subscriber as _subscriber
+
+
+class Runtime:
+    """The core module orchestrating interactions between key components of the system."""
+
+    def __init__(
+        self,
+        environment: _environment.Environment,
+        agent: _agent.Agent,
+        subscribers: list[_subscriber.Subscriber],
+        max_hz: float = 0,
+        num_episodes: int = 1,
+        max_episode_steps: int = 0,
+    ) -> None:
+        self._environment = environment
+        self._agent = agent
+        self._subscribers = subscribers
+        self._max_hz = max_hz
+        self._num_episodes = num_episodes
+        self._max_episode_steps = max_episode_steps
+
+        self._in_episode = False
+        self._episode_steps = 0
+
+    def run(self) -> None:
+        """Runs the runtime loop continuously until stop() is called or the environment is done."""
+        for _ in range(self._num_episodes):
+            self._run_episode()
+
+        # Final reset, this is important for real environments to move the robot to its home position.
+        self._environment.reset()
+
+    def run_in_new_thread(self) -> threading.Thread:
+        """Runs the runtime loop in a new thread."""
+        thread = threading.Thread(target=self.run)
+        thread.start()
+        return thread
+
+    def mark_episode_complete(self) -> None:
+        """Marks the end of an episode."""
+        self._in_episode = False
+
+    def _run_episode(self) -> None:
+        """Runs a single episode."""
+        logging.info("Starting episode...")
+        self._environment.reset()
+        self._agent.reset()
+        for subscriber in self._subscribers:
+            subscriber.on_episode_start()
+
+        self._in_episode = True
+        self._episode_steps = 0
+        step_time = 1 / self._max_hz if self._max_hz > 0 else 0
+        last_step_time = time.time()
+
+        while self._in_episode:
+            self._step()
+            self._episode_steps += 1
+
+            # Sleep to maintain the desired frame rate
+            now = time.time()
+            dt = now - last_step_time
+            if dt < step_time:
+                time.sleep(step_time - dt)
+                last_step_time = time.time()
+            else:
+                last_step_time = now
+
+        logging.info("Episode completed.")
+        for subscriber in self._subscribers:
+            subscriber.on_episode_end()
+
+    def _step(self) -> None:
+        """A single step of the runtime loop."""
+        observation = self._environment.get_observation()
+        action = self._agent.get_action(observation)
+        self._environment.apply_action(action)
+
+        for subscriber in self._subscribers:
+            subscriber.on_step(observation, action)
+
+        if self._environment.is_episode_complete() or (
+            self._max_episode_steps > 0 and self._episode_steps >= self._max_episode_steps
+        ):
+            self.mark_episode_complete()

packages/openpi-client/src/openpi_client/runtime/subscriber.py

@@ -0,0 +1,20 @@
+import abc
+
+
+class Subscriber(abc.ABC):
+    """Subscribes to events in the runtime.
+
+    Subscribers can be used to save data, visualize, etc.
+    """
+
+    @abc.abstractmethod
+    def on_episode_start(self) -> None:
+        """Called when an episode starts."""
+
+    @abc.abstractmethod
+    def on_step(self, observation: dict, action: dict) -> None:
+        """Append a step to the episode."""
+
+    @abc.abstractmethod
+    def on_episode_end(self) -> None:
+        """Called when an episode ends."""

packages/openpi-client/src/openpi_client/websocket_client_policy.py

@@ -0,0 +1,49 @@
+import logging
+import time
+from typing import Dict, Tuple
+
+import websockets.sync.client
+from typing_extensions import override
+
+from openpi_client import base_policy as _base_policy
+from openpi_client import msgpack_numpy
+
+
+class WebsocketClientPolicy(_base_policy.BasePolicy):
+    """Implements the Policy interface by communicating with a server over websocket.
+
+    See WebsocketPolicyServer for a corresponding server implementation.
+    """
+
+    def __init__(self, host: str = "0.0.0.0", port: int = 8000) -> None:
+        self._uri = f"ws://{host}:{port}"
+        self._packer = msgpack_numpy.Packer()
+        self._ws, self._server_metadata = self._wait_for_server()
+
+    def get_server_metadata(self) -> Dict:
+        return self._server_metadata
+
+    def _wait_for_server(self) -> Tuple[websockets.sync.client.ClientConnection, Dict]:
+        logging.info(f"Waiting for server at {self._uri}...")
+        while True:
+            try:
+                conn = websockets.sync.client.connect(self._uri, compression=None, max_size=None)
+                metadata = msgpack_numpy.unpackb(conn.recv())
+                return conn, metadata
+            except ConnectionRefusedError:
+                logging.info("Still waiting for server...")
+                time.sleep(5)
+
+    @override
+    def infer(self, obs: Dict) -> Dict:  # noqa: UP006
+        data = self._packer.pack(obs)
+        self._ws.send(data)
+        response = self._ws.recv()
+        if isinstance(response, str):
+            # we're expecting bytes; if the server sends a string, it's an error.
+            raise RuntimeError(f"Error in inference server:\n{response}")
+        return msgpack_numpy.unpackb(response)
+
+    @override
+    def reset(self) -> None:
+        pass

scripts/docker/compose.yml

@@ -0,0 +1,29 @@
+# Run with:
+# docker compose -f scripts/compose.yml up --build
+services:
+  openpi_server:
+    image: openpi_server
+    build:
+      context: ..
+      dockerfile: scripts/docker/serve_policy.Dockerfile
+    init: true
+    tty: true
+    network_mode: host
+    # Populate configured openpi data home to /openpi_assets inside the container.
+    # Populate aws credential inside the container.
+    volumes:
+      - $PWD:/app
+      - ${OPENPI_DATA_HOME:-~/.cache/openpi}:/openpi_assets
+    environment:
+      - SERVER_ARGS
+      - OPENPI_DATA_HOME=/openpi_assets
+      - IS_DOCKER=true
+
+    # Comment out this block if not running on a machine with GPUs.
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

scripts/docker/install_docker_ubuntu22.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+# Add Docker's official GPG key:
+sudo apt-get update
+sudo apt-get install -y ca-certificates curl
+sudo install -m 0755 -d /etc/apt/keyrings
+sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
+sudo chmod a+r /etc/apt/keyrings/docker.asc
+
+# Add the repository to Apt sources:
+echo \
+	"deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \
+  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" |
+	sudo tee /etc/apt/sources.list.d/docker.list >/dev/null
+sudo apt-get update
+
+sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
+
+# Add current user to the 'docker' group, which allows them to use docker commands (docker build, docker run, etc).
+# See https://docs.docker.com/engine/install/linux-postinstall/
+username=$(whoami)
+sudo usermod -aG docker $username
+
+# Configure docker to start automatically on system boot.
+sudo systemctl enable docker.service
+sudo systemctl enable containerd.service
+
+# https://forums.docker.com/t/docker-credential-desktop-exe-executable-file-not-found-in-path-using-wsl2/100225/5
+if [ ~/.docker/config.json ]; then
+	sed -i 's/credsStore/credStore/g' ~/.docker/config.json
+fi
+
+echo ""
+echo "********************************************************************"
+echo "**** Restart to allow Docker permission changes to take effect. ****"
+echo "********************************************************************"
+echo ""

scripts/docker/install_nvidia_container_toolkit.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+# Installs the NVIDIA Container Toolkit, which allows Docker containers to access NVIDIA GPUs.
+# NVIDIA's official documentation: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html
+
+curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg &&
+	curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list |
+	sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' |
+		sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
+
+# NVIDIA's documenation omits 'sudo' in the following command, but it is required.
+sudo sed -i -e '/experimental/ s/^#//g' /etc/apt/sources.list.d/nvidia-container-toolkit.list
+sudo apt-get update
+sudo apt-get install -y nvidia-container-toolkit
+
+sudo nvidia-ctk runtime configure --runtime=docker
+sudo systemctl restart docker

scripts/docker/serve_policy.Dockerfile

@@ -0,0 +1,34 @@
+# Dockerfile for serving a PI policy.
+# Based on UV's instructions: https://docs.astral.sh/uv/guides/integration/docker/#developing-in-a-container
+
+# Build the container:
+# docker build . -t openpi_server -f scripts/docker/serve_policy.Dockerfile
+
+# Run the container:
+# docker run --rm -it --network=host -v .:/app --gpus=all openpi_server /bin/bash
+
+FROM nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04@sha256:2d913b09e6be8387e1a10976933642c73c840c0b735f0bf3c28d97fc9bc422e0
+COPY --from=ghcr.io/astral-sh/uv:0.5.1 /uv /uvx /bin/
+
+WORKDIR /app
+
+# Needed because LeRobot uses git-lfs.
+RUN apt-get update && apt-get install -y git git-lfs
+
+# Copy from the cache instead of linking since it's a mounted volume
+ENV UV_LINK_MODE=copy
+
+# Write the virtual environment outside of the project directory so it doesn't
+# leak out of the container when we mount the application code.
+ENV UV_PROJECT_ENVIRONMENT=/.venv
+
+# Install the project's dependencies using the lockfile and settings
+RUN uv venv --python 3.11.9 $UV_PROJECT_ENVIRONMENT
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    --mount=type=bind,source=packages/openpi-client/pyproject.toml,target=packages/openpi-client/pyproject.toml \
+    --mount=type=bind,source=packages/openpi-client/src,target=packages/openpi-client/src \
+    GIT_LFS_SKIP_SMUDGE=1 uv sync --frozen --no-install-project --no-dev
+
+CMD /bin/bash -c "uv run scripts/serve_policy.py $SERVER_ARGS"

scripts/logs/eval_pi0fast_libero.out

@@ -0,0 +1,1243 @@
+================================================================================
+JobID = 12563538
+User = u24877, Account = kisski-spath
+Partition = kisski, Nodelist = ggpu170
+================================================================================
+===========================================
+Experiment: eval_pi0fast_libero
+Codebase: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid
+Checkpoint: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/outputs/train/2026-02-08/07-52-25_pi0fast_baseline_libero_30k/checkpoints/100000/pretrained_model
+Video Output: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/results/eval_pi0fast_libero/videos
+===========================================
+[robosuite WARNING] No private macro file found! (__init__.py:7)
+[robosuite WARNING] It is recommended to use a private macro file (__init__.py:8)
+[robosuite WARNING] To setup, run: python /user/hominhduy.nguyen01/u24877/.local/lib/python3.10/site-packages/robosuite/scripts/setup_macros.py (__init__.py:9)
+Gym has been unmaintained since 2022 and does not support NumPy 2.0 amongst other critical functionality.
+Please upgrade to Gymnasium, the maintained drop-in replacement of Gym, or contact the authors of your software and request that they upgrade.
+Users of this version of Gym should be able to simply replace 'import gym' with 'import gymnasium as gym' in the vast majority of cases.
+See the migration guide at https://gymnasium.farama.org/introduction/migration_guide/ for additional information.
+INFO:root:Task suite: libero_10
+<frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute
+<frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute
+Some kwargs in processor config are unused and will not have any effect: scale, action_dim, time_horizon, min_token, vocab_size. 
+/projects/extern/kisski/kisski-spath/dir.project/VLA_3D/miniconda3/envs/pitorch_doanh/lib/python3.10/site-packages/transformers/models/paligemma/configuration_paligemma.py:134: FutureWarning: The `ignore_index` attribute is deprecated and will be removed in v4.47.
+  warnings.warn(
+/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/libero_config/libero_config.yaml
+[info] using task orders [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+Loading weights from local directory
+[INFO] Missing keys (will be randomly initialized): ['model.pi0_paligemma.language_model.model.embed_tokens.weight']
+
  0%|          | 0/10 [00:00<?, ?it/s][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 1...
+You're using a GemmaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 1
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:40<06:08, 40.97s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 2
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:17<05:07, 38.42s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 3
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:54<04:24, 37.83s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 4
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:33<03:48, 38.11s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 5
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [03:11<03:11, 38.32s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 6
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:50<02:34, 38.56s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 7
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [04:28<01:54, 38.17s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 8
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [05:09<01:18, 39.02s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 9
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [05:49<00:39, 39.29s/it]INFO:root:
+Task: put both the alphabet soup and the tomato sauce in the basket
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 10
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 10%|█         | 1/10 [06:31<58:44, 391.59s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 11
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:38<05:44, 38.24s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 12
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:16<05:04, 38.11s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 13
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:55<04:29, 38.45s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 14
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:33<03:50, 38.38s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 15
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [03:11<03:11, 38.37s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 16
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:50<02:34, 38.58s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 17
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [04:29<01:55, 38.50s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 18
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [05:07<01:16, 38.43s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 19
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [05:46<00:38, 38.71s/it]INFO:root:
+Task: put both the cream cheese box and the butter in the basket
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 20
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 20%|██        | 2/10 [12:57<51:47, 388.42s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 21
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:31<04:47, 31.93s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 22
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:03<04:12, 31.57s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 23
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:34<03:39, 31.35s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 24
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:02<03:00, 30.07s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 25
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [02:38<02:40, 32.18s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 26
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:12<02:11, 32.96s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 27
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [03:46<01:40, 33.34s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 28
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [04:21<01:07, 33.60s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 29
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [04:51<00:32, 32.60s/it]INFO:root:
+Task: turn on the stove and put the moka pot on it
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 30
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 30%|███       | 3/10 [18:19<41:45, 357.93s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 31
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:31<04:46, 31.84s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 32
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:02<04:10, 31.26s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 33
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:34<03:39, 31.32s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 34
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:03<03:04, 30.71s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 35
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [02:36<02:37, 31.57s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 36
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:07<02:05, 31.33s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 37
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [03:39<01:34, 31.34s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 38
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [04:11<01:03, 31.67s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 39
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [04:43<00:31, 31.76s/it]INFO:root:
+Task: put the black bowl in the bottom drawer of the cabinet and close it
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 40
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 40%|████      | 4/10 [23:33<34:04, 340.77s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 41
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:35<05:22, 35.86s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 42
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:15<05:05, 38.18s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 43
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:55<04:32, 38.95s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 44
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:31<03:46, 37.70s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 45
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [03:09<03:09, 37.82s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 46
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:47<02:32, 38.03s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 47
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [04:28<01:56, 38.90s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 48
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [05:07<01:17, 38.94s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 49
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [05:46<00:38, 38.99s/it]INFO:root:
+Task: put the white mug on the left plate and put the yellow and white mug on the right plate
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 50
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 50%|█████     | 5/10 [29:57<29:40, 356.10s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 51
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:33<05:02, 33.65s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 52
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:05<04:21, 32.66s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 53
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:39<03:51, 33.06s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 54
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:11<03:16, 32.69s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 55
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [02:42<02:41, 32.23s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 56
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:15<02:09, 32.46s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 57
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [03:49<01:38, 32.82s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 58
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [04:21<01:05, 32.52s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 59
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [04:54<00:32, 32.74s/it]INFO:root:
+Task: pick up the book and place it in the back compartment of the caddy
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 60
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 60%|██████    | 6/10 [35:25<23:06, 346.70s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 61
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:34<05:09, 34.37s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 62
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:08<04:35, 34.44s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 63
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:42<03:57, 33.94s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 64
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:16<03:24, 34.01s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 65
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [02:53<02:55, 35.07s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 66
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:27<02:19, 34.87s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 67
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [04:05<01:47, 35.78s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 68
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [04:39<01:10, 35.26s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 69
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [05:15<00:35, 35.41s/it]INFO:root:
+Task: put the white mug on the plate and put the chocolate pudding to the right of the plate
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 70
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 70%|███████   | 7/10 [41:16<17:24, 348.18s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 71
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:37<05:38, 37.63s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 72
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [01:14<04:55, 36.90s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 73
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:50<04:18, 36.92s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 74
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [02:28<03:42, 37.15s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 75
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [03:06<03:07, 37.55s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 76
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:43<02:29, 37.33s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 77
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [04:21<01:52, 37.54s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 78
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [04:57<01:14, 37.14s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 79
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [05:33<00:36, 36.74s/it]INFO:root:
+Task: put both the alphabet soup and the cream cheese box in the basket
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 80
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 80%|████████  | 8/10 [47:29<11:51, 355.92s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 81
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:27<04:03, 27.01s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 82
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [00:54<03:36, 27.05s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 83
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:21<03:09, 27.08s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 84
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [01:48<02:41, 26.97s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 85
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [02:14<02:14, 26.97s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 86
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [02:43<01:49, 27.40s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 87
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [03:09<01:21, 27.19s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 88
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [03:37<00:54, 27.34s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 89
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [04:04<00:27, 27.32s/it]INFO:root:
+Task: put both moka pots on the stove
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 90
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
 90%|█████████ | 9/10 [52:02<05:29, 329.92s/it][Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+[Warning]: datasets path /mnt/vast-kisski/projects/kisski-spath/VLA_objectcentric/VLA-Humanoid/libero/libero/../datasets does not exist!
+
+
  0%|          | 0/10 [00:00<?, ?it/s]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 1...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 91
+INFO:root:# successes: 0 (0.0%)
+
+
 10%|█         | 1/10 [00:29<04:25, 29.54s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 2...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 92
+INFO:root:# successes: 0 (0.0%)
+
+
 20%|██        | 2/10 [00:59<03:58, 29.84s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 3...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 93
+INFO:root:# successes: 0 (0.0%)
+
+
 30%|███       | 3/10 [01:29<03:30, 30.07s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 4...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 94
+INFO:root:# successes: 0 (0.0%)
+
+
 40%|████      | 4/10 [01:59<02:59, 29.99s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 5...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 95
+INFO:root:# successes: 0 (0.0%)
+
+
 50%|█████     | 5/10 [02:30<02:31, 30.31s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 6...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 96
+INFO:root:# successes: 0 (0.0%)
+
+
 60%|██████    | 6/10 [03:01<02:02, 30.60s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 7...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 97
+INFO:root:# successes: 0 (0.0%)
+
+
 70%|███████   | 7/10 [03:32<01:31, 30.64s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 8...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 98
+INFO:root:# successes: 0 (0.0%)
+
+
 80%|████████  | 8/10 [04:02<01:00, 30.34s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 9...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 99
+INFO:root:# successes: 0 (0.0%)
+
+
 90%|█████████ | 9/10 [04:31<00:30, 30.02s/it]INFO:root:
+Task: put the yellow and white mug in the microwave and close it
+INFO:root:Starting episode 10...
+huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+To disable this warning, you can either:
+	- Avoid using `tokenizers` before the fork if possible
+	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
+INFO:root:Success: False
+INFO:root:# episodes completed so far: 100
+INFO:root:# successes: 0 (0.0%)
+
+
+INFO:root:Current task success rate: 0.0
+INFO:root:Current total success rate: 0.0
+
+INFO:root:Total success rate: 0.0
+INFO:root:Total episodes: 100
+Exception ignored in: <function MjRenderContext.__del__ at 0x14cd11c11ea0>
+Traceback (most recent call last):
+  File "/user/hominhduy.nguyen01/u24877/.local/lib/python3.10/site-packages/robosuite/utils/binding_utils.py", line 199, in __del__
+    self.gl_ctx.free()
+  File "/user/hominhduy.nguyen01/u24877/.local/lib/python3.10/site-packages/robosuite/renderers/context/egl_context.py", line 149, in free
+    EGL.eglMakeCurrent(EGL_DISPLAY, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, EGL.EGL_NO_CONTEXT)
+  File "/user/hominhduy.nguyen01/u24877/.local/lib/python3.10/site-packages/OpenGL/error.py", line 230, in glCheckError
+    raise self._errorClass(
+OpenGL.raw.EGL._errors.EGLError: EGLError(
+	err = EGL_NOT_INITIALIZED,
+	baseOperation = eglMakeCurrent,
+	cArguments = (
+		<OpenGL._opaque.EGLDisplay_pointer object at 0x14ccc11caa40>,
+		<OpenGL._opaque.EGLSurface_pointer object at 0x14cd11b98bc0>,
+		<OpenGL._opaque.EGLSurface_pointer object at 0x14cd11b98bc0>,
+		<OpenGL._opaque.EGLContext_pointer object at 0x14cd11b98b40>,
+	),
+	result = 0
+)
+Exception ignored in: <function EGLGLContext.__del__ at 0x14cd11c11cf0>
+Traceback (most recent call last):
+  File "/user/hominhduy.nguyen01/u24877/.local/lib/python3.10/site-packages/robosuite/renderers/context/egl_context.py", line 155, in __del__
+    self.free()
+  File "/user/hominhduy.nguyen01/u24877/.local/lib/python3.10/site-packages/robosuite/renderers/context/egl_context.py", line 149, in free
+    EGL.eglMakeCurrent(EGL_DISPLAY, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, EGL.EGL_NO_CONTEXT)
+  File "/user/hominhduy.nguyen01/u24877/.local/lib/python3.10/site-packages/OpenGL/error.py", line 230, in glCheckError
+    raise self._errorClass(
+OpenGL.raw.EGL._errors.EGLError: EGLError(
+	err = EGL_NOT_INITIALIZED,
+	baseOperation = eglMakeCurrent,
+	cArguments = (
+		<OpenGL._opaque.EGLDisplay_pointer object at 0x14ccc11caa40>,
+		<OpenGL._opaque.EGLSurface_pointer object at 0x14cd11b98bc0>,
+		<OpenGL._opaque.EGLSurface_pointer object at 0x14cd11b98bc0>,
+		<OpenGL._opaque.EGLContext_pointer object at 0x14cd11b98b40>,
+	),
+	result = 0
+)
+sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute
+===========================================
+Evaluation completed!
+Videos saved to: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/results/eval_pi0fast_libero/videos
+===========================================
+============ Job Information ===================================================
+Submitted: 2026-02-10T09:23:42
+Started: 2026-02-10T09:24:05
+Ended: 2026-02-10T10:22:43
+Elapsed: 59 min, Limit: 240 min, Difference: 181 min
+CPUs: 16, Nodes: 1
+Estimated Consumption: 590.00 core-hours
+================================================================================

scripts/logs/pi0fast_baseline_libero_h100.out

@@ -0,0 +1,569 @@
+================================================================================
+JobID = 12538472
+User = u24877, Account = kisski-spath
+Partition = kisski, Nodelist = ggpu174
+================================================================================
+===========================================
+Experiment: pi0fast_baseline_libero_30k
+Codebase: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid
+Data: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10
+Policy Config Source: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/configs/policy_config/pi0fast_baseline.json
+Temp Config Dir: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/temp/configs/pi0fast_baseline_libero_30k
+Output Dir: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/outputs/train/2026-02-08/07-52-25_pi0fast_baseline_libero_30k
+===========================================
+{'batch_size': 4,
+ 'dataset': {'episodes': None,
+             'image_transforms': {'enable': True,
+                                  'image_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'crop_resize': {'kwargs': {'ratio': [1,
+                                                                                     1],
+                                                                           'scale': [0.9,
+                                                                                     0.95],
+                                                                           'size': [256,
+                                                                                    256]},
+                                                                'type': 'RandomResizedCrop',
+                                                                'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'rotate': {'kwargs': {'degrees': [-5,
+                                                                                  5]},
+                                                           'type': 'RandomRotate',
+                                                           'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}},
+                                  'max_num_transforms': 3,
+                                  'random_order': False,
+                                  'wrist_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}}},
+             'repo_id': '.',
+             'revision': None,
+             'root': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10',
+             'use_imagenet_stats': True,
+             'video_backend': 'pyav',
+             'vqa_data_path': None},
+ 'env': None,
+ 'eval': {'batch_size': 50, 'n_episodes': 50, 'use_async_envs': False},
+ 'eval_freq': 20000,
+ 'job_name': 'pi0fast_baseline_libero_30k',
+ 'log_freq': 100,
+ 'num_workers': 8,
+ 'optimizer': {'betas': [0.9, 0.95],
+               'eps': 1e-08,
+               'grad_clip_norm': 1,
+               'lr': 0.0001,
+               'type': 'adamw',
+               'weight_decay': 1e-05},
+ 'output_dir': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/outputs/train/2026-02-08/07-52-25_pi0fast_baseline_libero_30k',
+ 'policy': {'adapt_to_pi_aloha': False,
+            'checkpoint_path': None,
+            'chunk_size': 10,
+            'device': 'cuda',
+            'empty_cameras': 0,
+            'fast_skip_tokens': 128,
+            'fast_tokenizer_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/fast',
+            'freeze_lm_head': True,
+            'freeze_vision_encoder': True,
+            'grad_clip_norm': 1,
+            'gradient_accumulation_steps': 4,
+            'input_features': {},
+            'interpolate_like_pi': False,
+            'max_action_dim': 32,
+            'max_decoding_steps': 256,
+            'max_input_seq_len': 256,
+            'max_state_dim': 32,
+            'n_action_steps': 10,
+            'n_obs_steps': 1,
+            'normalization_mapping': {'ACTION': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'STATE': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'VISUAL': <NormalizationMode.IDENTITY: 'IDENTITY'>},
+            'optimizer_betas': [0.9, 0.95],
+            'optimizer_eps': 1e-08,
+            'optimizer_lr': 0.0001,
+            'optimizer_weight_decay': 1e-05,
+            'output_features': {},
+            'padding_side': 'right',
+            'pi0_paligemma_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/paligemma-3b-pt-224',
+            'precision': 'bfloat16',
+            'proj_width': 1024,
+            'relaxed_action_decoding': True,
+            'resize_imgs_with_padding': [224, 224],
+            'scheduler_decay_lr': 2.5e-06,
+            'scheduler_decay_steps': 30000,
+            'scheduler_warmup_steps': 1000,
+            'tokenizer_max_length': 48,
+            'type': 'pi0fast',
+            'use_amp': False,
+            'use_cache': True,
+            'use_delta_joint_actions_aloha': False},
+ 'resume': False,
+ 'save_checkpoint': True,
+ 'save_freq': 5000,
+ 'scheduler': {'decay_lr': 2.5e-06,
+               'num_decay_steps': 30000,
+               'num_warmup_steps': 1000,
+               'peak_lr': 0.0001,
+               'type': 'cosine_decay_with_warmup'},
+ 'seed': 42,
+ 'steps': 100000,
+ 'use_policy_training_preset': True,
+ 'wandb': {'disable_artifact': True,
+           'enable': False,
+           'entity': 'Robotics_VLA',
+           'mode': 'online',
+           'notes': None,
+           'project': 'pi0_lerobot',
+           'run_id': None}}
+{'batch_size': 4,
+ 'dataset': {'episodes': None,
+             'image_transforms': {'enable': True,
+                                  'image_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'crop_resize': {'kwargs': {'ratio': [1,
+                                                                                     1],
+                                                                           'scale': [0.9,
+                                                                                     0.95],
+                                                                           'size': [256,
+                                                                                    256]},
+                                                                'type': 'RandomResizedCrop',
+                                                                'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'rotate': {'kwargs': {'degrees': [-5,
+                                                                                  5]},
+                                                           'type': 'RandomRotate',
+                                                           'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}},
+                                  'max_num_transforms': 3,
+                                  'random_order': False,
+                                  'wrist_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}}},
+             'repo_id': '.',
+             'revision': None,
+             'root': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10',
+             'use_imagenet_stats': True,
+             'video_backend': 'pyav',
+             'vqa_data_path': None},
+ 'env': None,
+ 'eval': {'batch_size': 50, 'n_episodes': 50, 'use_async_envs': False},
+ 'eval_freq': 20000,
+ 'job_name': 'pi0fast_baseline_libero_30k',
+ 'log_freq': 100,
+ 'num_workers': 8,
+ 'optimizer': {'betas': [0.9, 0.95],
+               'eps': 1e-08,
+               'grad_clip_norm': 1,
+               'lr': 0.0001,
+               'type': 'adamw',
+               'weight_decay': 1e-05},
+ 'output_dir': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/outputs/train/2026-02-08/07-52-25_pi0fast_baseline_libero_30k',
+ 'policy': {'adapt_to_pi_aloha': False,
+            'checkpoint_path': None,
+            'chunk_size': 10,
+            'device': 'cuda',
+            'empty_cameras': 0,
+            'fast_skip_tokens': 128,
+            'fast_tokenizer_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/fast',
+            'freeze_lm_head': True,
+            'freeze_vision_encoder': True,
+            'grad_clip_norm': 1,
+            'gradient_accumulation_steps': 4,
+            'input_features': {},
+            'interpolate_like_pi': False,
+            'max_action_dim': 32,
+            'max_decoding_steps': 256,
+            'max_input_seq_len': 256,
+            'max_state_dim': 32,
+            'n_action_steps': 10,
+            'n_obs_steps': 1,
+            'normalization_mapping': {'ACTION': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'STATE': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'VISUAL': <NormalizationMode.IDENTITY: 'IDENTITY'>},
+            'optimizer_betas': [0.9, 0.95],
+            'optimizer_eps': 1e-08,
+            'optimizer_lr': 0.0001,
+            'optimizer_weight_decay': 1e-05,
+            'output_features': {},
+            'padding_side': 'right',
+            'pi0_paligemma_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/paligemma-3b-pt-224',
+            'precision': 'bfloat16',
+            'proj_width': 1024,
+            'relaxed_action_decoding': True,
+            'resize_imgs_with_padding': [224, 224],
+            'scheduler_decay_lr': 2.5e-06,
+            'scheduler_decay_steps': 30000,
+            'scheduler_warmup_steps': 1000,
+            'tokenizer_max_length': 48,
+            'type': 'pi0fast',
+            'use_amp': False,
+            'use_cache': True,
+            'use_delta_joint_actions_aloha': False},
+ 'resume': False,
+ 'save_checkpoint': True,
+ 'save_freq': 5000,
+ 'scheduler': {'decay_lr': 2.5e-06,
+               'num_decay_steps': 30000,
+               'num_warmup_steps': 1000,
+               'peak_lr': 0.0001,
+               'type': 'cosine_decay_with_warmup'},
+ 'seed': 42,
+ 'steps': 100000,
+ 'use_policy_training_preset': True,
+ 'wandb': {'disable_artifact': True,
+           'enable': False,
+           'entity': 'Robotics_VLA',
+           'mode': 'online',
+           'notes': None,
+           'project': 'pi0_lerobot',
+           'run_id': None}}
+{'batch_size': 4,
+ 'dataset': {'episodes': None,
+             'image_transforms': {'enable': True,
+                                  'image_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'crop_resize': {'kwargs': {'ratio': [1,
+                                                                                     1],
+                                                                           'scale': [0.9,
+                                                                                     0.95],
+                                                                           'size': [256,
+                                                                                    256]},
+                                                                'type': 'RandomResizedCrop',
+                                                                'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'rotate': {'kwargs': {'degrees': [-5,
+                                                                                  5]},
+                                                           'type': 'RandomRotate',
+                                                           'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}},
+                                  'max_num_transforms': 3,
+                                  'random_order': False,
+                                  'wrist_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}}},
+             'repo_id': '.',
+             'revision': None,
+             'root': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10',
+             'use_imagenet_stats': True,
+             'video_backend': 'pyav',
+             'vqa_data_path': None},
+ 'env': None,
+ 'eval': {'batch_size': 50, 'n_episodes': 50, 'use_async_envs': False},
+ 'eval_freq': 20000,
+ 'job_name': 'pi0fast_baseline_libero_30k',
+ 'log_freq': 100,
+ 'num_workers': 8,
+ 'optimizer': {'betas': [0.9, 0.95],
+               'eps': 1e-08,
+               'grad_clip_norm': 1,
+               'lr': 0.0001,
+               'type': 'adamw',
+               'weight_decay': 1e-05},
+ 'output_dir': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/outputs/train/2026-02-08/07-52-25_pi0fast_baseline_libero_30k',
+ 'policy': {'adapt_to_pi_aloha': False,
+            'checkpoint_path': None,
+            'chunk_size': 10,
+            'device': 'cuda',
+            'empty_cameras': 0,
+            'fast_skip_tokens': 128,
+            'fast_tokenizer_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/fast',
+            'freeze_lm_head': True,
+            'freeze_vision_encoder': True,
+            'grad_clip_norm': 1,
+            'gradient_accumulation_steps': 4,
+            'input_features': {},
+            'interpolate_like_pi': False,
+            'max_action_dim': 32,
+            'max_decoding_steps': 256,
+            'max_input_seq_len': 256,
+            'max_state_dim': 32,
+            'n_action_steps': 10,
+            'n_obs_steps': 1,
+            'normalization_mapping': {'ACTION': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'STATE': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'VISUAL': <NormalizationMode.IDENTITY: 'IDENTITY'>},
+            'optimizer_betas': [0.9, 0.95],
+            'optimizer_eps': 1e-08,
+            'optimizer_lr': 0.0001,
+            'optimizer_weight_decay': 1e-05,
+            'output_features': {},
+            'padding_side': 'right',
+            'pi0_paligemma_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/paligemma-3b-pt-224',
+            'precision': 'bfloat16',
+            'proj_width': 1024,
+            'relaxed_action_decoding': True,
+            'resize_imgs_with_padding': [224, 224],
+            'scheduler_decay_lr': 2.5e-06,
+            'scheduler_decay_steps': 30000,
+            'scheduler_warmup_steps': 1000,
+            'tokenizer_max_length': 48,
+            'type': 'pi0fast',
+            'use_amp': False,
+            'use_cache': True,
+            'use_delta_joint_actions_aloha': False},
+ 'resume': False,
+ 'save_checkpoint': True,
+ 'save_freq': 5000,
+ 'scheduler': {'decay_lr': 2.5e-06,
+               'num_decay_steps': 30000,
+               'num_warmup_steps': 1000,
+               'peak_lr': 0.0001,
+               'type': 'cosine_decay_with_warmup'},
+ 'seed': 42,
+ 'steps': 100000,
+ 'use_policy_training_preset': True,
+ 'wandb': {'disable_artifact': True,
+           'enable': False,
+           'entity': 'Robotics_VLA',
+           'mode': 'online',
+           'notes': None,
+           'project': 'pi0_lerobot',
+           'run_id': None}}
+{'batch_size': 4,
+ 'dataset': {'episodes': None,
+             'image_transforms': {'enable': True,
+                                  'image_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'crop_resize': {'kwargs': {'ratio': [1,
+                                                                                     1],
+                                                                           'scale': [0.9,
+                                                                                     0.95],
+                                                                           'size': [256,
+                                                                                    256]},
+                                                                'type': 'RandomResizedCrop',
+                                                                'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'rotate': {'kwargs': {'degrees': [-5,
+                                                                                  5]},
+                                                           'type': 'RandomRotate',
+                                                           'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}},
+                                  'max_num_transforms': 3,
+                                  'random_order': False,
+                                  'wrist_tfs': {'brightness': {'kwargs': {'brightness': [0.8,
+                                                                                         1.2]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'contrast': {'kwargs': {'contrast': [0.8,
+                                                                                     1.2]},
+                                                             'type': 'ColorJitter',
+                                                             'weight': 1.0},
+                                                'hue': {'kwargs': {'hue': [-0.05,
+                                                                           0.05]},
+                                                        'type': 'ColorJitter',
+                                                        'weight': 1.0},
+                                                'saturation': {'kwargs': {'saturation': [0.5,
+                                                                                         1.5]},
+                                                               'type': 'ColorJitter',
+                                                               'weight': 1.0},
+                                                'sharpness': {'kwargs': {'sharpness': [0.5,
+                                                                                       1.5]},
+                                                              'type': 'SharpnessJitter',
+                                                              'weight': 1.0}}},
+             'repo_id': '.',
+             'revision': None,
+             'root': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10',
+             'use_imagenet_stats': True,
+             'video_backend': 'pyav',
+             'vqa_data_path': None},
+ 'env': None,
+ 'eval': {'batch_size': 50, 'n_episodes': 50, 'use_async_envs': False},
+ 'eval_freq': 20000,
+ 'job_name': 'pi0fast_baseline_libero_30k',
+ 'log_freq': 100,
+ 'num_workers': 8,
+ 'optimizer': {'betas': [0.9, 0.95],
+               'eps': 1e-08,
+               'grad_clip_norm': 1,
+               'lr': 0.0001,
+               'type': 'adamw',
+               'weight_decay': 1e-05},
+ 'output_dir': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/outputs/train/2026-02-08/07-52-25_pi0fast_baseline_libero_30k',
+ 'policy': {'adapt_to_pi_aloha': False,
+            'checkpoint_path': None,
+            'chunk_size': 10,
+            'device': 'cuda',
+            'empty_cameras': 0,
+            'fast_skip_tokens': 128,
+            'fast_tokenizer_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/fast',
+            'freeze_lm_head': True,
+            'freeze_vision_encoder': True,
+            'grad_clip_norm': 1,
+            'gradient_accumulation_steps': 4,
+            'input_features': {},
+            'interpolate_like_pi': False,
+            'max_action_dim': 32,
+            'max_decoding_steps': 256,
+            'max_input_seq_len': 256,
+            'max_state_dim': 32,
+            'n_action_steps': 10,
+            'n_obs_steps': 1,
+            'normalization_mapping': {'ACTION': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'STATE': <NormalizationMode.MEAN_STD: 'MEAN_STD'>,
+                                      'VISUAL': <NormalizationMode.IDENTITY: 'IDENTITY'>},
+            'optimizer_betas': [0.9, 0.95],
+            'optimizer_eps': 1e-08,
+            'optimizer_lr': 0.0001,
+            'optimizer_weight_decay': 1e-05,
+            'output_features': {},
+            'padding_side': 'right',
+            'pi0_paligemma_path': '/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/paligemma-3b-pt-224',
+            'precision': 'bfloat16',
+            'proj_width': 1024,
+            'relaxed_action_decoding': True,
+            'resize_imgs_with_padding': [224, 224],
+            'scheduler_decay_lr': 2.5e-06,
+            'scheduler_decay_steps': 30000,
+            'scheduler_warmup_steps': 1000,
+            'tokenizer_max_length': 48,
+            'type': 'pi0fast',
+            'use_amp': False,
+            'use_cache': True,
+            'use_delta_joint_actions_aloha': False},
+ 'resume': False,
+ 'save_checkpoint': True,
+ 'save_freq': 5000,
+ 'scheduler': {'decay_lr': 2.5e-06,
+               'num_decay_steps': 30000,
+               'num_warmup_steps': 1000,
+               'peak_lr': 0.0001,
+               'type': 'cosine_decay_with_warmup'},
+ 'seed': 42,
+ 'steps': 100000,
+ 'use_policy_training_preset': True,
+ 'wandb': {'disable_artifact': True,
+           'enable': False,
+           'entity': 'Robotics_VLA',
+           'mode': 'online',
+           'notes': None,
+           'project': 'pi0_lerobot',
+           'run_id': None}}
+load fresh policy
+load fresh policy
+load fresh policy
+load fresh policy
+===========================================
+Training completed!
+Outputs in: /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/outputs/train/2026-02-08/07-52-25_pi0fast_baseline_libero_30k
+===========================================
+============ Job Information ===================================================
+Submitted: 2026-02-08T07:52:07
+Started: 2026-02-08T07:52:16
+Ended: 2026-02-09T07:22:29
+Elapsed: 1411 min, Limit: 2880 min, Difference: 1469 min
+CPUs: 32, Nodes: 1
+Estimated Consumption: 14110.00 core-hours
+================================================================================

scripts/serve_policy_ricl.py

@@ -0,0 +1,90 @@
+import dataclasses
+import enum
+import logging
+import socket
+
+import tyro
+
+from openpi.policies import policy as _policy
+from openpi.policies import policy_config as _policy_config
+from openpi.serving import websocket_policy_server
+from openpi.training import config as _config
+
+
+class EnvMode(enum.Enum):
+    """Supported environments."""
+
+    ALOHA = "aloha"
+    ALOHA_SIM = "aloha_sim"
+    DROID = "droid"
+    LIBERO = "libero"
+
+
+@dataclasses.dataclass
+class Checkpoint:
+    """Load a policy from a trained checkpoint."""
+
+    # Training config name (e.g., "pi0_aloha_sim").
+    config: str
+    # Checkpoint directory (e.g., "checkpoints/pi0_aloha_sim/exp/10000").
+    dir: str
+    # Demos directory (e.g., "ricl_droid_preprocessing/collected_demos/2025-03-04").
+    demos_dir: str
+
+
+@dataclasses.dataclass
+class Default:
+    """Use the default policy for the given environment."""
+
+
+@dataclasses.dataclass
+class Args:
+    """Arguments for the serve_policy script."""
+
+    # Environment to serve the policy for. This is only used when serving default policies.
+    env: EnvMode = EnvMode.DROID
+
+    # If provided, will be used in case the "prompt" key is not present in the data, or if the model doesn't have a default
+    # prompt.
+    default_prompt: str | None = None
+
+    # Port to serve the policy on.
+    port: int = 8000
+    # Record the policy's behavior for debugging.
+    record: bool = False
+
+    # Specifies how to load the policy. If not provided, the default policy for the environment will be used.
+    policy: Checkpoint | Default = dataclasses.field(default_factory=Default)
+
+
+def create_policy(args: Args) -> _policy.Policy:
+    """Create a policy from the given arguments."""
+    return _policy_config.create_trained_ricl_policy(
+        _config.get_config(args.policy.config), args.policy.dir, demos_dir=args.policy.demos_dir
+    )
+
+
+def main(args: Args) -> None:
+    policy = create_policy(args)
+    policy_metadata = policy.metadata
+
+    # Record the policy's behavior.
+    if args.record:
+        policy = _policy.PolicyRecorder(policy, "policy_records")
+
+    hostname = socket.gethostname()
+    local_ip = socket.gethostbyname(hostname)
+    logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip)
+
+    server = websocket_policy_server.WebsocketPolicyServer(
+        policy=policy,
+        host="0.0.0.0",
+        port=args.port,
+        metadata=policy_metadata,
+    )
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, force=True)
+    main(tyro.cli(Args))

scripts/test_ricl_data_pipeline.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+"""
+Quick test script to validate LIBERO dataset loading and RICL precomputation logic.
+
+This runs a minimal version of the precomputation to verify:
+1. Parquet files can be loaded
+2. Video frames can be extracted
+3. DINOv2 embeddings can be built
+4. FAISS index works correctly
+5. Self-exclusion in retrieval works
+
+Usage:
+    python scripts/test_ricl_data_pipeline.py
+"""
+
+import sys
+from pathlib import Path
+
+# Add src to path
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
+
+import numpy as np
+import pandas as pd
+from glob import glob
+from tqdm import tqdm
+import torch
+import torchvision
+
+print("=" * 60)
+print("RICL Data Pipeline Test")
+print("=" * 60)
+
+# Paths - Using smaller test dataset
+DATA_DIR = Path("/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10")
+VIDEOS_DIR = DATA_DIR / "videos"
+
+print(f"Testing with dataset: {DATA_DIR}")
+print(f"Dataset exists: {DATA_DIR.exists()}")
+
+# Test 1: Load parquet files
+print("\nTest 1: Loading parquet files...")
+parquet_files = sorted(glob(str(DATA_DIR / "data" / "**" / "*.parquet"), recursive=True))
+print(f"  Found {len(parquet_files)} parquet files")
+
+# Load first few files
+dfs = []
+for f in parquet_files[:2]:  # Just first 2 chunks
+    df = pd.read_parquet(f)
+    dfs.append(df)
+    print(f"  Loaded {f.split('/')[-1]}: {len(df)} rows")
+
+df = pd.concat(dfs, ignore_index=True)
+print(f"  Total frames: {len(df)}")
+print(f"  ✓ PASSED")
+
+# Test 2: Extract state and action
+print("\nTest 2: Extracting state and action...")
+row = df.iloc[0]
+state_ee = np.array(row["observation.states.ee_state"], dtype=np.float32)
+action = np.array(row["action"], dtype=np.float32)
+
+print(f"  State (ee_state) shape: {state_ee.shape}")
+print(f"  Action shape: {action.shape}")
+print(f"  State sample: {state_ee[:3]}")
+print(f"  Action sample: {action[:3]}")
+
+assert state_ee.shape == (6,), f"Expected state shape (6,), got {state_ee.shape}"
+assert action.shape == (7,), f"Expected action shape (7,), got {action.shape}"
+print(f"  ✓ PASSED")
+
+# Test 3: Load video frame
+print("\nTest 3: Loading video frame...")
+ep_idx = int(row["episode_index"])
+frame_idx = int(row["frame_index"])
+
+video_file = f"episode_{ep_idx:06d}.mp4"
+image_key = "observation.images.image"
+
+# Find video file
+video_path = None
+for chunk_dir in sorted(VIDEOS_DIR.glob("chunk-*")):
+    candidate = chunk_dir / image_key / video_file
+    if candidate.exists():
+        video_path = candidate
+        break
+
+if video_path is None:
+    print(f"  ✗ FAILED: Video not found for episode {ep_idx}")
+else:
+    print(f"  Video path: {video_path}")
+    
+    # Load frame
+    vframes, _, _ = torchvision.io.read_video(str(video_path), pts_unit='sec')
+    if frame_idx >= len(vframes):
+        frame_idx = len(vframes) - 1
+    
+    frame = vframes[frame_idx].numpy()
+    print(f"  Frame shape: {frame.shape}")
+    print(f"  Frame dtype: {frame.dtype}")
+    print(f"  Frame range: [{frame.min()}, {frame.max()}]")
+    
+    assert frame.shape[2] == 3, f"Expected RGB image, got shape {frame.shape}"
+    print(f"  ✓ PASSED")
+
+# Test 4: DINOv2 encoding (if model available)
+print("\nTest 4: DINOv2 encoding...")
+DINO_MODEL_PATH = Path("/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/VLA-Humanoid/rag/dinov2-base")
+
+if not DINO_MODEL_PATH.exists():
+    print(f"  SKIPPED: DINOv2 model not found at {DINO_MODEL_PATH}")
+else:
+    from transformers import AutoImageProcessor, AutoModel
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    processor = AutoImageProcessor.from_pretrained(str(DINO_MODEL_PATH), local_files_only=True)
+    model = AutoModel.from_pretrained(str(DINO_MODEL_PATH), local_files_only=True).to(device)
+    model.eval()
+    
+    with torch.no_grad():
+        inputs = processor(images=frame, return_tensors="pt")
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        outputs = model(**inputs)
+        embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy().squeeze()
+    
+    print(f"  Embedding shape: {embedding.shape}")
+    print(f"  Embedding norm: {np.linalg.norm(embedding):.4f}")
+    
+    assert embedding.shape == (768,), f"Expected embedding shape (768,), got {embedding.shape}"
+    print(f"  ✓ PASSED")
+
+# Test 5: Action chunk building
+print("\nTest 5: Action chunk building...")
+ACTION_HORIZON = 50
+
+# Get all actions for this episode
+episode_mask = df["episode_index"] == ep_idx
+ep_df = df[episode_mask]
+ep_actions = np.array([
+    np.array(row["action"], dtype=np.float32)
+    for _, row in ep_df.iterrows()
+])
+
+print(f"  Episode {ep_idx} length: {len(ep_df)} frames")
+print(f"  Episode actions shape: {ep_actions.shape}")
+
+# Build chunk
+local_frame_idx = 0
+action_chunk = []
+for i in range(ACTION_HORIZON):
+    if local_frame_idx + i < len(ep_actions):
+        action_chunk.append(ep_actions[local_frame_idx + i])
+    else:
+        action_chunk.append(ep_actions[-1])  # Pad with last action
+
+action_chunk = np.stack(action_chunk, axis=0)
+print(f"  Action chunk shape: {action_chunk.shape}")
+
+assert action_chunk.shape == (ACTION_HORIZON, 7), f"Expected shape ({ACTION_HORIZON}, 7), got {action_chunk.shape}"
+print(f"  ✓ PASSED")
+
+# Test 6: FAISS index (if faiss available)
+print("\nTest 6: FAISS index and retrieval...")
+try:
+    import faiss
+    
+    # Create dummy embeddings
+    N = 100
+    D = 768
+    embeddings = np.random.randn(N, D).astype(np.float32)
+    
+    # L2 normalize
+    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+    embeddings = embeddings / (norms + 1e-8)
+    
+    # Build index
+    index = faiss.IndexFlatL2(D)
+    index.add(embeddings)
+    
+    # Search for top-3 nearest neighbors
+    top_k = 3
+    search_k = top_k + 1  # +1 to exclude self
+    
+    query_idx = 0
+    query_emb = embeddings[query_idx:query_idx+1]
+    
+    D_result, I_result = index.search(query_emb, search_k)
+    
+    print(f"  Index built with {index.ntotal} vectors")
+    print(f"  Query index: {query_idx}")
+    print(f"  Retrieved indices: {I_result[0]}")
+    print(f"  Retrieved distances: {D_result[0]}")
+    
+    # Verify self is first result (distance ~0)
+    assert I_result[0, 0] == query_idx, "Self should be first result"
+    assert D_result[0, 0] < 1e-5, f"Self-distance should be ~0, got {D_result[0, 0]}"
+    
+    # Filter out self
+    mask = I_result[0] != query_idx
+    nn_indices = I_result[0][mask][:top_k]
+    nn_distances = D_result[0][mask][:top_k]
+    
+    print(f"  After self-exclusion:")
+    print(f"    NN indices: {nn_indices}")
+    print(f"    NN distances: {nn_distances}")
+    
+    assert query_idx not in nn_indices, "Self should be excluded"
+    assert len(nn_indices) == top_k, f"Should have {top_k} NNs, got {len(nn_indices)}"
+    
+    print(f"  ✓ PASSED")
+
+except ImportError:
+    print(f"  SKIPPED: faiss not available")
+
+print("\n" + "=" * 60)
+print("✓ All tests passed!")
+print("=" * 60)
+print("\nData pipeline is ready. You can now run:")
+print("  sbatch slurm/precompute_ricl_context_libero.slurm")
+print("=" * 60)

scripts/test_ricl_data_pipeline_v2.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+"""
+Quick test script for RICL LIBERO data pipeline.
+
+Tests:
+1. Dataset loading
+2. Video frame loading
+3. Action chunk building
+4. Transform pipeline
+5. Batch creation
+"""
+
+import sys
+from pathlib import Path
+
+# Add src to path
+sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
+
+import numpy as np
+
+def test_video_loading():
+    """Test that videos can be loaded."""
+    print("\n" + "="*60)
+    print("TEST 1: Video Loading")
+    print("="*60)
+    
+    import av
+    from pathlib import Path
+    
+    data_dir = Path("/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10")
+    
+    # Try to load first frame from first episode
+    video_path = data_dir / "videos" / "chunk-000" / "observation.images.image" / "episode_000000.mp4"
+    
+    if not video_path.exists():
+        print(f"✗ Video not found: {video_path}")
+        return False
+    
+    try:
+        with av.open(str(video_path)) as container:
+            frame = next(container.decode(video=0))
+            img = frame.to_ndarray(format='rgb24')
+            print(f"✓ Loaded frame shape: {img.shape}")
+            print(f"✓ Frame dtype: {img.dtype}")
+            print(f"✓ Frame range: [{img.min()}, {img.max()}]")
+            return True
+    except Exception as e:
+        print(f"✗ Failed to load video: {e}")
+        return False
+
+
+def test_dataset_loading():
+    """Test RiclLiberoDataset loading."""
+    print("\n" + "="*60)
+    print("TEST 2: Dataset Loading")
+    print("="*60)
+    
+    from openpi.data.ricl_libero_dataset import RiclLiberoDataset
+    
+    data_dir = "/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10"
+    context_dir = "rag/ricl_training_context_libero_10_test"
+    
+    try:
+        dataset = RiclLiberoDataset(
+            data_dir=data_dir,
+            context_dir=context_dir,
+            action_horizon=15,  # Use 15 like DROID
+            lambda_decay=10.0,
+            num_retrieved_observations=1,
+        )
+        
+        print(f"✓ Dataset loaded: {len(dataset)} samples")
+        return dataset
+    except Exception as e:
+        print(f"✗ Failed to load dataset: {e}")
+        import traceback
+        traceback.print_exc()
+        return None
+
+
+def test_sample_loading(dataset):
+    """Test loading a single sample."""
+    print("\n" + "="*60)
+    print("TEST 3: Sample Loading")
+    print("="*60)
+    
+    if dataset is None:
+        print("✗ Skipping (dataset not loaded)")
+        return False
+    
+    try:
+        sample = dataset[0]
+        
+        print("✓ Sample loaded successfully")
+        print(f"  Keys: {list(sample.keys())}")
+        
+        # Check images
+        for key in ['query_top_image', 'query_wrist_image', 'retrieved_0_top_image', 'retrieved_0_wrist_image']:
+            if key in sample:
+                img = sample[key]
+                print(f"  {key}: shape={img.shape}, dtype={img.dtype}")
+                assert img.shape[2] == 3, f"Expected RGB image, got shape {img.shape}"
+                assert img.dtype == np.uint8, f"Expected uint8, got {img.dtype}"
+        
+        # Check states
+        for key in ['query_state', 'retrieved_0_state']:
+            if key in sample:
+                state = sample[key]
+                print(f"  {key}: shape={state.shape}")
+        
+        # Check actions
+        for key in ['query_actions', 'retrieved_0_actions']:
+            if key in sample:
+                actions = sample[key]
+                print(f"  {key}: shape={actions.shape}")
+                assert actions.shape[1] == 7, f"Expected 7-dim actions, got {actions.shape}"
+        
+        return True
+        
+    except Exception as e:
+        print(f"✗ Failed to load sample: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+
+def test_batch_loading(dataset):
+    """Test loading a batch."""
+    print("\n" + "="*60)
+    print("TEST 4: Batch Loading")
+    print("="*60)
+    
+    if dataset is None:
+        print("✗ Skipping (dataset not loaded)")
+        return False
+    
+    try:
+        batch_size = 4
+        batch = [dataset[i] for i in range(min(batch_size, len(dataset)))]
+        
+        print(f"✓ Loaded batch of {len(batch)} samples")
+        
+        # Check consistency
+        keys = set(batch[0].keys())
+        for i, sample in enumerate(batch[1:], 1):
+            assert set(sample.keys()) == keys, f"Sample {i} has different keys"
+        
+        print(f"  All samples have consistent keys: {len(keys)} keys")
+        return True
+        
+    except Exception as e:
+        print(f"✗ Failed to load batch: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+
+def main():
+    print("\n" + "#"*60)
+    print("# RICL LIBERO Data Pipeline Test")
+    print("#"*60)
+    
+    results = {}
+    
+    # Run tests
+    results['video_loading'] = test_video_loading()
+    results['dataset_loading'] = test_dataset_loading()
+    
+    # Get dataset for subsequent tests
+    if results['dataset_loading']:
+        from openpi.data.ricl_libero_dataset import RiclLiberoDataset
+        dataset = RiclLiberoDataset(
+            data_dir="/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10",
+            context_dir="rag/ricl_training_context_libero_10_test",
+            action_horizon=15,
+            lambda_decay=10.0,
+            num_retrieved_observations=1,
+        )
+    else:
+        dataset = None
+    
+    results['sample_loading'] = test_sample_loading(dataset)
+    results['batch_loading'] = test_batch_loading(dataset)
+    
+    # Print summary
+    print("\n" + "="*60)
+    print("TEST SUMMARY")
+    print("="*60)
+    
+    for test_name, passed in results.items():
+        status = "✓ PASS" if passed else "✗ FAIL"
+        print(f"  {test_name:20s}: {status}")
+    
+    all_passed = all(results.values())
+    print()
+    if all_passed:
+        print("🎉 ALL TESTS PASSED!")
+    else:
+        print("❌ SOME TESTS FAILED")
+    
+    print("="*60)
+    return 0 if all_passed else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/train.py

@@ -0,0 +1,273 @@
+import dataclasses
+import functools
+import logging
+import platform
+from typing import Any
+
+import etils.epath as epath
+import flax.nnx as nnx
+from flax.training import common_utils
+import flax.traverse_util as traverse_util
+import jax
+import jax.experimental
+import jax.numpy as jnp
+import optax
+import tqdm_loggable.auto as tqdm
+import wandb
+
+import openpi.models.model as _model
+import openpi.shared.array_typing as at
+import openpi.shared.nnx_utils as nnx_utils
+import openpi.training.checkpoints as _checkpoints
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.optimizer as _optimizer
+import openpi.training.sharding as sharding
+import openpi.training.utils as training_utils
+import openpi.training.weight_loaders as _weight_loaders
+
+
+def init_logging():
+    """Custom logging format for better readability."""
+    level_mapping = {"DEBUG": "D", "INFO": "I", "WARNING": "W", "ERROR": "E", "CRITICAL": "C"}
+
+    class CustomFormatter(logging.Formatter):
+        def format(self, record):
+            record.levelname = level_mapping.get(record.levelname, record.levelname)
+            return super().format(record)
+
+    formatter = CustomFormatter(
+        fmt="%(asctime)s.%(msecs)03d [%(levelname)s] %(message)-80s (%(process)d:%(filename)s:%(lineno)s)",
+        datefmt="%H:%M:%S",
+    )
+
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    logger.handlers[0].setFormatter(formatter)
+
+
+def init_wandb(config: _config.TrainConfig, *, resuming: bool, log_code: bool = False, enabled: bool = True):
+    if not enabled:
+        wandb.init(mode="disabled")
+        return
+
+    ckpt_dir = config.checkpoint_dir
+    if not ckpt_dir.exists():
+        raise FileNotFoundError(f"Checkpoint directory {ckpt_dir} does not exist.")
+    if resuming:
+        run_id = (ckpt_dir / "wandb_id.txt").read_text().strip()
+        wandb.init(id=run_id, resume="must", project=config.project_name)
+    else:
+        wandb.init(
+            name=config.exp_name,
+            config=dataclasses.asdict(config),
+            project=config.project_name,
+        )
+        (ckpt_dir / "wandb_id.txt").write_text(wandb.run.id)
+
+    if log_code:
+        wandb.run.log_code(epath.Path(__file__).parent.parent)
+
+
+def _load_weights_and_validate(loader: _weight_loaders.WeightLoader, params_shape: at.Params) -> at.Params:
+    """Loads and validates the weights. Returns a loaded subset of the weights."""
+    loaded_params = loader.load(params_shape)
+    at.check_pytree_equality(expected=params_shape, got=loaded_params, check_shapes=True, check_dtypes=True)
+
+    # Remove jax.ShapeDtypeStruct from the loaded params. This makes sure that only the loaded params are returned.
+    return traverse_util.unflatten_dict(
+        {k: v for k, v in traverse_util.flatten_dict(loaded_params).items() if not isinstance(v, jax.ShapeDtypeStruct)}
+    )
+
+
+@at.typecheck
+def init_train_state(
+    config: _config.TrainConfig, init_rng: at.KeyArrayLike, mesh: jax.sharding.Mesh, *, resume: bool
+) -> tuple[training_utils.TrainState, Any]:
+    tx = _optimizer.create_optimizer(config.optimizer, config.lr_schedule, weight_decay_mask=None)
+
+    def init(rng: at.KeyArrayLike, partial_params: at.Params | None = None) -> training_utils.TrainState:
+        rng, model_rng = jax.random.split(rng)
+        # initialize the model (and its parameters).
+        model = config.model.create(model_rng)
+
+        # Merge the partial params into the model.
+        if partial_params is not None:
+            graphdef, state = nnx.split(model)
+            # This will produce an error if the partial params are not a subset of the state.
+            state.replace_by_pure_dict(partial_params)
+            model = nnx.merge(graphdef, state)
+
+        params = nnx.state(model)
+        # Convert frozen params to bfloat16.
+        params = nnx_utils.state_map(params, config.freeze_filter, lambda p: p.replace(p.value.astype(jnp.bfloat16)))
+
+        return training_utils.TrainState(
+            step=0,
+            params=params,
+            model_def=nnx.graphdef(model),
+            tx=tx,
+            opt_state=tx.init(params.filter(config.trainable_filter)),
+            ema_decay=config.ema_decay,
+            ema_params=None if config.ema_decay is None else params,
+        )
+
+    train_state_shape = jax.eval_shape(init, init_rng)
+    state_sharding = sharding.fsdp_sharding(train_state_shape, mesh, log=True)
+
+    if resume:
+        return train_state_shape, state_sharding
+
+    partial_params = _load_weights_and_validate(config.weight_loader, train_state_shape.params.to_pure_dict())
+    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+    # Initialize the train state and mix in the partial params.
+    train_state = jax.jit(
+        init,
+        donate_argnums=(1,),  # donate the partial params buffer.
+        in_shardings=replicated_sharding,
+        out_shardings=state_sharding,
+    )(init_rng, partial_params)
+
+    return train_state, state_sharding
+
+
+@at.typecheck
+def train_step(
+    config: _config.TrainConfig,
+    rng: at.KeyArrayLike,
+    state: training_utils.TrainState,
+    batch: tuple[_model.Observation, _model.Actions],
+) -> tuple[training_utils.TrainState, dict[str, at.Array]]:
+    model = nnx.merge(state.model_def, state.params)
+    model.train()
+
+    @at.typecheck
+    def loss_fn(
+        model: _model.BaseModel, rng: at.KeyArrayLike, observation: _model.Observation, actions: _model.Actions
+    ):
+        chunked_loss = model.compute_loss(rng, observation, actions, train=True)
+        return jnp.mean(chunked_loss)
+
+    train_rng = jax.random.fold_in(rng, state.step)
+    observation, actions = batch
+
+    # Filter out frozen params.
+    diff_state = nnx.DiffState(0, config.trainable_filter)
+    loss, grads = nnx.value_and_grad(loss_fn, argnums=diff_state)(model, train_rng, observation, actions)
+
+    params = state.params.filter(config.trainable_filter)
+    updates, new_opt_state = state.tx.update(grads, state.opt_state, params)
+    new_params = optax.apply_updates(params, updates)
+
+    # Update the model in place and return the new full state.
+    nnx.update(model, new_params)
+    new_params = nnx.state(model)
+
+    new_state = dataclasses.replace(state, step=state.step + 1, params=new_params, opt_state=new_opt_state)
+    if state.ema_decay is not None:
+        new_state = dataclasses.replace(
+            new_state,
+            ema_params=jax.tree.map(
+                lambda old, new: state.ema_decay * old + (1 - state.ema_decay) * new, state.ema_params, new_params
+            ),
+        )
+
+    # Filter out params that aren't kernels.
+    kernel_params = nnx.state(
+        model,
+        nnx.All(
+            nnx.Param,
+            nnx.Not(nnx_utils.PathRegex(".*/(bias|scale|pos_embedding|input_embedding)")),
+            lambda _, x: x.value.ndim > 1,
+        ),
+    )
+    info = {
+        "loss": loss,
+        "grad_norm": optax.global_norm(grads),
+        "param_norm": optax.global_norm(kernel_params),
+    }
+    return new_state, info
+
+
+def main(config: _config.TrainConfig):
+    init_logging()
+    logging.info(f"Running on: {platform.node()}")
+
+    if config.batch_size % jax.device_count() != 0:
+        raise ValueError(
+            f"Batch size {config.batch_size} must be divisible by the number of devices {jax.device_count()}."
+        )
+
+    jax.config.update("jax_compilation_cache_dir", str(epath.Path("~/.cache/jax").expanduser()))
+
+    rng = jax.random.key(config.seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    mesh = sharding.make_mesh(config.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+    checkpoint_manager, resuming = _checkpoints.initialize_checkpoint_dir(
+        config.checkpoint_dir,
+        keep_period=config.keep_period,
+        overwrite=config.overwrite,
+        resume=config.resume,
+    )
+    init_wandb(config, resuming=resuming, enabled=config.wandb_enabled)
+
+    data_loader = _data_loader.create_data_loader(
+        config,
+        sharding=data_sharding,
+        num_workers=config.num_workers,
+        shuffle=True,
+    )
+    data_iter = iter(data_loader)
+    batch = next(data_iter)
+    logging.info(f"Initialized data loader:\n{training_utils.array_tree_to_info(batch)}")
+
+    train_state, train_state_sharding = init_train_state(config, init_rng, mesh, resume=resuming)
+    jax.block_until_ready(train_state)
+    logging.info(f"Initialized train state:\n{training_utils.array_tree_to_info(train_state.params)}")
+
+    if resuming:
+        train_state = _checkpoints.restore_state(checkpoint_manager, train_state, data_loader)
+
+    ptrain_step = jax.jit(
+        functools.partial(train_step, config),
+        in_shardings=(replicated_sharding, train_state_sharding, data_sharding),
+        out_shardings=(train_state_sharding, replicated_sharding),
+        donate_argnums=(1,),
+    )
+
+    start_step = int(train_state.step)
+    pbar = tqdm.tqdm(
+        range(start_step, config.num_train_steps),
+        initial=start_step,
+        total=config.num_train_steps,
+        dynamic_ncols=True,
+    )
+
+    infos = []
+    for step in pbar:
+        with sharding.set_mesh(mesh):
+            train_state, info = ptrain_step(train_rng, train_state, batch)
+        infos.append(info)
+        if step % config.log_interval == 0:
+            stacked_infos = common_utils.stack_forest(infos)
+            reduced_info = jax.device_get(jax.tree.map(jnp.mean, stacked_infos))
+            info_str = ", ".join(f"{k}={v:.4f}" for k, v in reduced_info.items())
+            pbar.write(f"Step {step}: {info_str}")
+            wandb.log(reduced_info, step=step)
+            infos = []
+        batch = next(data_iter)
+
+        if (step % config.save_interval == 0 and step > start_step) or step == config.num_train_steps - 1:
+            _checkpoints.save_state(checkpoint_manager, train_state, data_loader, step)
+
+    logging.info("Waiting for checkpoint manager to finish")
+    checkpoint_manager.wait_until_finished()
+
+
+if __name__ == "__main__":
+    main(_config.cli())

scripts/train_debug_interactive.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+
+# Interactive debug training (no SLURM - runs directly on allocated node)
+
+cd /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi
+
+export PYTHONPATH="/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/src:$PYTHONPATH"
+export HF_HOME="/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/hf_cache"
+
+# Memory optimization for JAX/XLA
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+
+echo "=========================================="
+echo "RICL Interactive Debug Training"
+echo "=========================================="
+echo "Config:"
+echo "  Steps: 500 (quick test)"
+echo "  Batch: 2 (reduced to fit single GPU memory)"
+echo "  Action Horizon: 15"
+echo "  Lambda: 10.0"
+echo "=========================================="
+
+python scripts/train_pi0_fast_ricl.py \
+    pi0_fast_ricl_libero \
+    --exp-name="ricl_libero_debug_interactive" \
+    --project-name="ricl_openpi" \
+    --libero-data-dir="/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10" \
+    --libero-context-dir="rag/ricl_training_context_libero_10_test" \
+    --model.lamda=10.0 \
+    --model.use-action-interpolation \
+    --model.num-retrieved-observations=1 \
+    --model.action-horizon=15 \
+    --model.action-dim=7 \
+    --model.max-token-len=250 \
+    --batch-size=2 \
+    --num-train-steps=500 \
+    --save-interval=500 \
+    --log-interval=10 \
+    --lr-schedule.peak-lr=1e-4 \
+    --seed=42 \
+    --no-wandb-enabled \
+    --overwrite \
+    --no-resume
+
+echo ""
+echo "=========================================="
+echo "Debug Training Complete!"
+echo "=========================================="

scripts/train_pi0_fast_baseline_libero.slurm

@@ -0,0 +1,85 @@
+#!/bin/bash
+#SBATCH --job-name=pi0fast_baseline_libero
+#SBATCH --output=logs/pi0fast_baseline_libero_h100.out
+#SBATCH --error=logs/pi0fast_baseline_libero_h100.err
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:4
+#SBATCH --cpus-per-task=32
+#SBATCH --time=48:00:00
+#SBATCH --mem=400G
+#SBATCH --partition=kisski
+
+# ==========================================
+# Baseline Pi0-FAST Training for LIBERO
+# Using VLA-Humanoid Codebase (LeRobot)
+# ==========================================
+
+# Activate Environment
+source /projects/extern/kisski/kisski-spath/dir.project/VLA_3D/miniconda3/etc/profile.d/conda.sh
+conda activate pitorch_doanh
+
+# Environment Variables
+export PROJECT_DIR="/projects/extern/kisski/kisski-spath/dir.project"
+export PYTHONPATH="$PROJECT_DIR/VLA_Groot/in_context_learning/VLA-Humanoid:$PYTHONPATH"
+export HF_HOME="$PROJECT_DIR/VLA_Groot/hf_cache" 
+
+# Paths
+VLA_HUMANOID_DIR="$PROJECT_DIR/VLA_Groot/in_context_learning/VLA-Humanoid"
+DATA_DIR="$PROJECT_DIR/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10"
+POLICY_CONFIG_SOURCE="$VLA_HUMANOID_DIR/configs/policy_config/pi0fast_baseline.json"
+LIBERO_CONFIG_PATH="$VLA_HUMANOID_DIR/configs/libero_config/default.json"
+
+EXP_NAME="pi0fast_baseline_libero_pretrained_fix"
+# Output directory
+OUTPUT_DIR="$VLA_HUMANOID_DIR/outputs/train/$(date +%Y-%m-%d)/$(date +%H-%M-%S)_$EXP_NAME"
+# Temp directory for initializing config (must separate from output to avoid FileExistsError)
+TEMP_CONFIG_DIR="$VLA_HUMANOID_DIR/temp/configs/$EXP_NAME"
+
+echo "==========================================="
+echo "Experiment: $EXP_NAME"
+echo "Codebase: $VLA_HUMANOID_DIR"
+echo "Data: $DATA_DIR"
+echo "Policy Config Source: $POLICY_CONFIG_SOURCE"
+echo "Temp Config Dir: $TEMP_CONFIG_DIR"
+echo "Output Dir: $OUTPUT_DIR"
+echo "==========================================="
+
+# Move to codebase directory
+cd $VLA_HUMANOID_DIR
+
+# 1. Run Verification Script (CRITICAL: Fail if setup is wrong)
+echo "==========================================="
+echo "Running verification script..."
+echo "==========================================="
+python scripts/verify_training_setup.py
+if [ $? -ne 0 ]; then
+    echo "❌ ERROR: Verification failed! Aborting training."
+    exit 1
+fi
+echo "✓ Verification passed. Proceeding to training..."
+
+# 2. Setup Temp Config Directory (Source for initialization)
+mkdir -p $TEMP_CONFIG_DIR
+cp $POLICY_CONFIG_SOURCE $TEMP_CONFIG_DIR/config.json
+
+# 3. Run Training
+# Note: OUTPUT_DIR is created by the script, it should not exist beforehand if resume=False
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+accelerate launch --num_processes=4 --main_process_port 29600 lerobot/scripts/train_accelerate.py \
+  --policy.type=pi0fast \
+  --dataset.root=$DATA_DIR \
+  --dataset.video_backend=pyav \
+  --output_dir=$OUTPUT_DIR \
+  --job_name=$EXP_NAME \
+  --config_path=$LIBERO_CONFIG_PATH \
+  --batch_size=4 \
+  --policy.gradient_accumulation_steps=4 \
+  --save_freq=5000 \
+  --log_freq=100 \
+  --wandb.enable=false
+
+echo "==========================================="
+echo "Training completed!"
+echo "Outputs in: $OUTPUT_DIR"
+echo "==========================================="

scripts/train_pi0_fast_ricl.py

@@ -0,0 +1,317 @@
+import dataclasses
+import functools
+import logging
+import platform
+from typing import Any
+
+import etils.epath as epath
+import flax.nnx as nnx
+from flax.training import common_utils
+import flax.traverse_util as traverse_util
+import jax
+import jax.experimental
+import jax.numpy as jnp
+import optax
+import tqdm_loggable.auto as tqdm
+import wandb
+from torch.utils.tensorboard import SummaryWriter
+
+import openpi.models.model as _model
+import openpi.shared.array_typing as at
+import openpi.shared.nnx_utils as nnx_utils
+import openpi.training.checkpoints as _checkpoints
+import openpi.training.config as _config
+import openpi.training.data_loader as _data_loader
+import openpi.training.optimizer as _optimizer
+import openpi.training.sharding as sharding
+import openpi.training.utils as training_utils
+import openpi.training.weight_loaders as _weight_loaders
+
+
+def init_logging():
+    """Custom logging format for better readability."""
+    level_mapping = {"DEBUG": "D", "INFO": "I", "WARNING": "W", "ERROR": "E", "CRITICAL": "C"}
+
+    class CustomFormatter(logging.Formatter):
+        def format(self, record):
+            record.levelname = level_mapping.get(record.levelname, record.levelname)
+            return super().format(record)
+
+    formatter = CustomFormatter(
+        fmt="%(asctime)s.%(msecs)03d [%(levelname)s] %(message)-80s (%(process)d:%(filename)s:%(lineno)s)",
+        datefmt="%H:%M:%S",
+    )
+
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    logger.handlers[0].setFormatter(formatter)
+
+
+def init_wandb(config: _config.TrainConfig, *, resuming: bool, log_code: bool = False, enabled: bool = True):
+    if not enabled:
+        wandb.init(mode="disabled")
+        return
+
+    ckpt_dir = config.checkpoint_dir
+    if not ckpt_dir.exists():
+        raise FileNotFoundError(f"Checkpoint directory {ckpt_dir} does not exist.")
+    if resuming:
+        run_id = (ckpt_dir / "wandb_id.txt").read_text().strip()
+        wandb.init(id=run_id, resume="must", project=config.project_name)
+    else:
+        wandb.init(
+            name=config.exp_name,
+            config=dataclasses.asdict(config),
+            project=config.project_name,
+        )
+        (ckpt_dir / "wandb_id.txt").write_text(wandb.run.id)
+
+    if log_code:
+        wandb.run.log_code(epath.Path(__file__).parent.parent)
+
+
+def _load_weights_and_validate(loader: _weight_loaders.WeightLoader, params_shape: at.Params) -> at.Params:
+    """Loads and validates the weights. Returns a loaded subset of the weights."""
+    loaded_params = loader.load(params_shape)
+    at.check_pytree_equality(expected=params_shape, got=loaded_params, check_shapes=True, check_dtypes=True)
+
+    # Remove jax.ShapeDtypeStruct from the loaded params. This makes sure that only the loaded params are returned.
+    return traverse_util.unflatten_dict(
+        {k: v for k, v in traverse_util.flatten_dict(loaded_params).items() if not isinstance(v, jax.ShapeDtypeStruct)}
+    )
+
+
+def count_parameters(params, trainable_filter):
+    """Counts total and trainable parameters."""
+    total_params = sum(p.size for p in jax.tree_util.tree_leaves(params))
+    
+    trainable_params = sum(
+        p.size for p in jax.tree_util.tree_leaves(params.filter(trainable_filter))
+    )
+    
+    return total_params, trainable_params
+
+
+@at.typecheck
+def init_train_state(
+    config: _config.TrainConfig, init_rng: at.KeyArrayLike, mesh: jax.sharding.Mesh, *, resume: bool
+) -> tuple[training_utils.TrainState, Any]:
+    tx = _optimizer.create_optimizer(config.optimizer, config.lr_schedule, weight_decay_mask=None)
+
+    def init(rng: at.KeyArrayLike, partial_params: at.Params | None = None) -> training_utils.TrainState:
+        rng, model_rng = jax.random.split(rng)
+        # initialize the model (and its parameters).
+        model = config.model.create(model_rng)
+
+        # Merge the partial params into the model.
+        if partial_params is not None:
+            graphdef, state = nnx.split(model)
+            # This will produce an error if the partial params are not a subset of the state.
+            state.replace_by_pure_dict(partial_params)
+            model = nnx.merge(graphdef, state)
+
+        params = nnx.state(model)
+        # Convert frozen params to bfloat16.
+        params = nnx_utils.state_map(params, config.freeze_filter, lambda p: p.replace(p.value.astype(jnp.bfloat16)))
+
+        # Count total and trainable parameters
+        total_params, trainable_params = count_parameters(params, config.trainable_filter)
+        print(f"Total Parameters: {total_params // 1e6}M")
+        print(f"Trainable Parameters: {trainable_params // 1e6}M")
+        print(f"Trainable Parameters %: {100 * trainable_params / total_params:.2f}%")
+
+        return training_utils.TrainState(
+            step=0,
+            params=params,
+            model_def=nnx.graphdef(model),
+            tx=tx,
+            opt_state=tx.init(params.filter(config.trainable_filter)),
+            ema_decay=config.ema_decay,
+            ema_params=None if config.ema_decay is None else params,
+        )
+
+    train_state_shape = jax.eval_shape(init, init_rng)
+    state_sharding = sharding.fsdp_sharding(train_state_shape, mesh, log=True)
+
+    if resume:
+        return train_state_shape, state_sharding
+
+    partial_params = _load_weights_and_validate(config.weight_loader, train_state_shape.params.to_pure_dict())
+    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+    # Initialize the train state and mix in the partial params.
+    train_state = jax.jit(
+        init,
+        donate_argnums=(1,),  # donate the partial params buffer.
+        in_shardings=replicated_sharding,
+        out_shardings=state_sharding,
+    )(init_rng, partial_params)
+
+    return train_state, state_sharding
+
+
+def create_decode_indices(config: _config.TrainConfig) -> at.Int[at.Array, "decode_len"]:
+    # create the indices to decide which tokens to decode: i.e. only those belonging to each retrieved/query "prompt, state, action" prompt and not the images
+    image_token_len = 256*2 # number of image tokens times number of images
+    prompt_token_len = config.model.max_token_len # max token len for each retrieved/query "prompt, state, action" prompt
+    total_token_len = image_token_len + prompt_token_len
+    decode_indices = []
+    for i in range(config.model.num_retrieved_observations + 1):
+        decode_indices.extend(list(range(i * total_token_len + image_token_len + 1, (i+1) * total_token_len)))
+    decode_indices = jnp.asarray(decode_indices)
+    print(f'decode_indices shape: {decode_indices.shape}')
+    return decode_indices
+
+
+# @at.typecheck
+def train_step(
+    config: _config.TrainConfig,
+    rng: at.KeyArrayLike,
+    state: training_utils.TrainState,
+    batch: tuple[_model.RiclObservation, _model.Actions],
+    decode_indices: at.Int[at.Array, "decode_len"],
+) -> tuple[training_utils.TrainState, dict[str, at.Array]]:
+    model = nnx.merge(state.model_def, state.params)
+    model.train()
+
+    # @at.typecheck
+    def loss_fn(
+        model: _model.BaseModel, rng: at.KeyArrayLike, observation: _model.RiclObservation, actions: _model.Actions
+    ):
+        chunked_loss = model.compute_loss(rng, observation, actions, train=True, decode_indices=decode_indices)
+        return jnp.mean(chunked_loss)
+
+    train_rng = jax.random.fold_in(rng, state.step)
+    observation, actions = batch
+
+    # Filter out frozen params.
+    diff_state = nnx.DiffState(0, config.trainable_filter)
+    loss, grads = nnx.value_and_grad(loss_fn, argnums=diff_state)(model, train_rng, observation, actions)
+
+    params = state.params.filter(config.trainable_filter)
+    updates, new_opt_state = state.tx.update(grads, state.opt_state, params)
+    new_params = optax.apply_updates(params, updates)
+
+    # Update the model in place and return the new full state.
+    nnx.update(model, new_params)
+    new_params = nnx.state(model)
+
+    new_state = dataclasses.replace(state, step=state.step + 1, params=new_params, opt_state=new_opt_state)
+    if state.ema_decay is not None:
+        new_state = dataclasses.replace(
+            new_state,
+            ema_params=jax.tree.map(
+                lambda old, new: state.ema_decay * old + (1 - state.ema_decay) * new, state.ema_params, new_params
+            ),
+        )
+
+    # Filter out params that aren't kernels.
+    kernel_params = nnx.state(
+        model,
+        nnx.All(
+            nnx.Param,
+            nnx.Not(nnx_utils.PathRegex(".*/(bias|scale|pos_embedding|input_embedding)")),
+            lambda _, x: x.value.ndim > 1,
+        ),
+    )
+    info = {
+        "loss": loss,
+        "grad_norm": optax.global_norm(grads),
+        "param_norm": optax.global_norm(kernel_params),
+    }
+    return new_state, info
+
+
+def main(config: _config.TrainConfig):
+    init_logging()
+    logging.info(f"Running on: {platform.node()}")
+
+    if config.batch_size % jax.device_count() != 0:
+        raise ValueError(
+            f"Batch size {config.batch_size} must be divisible by the number of devices {jax.device_count()}."
+        )
+
+    jax.config.update("jax_compilation_cache_dir", str(epath.Path("~/.cache/jax").expanduser()))
+
+    rng = jax.random.key(config.seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    mesh = sharding.make_mesh(config.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+    checkpoint_manager, resuming = _checkpoints.initialize_checkpoint_dir(
+        config.checkpoint_dir,
+        keep_period=config.keep_period,
+        overwrite=config.overwrite,
+        resume=config.resume,
+    )
+    init_wandb(config, resuming=resuming, enabled=config.wandb_enabled)
+    
+    writer = None
+    if not config.wandb_enabled:
+        log_dir = config.checkpoint_dir / "tensorboard"
+        logging.info(f"WandB disabled. Using TensorBoard logging to {log_dir}")
+        writer = SummaryWriter(log_dir=str(log_dir))
+
+    data_loader = _data_loader.create_data_loader(
+        config,
+        sharding=data_sharding,
+        num_workers=config.num_workers,
+        shuffle=True,
+    )
+    data_iter = iter(data_loader)
+    batch = next(data_iter)
+    logging.info(f"Initialized data loader:\n{training_utils.array_tree_to_info(batch)}")
+
+    train_state, train_state_sharding = init_train_state(config, init_rng, mesh, resume=resuming)
+    jax.block_until_ready(train_state)
+    logging.info(f"Initialized train state:\n{training_utils.array_tree_to_info(train_state.params)}")
+
+    if resuming:
+        train_state = _checkpoints.restore_state(checkpoint_manager, train_state, data_loader)
+
+    decode_indices = create_decode_indices(config)
+
+    ptrain_step = jax.jit(
+        functools.partial(train_step, config, decode_indices=decode_indices),
+        in_shardings=(replicated_sharding, train_state_sharding, data_sharding),
+        out_shardings=(train_state_sharding, replicated_sharding),
+        donate_argnums=(1,),
+    )
+
+    start_step = int(train_state.step)
+    pbar = tqdm.tqdm(
+        range(start_step, config.num_train_steps),
+        initial=start_step,
+        total=config.num_train_steps,
+        dynamic_ncols=True,
+    )
+
+    infos = []
+    for step in pbar:
+        with sharding.set_mesh(mesh):
+            train_state, info = ptrain_step(train_rng, train_state, batch)
+        infos.append(info)
+        if step % config.log_interval == 0:
+            stacked_infos = common_utils.stack_forest(infos)
+            reduced_info = jax.device_get(jax.tree.map(jnp.mean, stacked_infos))
+            info_str = ", ".join(f"{k}={v}" for k, v in reduced_info.items())
+            pbar.write(f"Step {step}: {info_str}")
+            if config.wandb_enabled:
+                wandb.log(reduced_info, step=step)
+            elif writer is not None:
+                for k, v in reduced_info.items():
+                    writer.add_scalar(k, v, step)
+            infos = []
+        batch = next(data_iter)
+
+        if (step % config.save_interval == 0 and step > start_step) or step == config.num_train_steps - 1:
+            _checkpoints.save_state(checkpoint_manager, train_state, data_loader, step)
+
+    logging.info("Waiting for checkpoint manager to finish")
+    checkpoint_manager.wait_until_finished()
+
+
+if __name__ == "__main__":
+    main(_config.cli())

scripts/train_pi0_fast_ricl_libero_80k.slurm

@@ -0,0 +1,69 @@
+#!/bin/bash
+#SBATCH --job-name=ricl_libero_80k
+#SBATCH --output=logs/ricl_libero_80k_%j.out
+#SBATCH --error=logs/ricl_libero_80k_%j.err
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:4
+#SBATCH --cpus-per-task=32
+#SBATCH --time=48:00:00
+#SBATCH --mem=400G
+#SBATCH --partition=kisski
+
+# Activate Conda environment
+source /projects/extern/kisski/kisski-spath/dir.project/VLA_3D/miniconda3/etc/profile.d/conda.sh
+conda activate pitorch_doanh
+
+# Set PYTHONPATH to include the current directory
+export PYTHONPATH=/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/src:$PYTHONPATH
+
+# Workdir
+cd /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi
+
+# Create logs directory
+mkdir -p logs/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/scripts/train_pi0_fast_ricl_libero_80k.slurm
+
+# Configuration
+EXP_NAME="ricl_libero_lambda10_80k_bs16"
+PROJECT_NAME="ricl_libero_bs16"
+LIBERO_DATA_DIR="/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10"
+LIBERO_CONTEXT_DIR="rag/ricl_training_context_libero_10_test/"
+
+# Hyperparameters (Recommended Config)
+LAMBDA=10.0
+ACTION_HORIZON=15  # Reduced from 50 to match DROID (prevents token overflow)
+NUM_RETRIEVED=1
+ACTION_DIM=7
+MAX_TOKEN_LEN=250
+
+# Training parameters
+BATCH_SIZE=16   # Global batch size (4 per GPU for 4 GPUs)
+NUM_STEPS=80000
+SAVE_INTERVAL=10000
+LOG_INTERVAL=100
+
+echo "Starting training with Experiment Name: ${EXP_NAME}"
+echo "Data Dir: ${LIBERO_DATA_DIR}"
+echo "Context Dir: ${LIBERO_CONTEXT_DIR}"
+
+python scripts/train_pi0_fast_ricl.py pi0_fast_ricl_libero \
+    --exp-name="${EXP_NAME}" \
+    --project-name="${PROJECT_NAME}" \
+    --libero-data-dir="${LIBERO_DATA_DIR}" \
+    --libero-context-dir="${LIBERO_CONTEXT_DIR}" \
+    --model.lamda=${LAMBDA} \
+    --model.use-action-interpolation \
+    --model.num-retrieved-observations=${NUM_RETRIEVED} \
+    --model.action-horizon=${ACTION_HORIZON} \
+    --model.action-dim=${ACTION_DIM} \
+    --model.max-token-len=${MAX_TOKEN_LEN} \
+    --batch-size=${BATCH_SIZE} \
+    --num-train-steps=${NUM_STEPS} \
+    --save-interval=${SAVE_INTERVAL} \
+    --log-interval=${LOG_INTERVAL} \
+    --lr-schedule.peak-lr=1e-4 \
+    --seed=42 \
+    --no-wandb-enabled \
+    --overwrite \
+    --no-resume \
+    --fsdp_devices=4

scripts/train_pi0_fast_ricl_libero_80k_2nn.slurm

@@ -0,0 +1,70 @@
+#!/bin/bash
+#SBATCH --job-name=ricl_libero_2nn_80k
+#SBATCH --output=logs/ricl_libero_2nn_80k_%j.out
+#SBATCH --error=logs/ricl_libero_2nn_80k_%j.err
+#SBATCH --mem=128G
+#SBATCH --cpus-per-task=32
+#SBATCH --partition=kisski
+#SBATCH --time=48:00:00
+#SBATCH --gres=gpu:4
+
+# Environment setup
+source /projects/extern/kisski/kisski-spath/dir.project/VLA_3D/miniconda3/etc/profile.d/conda.sh
+conda activate pitorch_doanh
+
+export PYTHONPATH=/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/src:$PYTHONPATH
+
+# Create logs directory
+mkdir -p logs/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi/scripts/train_pi0_fast_ricl_libero_80k_2nn.slurm
+
+# Configuration
+EXP_NAME="ricl_libero_lambda10_80k_2nn_fsdp"
+PROJECT_NAME="ricl_libero_2nn"
+LIBERO_DATA_DIR="/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10"
+LIBERO_CONTEXT_DIR="rag/ricl_training_context_libero_10_test/"
+
+# Model hyperparameters
+LAMBDA=10.0
+NUM_RETRIEVED=2  # Increased to 2 neighbors
+ACTION_HORIZON=50
+ACTION_DIM=7
+MAX_TOKEN_LEN=250
+
+# Training parameters
+BATCH_SIZE=16   # Global batch size (4 per GPU for 4 GPUs)
+NUM_STEPS=80000
+SAVE_INTERVAL=10000
+LOG_INTERVAL=100
+
+echo "Starting training job: $SLURM_JOB_ID"
+echo "Experiment: $EXP_NAME"
+echo "Batch Size: $BATCH_SIZE"
+echo "Num Retrieved: $NUM_RETRIEVED"
+
+cd /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi
+
+# Run training
+# Note: FSDP is enabled to handle memory usage
+python scripts/train_pi0_fast_ricl.py pi0_fast_ricl_libero \
+    --exp-name="$EXP_NAME" \
+    --project-name="$PROJECT_NAME" \
+    --libero-data-dir="$LIBERO_DATA_DIR" \
+    --libero-context-dir="$LIBERO_CONTEXT_DIR" \
+    --model.lamda=$LAMBDA \
+    --model.use-action-interpolation \
+    --model.num-retrieved-observations=$NUM_RETRIEVED \
+    --model.action-horizon=$ACTION_HORIZON \
+    --model.action-dim=$ACTION_DIM \
+    --model.max-token-len=$MAX_TOKEN_LEN \
+    --batch-size=$BATCH_SIZE \
+    --num-train-steps=$NUM_STEPS \
+    --save-interval=$SAVE_INTERVAL \
+    --log-interval=$LOG_INTERVAL \
+    --lr-schedule.peak-lr=1e-4 \
+    --seed=42 \
+    --no-wandb-enabled \
+    --overwrite \
+    --no-resume \
+    --fsdp_devices=4 
+
+echo "Training completed"

scripts/train_test.py

@@ -0,0 +1,30 @@
+import dataclasses
+import os
+import pathlib
+
+import pytest
+
+os.environ["JAX_PLATFORMS"] = "cpu"
+
+from openpi.training import config as _config
+
+from . import train
+
+
+@pytest.mark.parametrize("config_name", ["debug"])
+def test_train(tmp_path: pathlib.Path, config_name: str):
+    config = dataclasses.replace(
+        _config._CONFIGS_DICT[config_name],  # noqa: SLF001
+        batch_size=2,
+        checkpoint_base_dir=tmp_path / "checkpoint",
+        exp_name="test",
+        overwrite=False,
+        resume=False,
+        num_train_steps=2,
+        log_interval=1,
+    )
+    train.main(config)
+
+    # test resuming
+    config = dataclasses.replace(config, resume=True, num_train_steps=4)
+    train.main(config)

scripts/verify_ricl_training_setup.py

@@ -0,0 +1,514 @@
+#!/usr/bin/env python3
+"""
+RICL Training Verification Script
+Run this BEFORE training to ensure everything is configured correctly.
+
+Usage:
+    cd /projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi
+    python scripts/verify_ricl_training_setup.py
+"""
+
+import os
+import sys
+import json
+import numpy as np
+from pathlib import Path
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+WORK_DIR = "/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/in_context_learning/ricl_openpi"
+DATA_DIR = "/projects/extern/kisski/kisski-spath/dir.project/VLA_Groot/merged_libero_mask_depth_noops_lerobot_10"
+CONTEXT_DIR = os.path.join(WORK_DIR, "rag/ricl_training_context_libero_10_test")
+CHECKPOINT_PATH = os.path.join(WORK_DIR, "pi0_fast_base_params")
+ASSETS_DIR = os.path.join(WORK_DIR, "assets")
+
+sys.path.insert(0, WORK_DIR)
+
+passed = 0
+failed = 0
+
+def check_pass(msg):
+    global passed
+    passed += 1
+    print(f"  ✓ {msg}")
+
+def check_fail(msg):
+    global failed
+    failed += 1
+    print(f"  ❌ {msg}")
+
+# ============================================================
+# CHECK 1: Verify all paths exist
+# ============================================================
+print("=" * 60)
+print("CHECK 1: Verify all paths exist")
+print("=" * 60)
+
+paths_to_check = {
+    "WORK_DIR": WORK_DIR,
+    "DATA_DIR": DATA_DIR,
+    "CONTEXT_DIR": CONTEXT_DIR,
+    "CHECKPOINT_PATH (pi0_fast_base_params)": CHECKPOINT_PATH,
+    "ASSETS_DIR": ASSETS_DIR,
+}
+
+for name, path in paths_to_check.items():
+    if os.path.exists(path):
+        check_pass(f"{name}: {path}")
+    else:
+        check_fail(f"{name} MISSING: {path}")
+print()
+
+# ============================================================
+# CHECK 2: Verify JAX/Orbax checkpoint structure
+# ============================================================
+print("=" * 60)
+print("CHECK 2: Verify base model checkpoint (JAX/Orbax)")
+print("=" * 60)
+
+required_ckpt_files = ["_METADATA", "_sharding", "manifest.ocdbt"]
+for f in required_ckpt_files:
+    fpath = os.path.join(CHECKPOINT_PATH, f)
+    if os.path.exists(fpath):
+        size = os.path.getsize(fpath)
+        check_pass(f"{f}: {size} bytes")
+    else:
+        check_fail(f"{f} MISSING")
+
+ocdbt_dir = os.path.join(CHECKPOINT_PATH, "ocdbt.process_0")
+if os.path.exists(ocdbt_dir):
+    num_files = len(os.listdir(ocdbt_dir))
+    check_pass(f"ocdbt.process_0/: {num_files} shard files")
+else:
+    check_fail("ocdbt.process_0/ directory MISSING")
+print()
+
+# ============================================================
+# CHECK 3: Verify RICL context directory
+# ============================================================
+print("=" * 60)
+print("CHECK 3: Verify RICL retrieval context")
+print("=" * 60)
+
+required_context_files = {
+    "nn_indices.npy": "Nearest neighbor indices",
+    "nn_distances.npy": "Nearest neighbor distances",
+    "actions.npy": "Action chunks",
+    "states.npy": "State vectors",
+    "metadata.json": "Frame metadata",
+    "embeddings.npy": "Visual embeddings",
+    "index.faiss": "FAISS index",
+}
+
+for fname, desc in required_context_files.items():
+    fpath = os.path.join(CONTEXT_DIR, fname)
+    if os.path.exists(fpath):
+        size_mb = os.path.getsize(fpath) / 1e6
+        check_pass(f"{fname} ({desc}): {size_mb:.1f} MB")
+    else:
+        check_fail(f"{fname} ({desc}) MISSING")
+print()
+
+# ============================================================
+# CHECK 4: Load and validate context data shapes
+# ============================================================
+print("=" * 60)
+print("CHECK 4: Validate context data shapes and contents")
+print("=" * 60)
+
+try:
+    nn_indices = np.load(os.path.join(CONTEXT_DIR, "nn_indices.npy"))
+    nn_distances = np.load(os.path.join(CONTEXT_DIR, "nn_distances.npy"))
+    actions = np.load(os.path.join(CONTEXT_DIR, "actions.npy"))
+    states = np.load(os.path.join(CONTEXT_DIR, "states.npy"))
+    
+    with open(os.path.join(CONTEXT_DIR, "metadata.json"), "r") as f:
+        metadata = json.load(f)
+    
+    num_frames = len(metadata)
+    print(f"  Total frames: {num_frames}")
+    print(f"  nn_indices shape: {nn_indices.shape}")
+    print(f"  nn_distances shape: {nn_distances.shape}")
+    print(f"  actions shape: {actions.shape}")
+    print(f"  states shape: {states.shape}")
+    
+    # Validate shapes match
+    if nn_indices.shape[0] == num_frames:
+        check_pass(f"nn_indices rows ({nn_indices.shape[0]}) match metadata ({num_frames})")
+    else:
+        check_fail(f"nn_indices rows ({nn_indices.shape[0]}) != metadata ({num_frames})")
+    
+    if nn_distances.shape[0] == num_frames:
+        check_pass(f"nn_distances rows ({nn_distances.shape[0]}) match metadata ({num_frames})")
+    else:
+        check_fail(f"nn_distances rows ({nn_distances.shape[0]}) != metadata ({num_frames})")
+    
+    if actions.shape[0] == num_frames:
+        check_pass(f"actions rows ({actions.shape[0]}) match metadata ({num_frames})")
+    else:
+        check_fail(f"actions rows ({actions.shape[0]}) != metadata ({num_frames})")
+    
+    # Check action dimensions
+    action_dim = actions.shape[-1] if len(actions.shape) > 1 else 0
+    print(f"  Action dimension: {action_dim}")
+    if action_dim == 7:
+        check_pass(f"Action dimension is 7 (LIBERO standard)")
+    else:
+        check_fail(f"Action dimension is {action_dim}, expected 7")
+    
+    # Action horizon
+    if len(actions.shape) == 3:
+        action_horizon = actions.shape[1]
+        print(f"  Action horizon: {action_horizon}")
+    
+    # Check nn_indices are valid (no out-of-range)
+    max_idx = nn_indices.max()
+    if max_idx < num_frames:
+        check_pass(f"nn_indices max ({max_idx}) < num_frames ({num_frames})")
+    else:
+        check_fail(f"nn_indices max ({max_idx}) >= num_frames ({num_frames}) - OUT OF RANGE!")
+    
+    # Check no self-retrieval (query != demo)
+    if nn_indices.shape[1] >= 1:
+        self_retrieval_count = np.sum(nn_indices[:, 0] == np.arange(num_frames))
+        self_pct = 100 * self_retrieval_count / num_frames
+        if self_pct < 5:
+            check_pass(f"Self-retrieval rate: {self_pct:.1f}% (low, good)")
+        else:
+            check_fail(f"Self-retrieval rate: {self_pct:.1f}% (HIGH - might be a bug!)")
+    
+    # Distance statistics
+    top1_dist = nn_distances[:, 0]
+    print(f"\n  Top-1 distance statistics:")
+    print(f"    Min: {top1_dist.min():.4f}")
+    print(f"    Max: {top1_dist.max():.4f}")
+    print(f"    Mean: {top1_dist.mean():.4f}")
+    print(f"    Median: {np.median(top1_dist):.4f}")
+    print(f"    Std: {top1_dist.std():.4f}")
+    
+    if top1_dist.mean() > 0:
+        check_pass(f"Distances are non-zero (mean={top1_dist.mean():.4f})")
+    else:
+        check_fail(f"Distances are all zero - retrieval might be broken!")
+
+except Exception as e:
+    check_fail(f"Error loading context: {e}")
+    import traceback
+    traceback.print_exc()
+
+print()
+
+# ============================================================
+# CHECK 5: Verify dataset structure
+# ============================================================
+print("=" * 60)
+print("CHECK 5: Verify dataset structure")
+print("=" * 60)
+
+meta_path = os.path.join(DATA_DIR, "meta")
+if os.path.exists(meta_path):
+    info_path = os.path.join(meta_path, "info.json")
+    if os.path.exists(info_path):
+        with open(info_path, "r") as f:
+            info = json.load(f)
+        check_pass(f"Dataset info found")
+        print(f"    Total episodes: {info.get('total_episodes', 'N/A')}")
+        print(f"    Total frames: {info.get('total_frames', 'N/A')}")
+        
+        features = info.get("features", {})
+        image_keys = [k for k in features if "image" in k.lower() and "mask" not in k.lower() and "depth" not in k.lower()]
+        print(f"    RGB image keys: {image_keys}")
+        
+        # Check video files exist
+        video_dir = os.path.join(DATA_DIR, "videos", "chunk-000")
+        if os.path.exists(video_dir):
+            check_pass(f"Video directory exists: {video_dir}")
+            for img_key in ["observation.images.image", "observation.images.wrist_image"]:
+                key_dir = os.path.join(video_dir, img_key)
+                if os.path.exists(key_dir):
+                    num_vids = len([f for f in os.listdir(key_dir) if f.endswith(".mp4")])
+                    check_pass(f"{img_key}: {num_vids} video files")
+                else:
+                    check_fail(f"{img_key} video directory MISSING")
+        else:
+            check_fail(f"Video directory MISSING: {video_dir}")
+    else:
+        check_fail(f"info.json MISSING")
+else:
+    check_fail(f"meta directory MISSING")
+
+# Check tasks.jsonl
+tasks_file = os.path.join(DATA_DIR, "meta", "tasks.jsonl")
+if os.path.exists(tasks_file):
+    task_count = 0
+    with open(tasks_file, "r") as f:
+        for line in f:
+            task_count += 1
+    check_pass(f"tasks.jsonl: {task_count} tasks")
+else:
+    check_fail(f"tasks.jsonl MISSING")
+
+print()
+
+# ============================================================
+# CHECK 6: Retrieval quality debug (sample pairs)
+# ============================================================
+print("=" * 60)
+print("CHECK 6: Retrieval quality debug (sample pairs)")
+print("=" * 60)
+
+try:
+    import random
+    random.seed(42)
+    
+    # Load task mapping
+    task_mapping = {}
+    with open(os.path.join(DATA_DIR, "meta", "tasks.jsonl"), "r") as f:
+        for line in f:
+            item = json.loads(line)
+            if "task_index" in item and "task" in item:
+                task_mapping[item["task_index"]] = item["task"]
+    
+    # Load parquet for task indices
+    import pandas as pd
+    parquet_files = sorted(list(Path(DATA_DIR, "data").rglob("*.parquet")))
+    dfs = [pd.read_parquet(f) for f in parquet_files[:5]]  # Load a subset for speed
+    df = pd.concat(dfs, ignore_index=True)
+    
+    # Sample and display pairs
+    num_samples = 5
+    sample_indices = random.sample(range(min(len(metadata), len(df))), min(num_samples, len(metadata), len(df)))
+    
+    same_task_count = 0
+    for i, query_idx in enumerate(sample_indices):
+        demo_idx = int(nn_indices[query_idx, 0])
+        distance = nn_distances[query_idx, 0]
+        
+        query_meta = metadata[query_idx]
+        demo_meta = metadata[demo_idx]
+        
+        # Get tasks
+        query_global = query_meta.get("global_frame_idx", query_idx)
+        demo_global = demo_meta.get("global_frame_idx", demo_idx)
+        
+        query_task_idx = -1
+        demo_task_idx = -1
+        if query_global < len(df):
+            query_row = df.iloc[query_global]
+            query_task_idx = int(query_row.get("task_index", -1))
+        if demo_global < len(df):
+            demo_row = df.iloc[demo_global]
+            demo_task_idx = int(demo_row.get("task_index", -1))
+        
+        query_task = task_mapping.get(query_task_idx, "Unknown")
+        demo_task = task_mapping.get(demo_task_idx, "Unknown")
+        same_task = query_task_idx == demo_task_idx
+        if same_task:
+            same_task_count += 1
+        
+        # Action similarity
+        action_mse = np.mean((actions[query_idx] - actions[demo_idx]) ** 2)
+        
+        print(f"\n  Pair {i+1}/{num_samples}:")
+        print(f"    Query: ep={query_meta.get('episode_idx', '?')}, frame={query_meta.get('frame_idx', '?')}")
+        print(f"      Task: '{query_task[:70]}'")
+        print(f"    Demo:  ep={demo_meta.get('episode_idx', '?')}, frame={demo_meta.get('frame_idx', '?')}")
+        print(f"      Task: '{demo_task[:70]}'")
+        print(f"    Distance: {distance:.4f} | Action MSE: {action_mse:.4f} | Same task: {same_task}")
+    
+    same_task_pct = 100 * same_task_count / num_samples
+    print(f"\n  Same-task retrieval rate (in sample): {same_task_pct:.0f}%")
+    if same_task_pct >= 50:
+        check_pass(f"Same-task retrieval is reasonable ({same_task_pct:.0f}%)")
+    else:
+        print(f"  ⚠ Low same-task retrieval - this could indicate cross-task retrieval (may be intended)")
+
+except Exception as e:
+    check_fail(f"Error in retrieval debug: {e}")
+    import traceback
+    traceback.print_exc()
+
+print()
+
+# ============================================================
+# CHECK 7: Verify RiclLiberoDataset can load
+# ============================================================
+print("=" * 60)
+print("CHECK 7: Verify RiclLiberoDataset loads correctly")
+print("=" * 60)
+
+try:
+    from openpi.data.ricl_libero_dataset import RiclLiberoDataset
+    
+    TARGET_ACTION_HORIZON = 10  # Must match pi0fast-LIBERO default
+    
+    dataset = RiclLiberoDataset(
+        data_dir=DATA_DIR,
+        context_dir=CONTEXT_DIR,
+        action_horizon=TARGET_ACTION_HORIZON,  # Truncate from precomputed 50 → 10
+        use_action_interpolation=True,
+        lambda_decay=10.0,
+        num_retrieved_observations=1,
+    )
+    
+    check_pass(f"Dataset created successfully with {len(dataset)} samples")
+    
+    # Try loading one sample
+    print("  Loading sample [0]...")
+    sample = dataset[0]
+    print(f"  Sample keys: {sorted(sample.keys())}")
+    
+    # Check expected keys (matching actual RiclLiberoDataset output format)
+    expected_keys = [
+        "query_observation.images.image", 
+        "query_observation.images.wrist_image", 
+        "query_observation.state", 
+        "query_actions", 
+        "query_prompt",
+    ]
+    for key in expected_keys:
+        if key in sample:
+            val = sample[key]
+            if isinstance(val, np.ndarray):
+                check_pass(f"{key}: shape={val.shape}, dtype={val.dtype}")
+            else:
+                check_pass(f"{key}: type={type(val).__name__}")
+        else:
+            check_fail(f"{key} MISSING from sample")
+    
+    # Check demo keys
+    demo_keys = [k for k in sample.keys() if "retrieved" in k or "demo" in k]
+    print(f"  Demo keys: {demo_keys}")
+    for key in demo_keys:
+        val = sample[key]
+        if isinstance(val, np.ndarray):
+            check_pass(f"{key}: shape={val.shape}, dtype={val.dtype}")
+        else:
+            check_pass(f"{key}: type={type(val).__name__}")
+    
+    # Check interpolation weights
+    interp_keys = [k for k in sample.keys() if "lamda" in k or "lambda" in k or "interp" in k]
+    if interp_keys:
+        for key in interp_keys:
+            val = sample[key]
+            if isinstance(val, np.ndarray):
+                print(f"  {key}: shape={val.shape}, range=[{val.min():.4f}, {val.max():.4f}]")
+            else:
+                print(f"  {key}: {val}")
+        check_pass("Interpolation weights present")
+    else:
+        print(f"  ⚠ No interpolation weight keys found (checked: lamda, lambda, interp)")
+    
+    # ============================================================
+    # CHECK 7b: Verify action truncation (50 → 10)
+    # ============================================================
+    print()
+    print("=" * 60)
+    print(f"CHECK 7b: Verify action truncation (precomputed=50 → target={TARGET_ACTION_HORIZON})")
+    print("=" * 60)
+    
+    # Precomputed actions shape
+    precomputed_horizon = actions.shape[1]
+    print(f"  Precomputed context action_horizon: {precomputed_horizon}")
+    print(f"  Target action_horizon: {TARGET_ACTION_HORIZON}")
+    
+    # Check query_actions shape
+    query_actions = sample["query_actions"]
+    if query_actions.shape[0] == TARGET_ACTION_HORIZON:
+        check_pass(f"query_actions truncated correctly: shape={query_actions.shape} (horizon={TARGET_ACTION_HORIZON})")
+    else:
+        check_fail(f"query_actions NOT truncated: shape={query_actions.shape} (expected horizon={TARGET_ACTION_HORIZON})")
+    
+    # Check demo actions shape
+    demo_actions = sample["retrieved_0_actions"]
+    if demo_actions.shape[0] == TARGET_ACTION_HORIZON:
+        check_pass(f"retrieved_0_actions truncated correctly: shape={demo_actions.shape} (horizon={TARGET_ACTION_HORIZON})")
+    else:
+        check_fail(f"retrieved_0_actions NOT truncated: shape={demo_actions.shape} (expected horizon={TARGET_ACTION_HORIZON})")
+    
+    # Verify truncation preserves data (first 10 of 50 should match)
+    raw_demo_actions_full = actions[int(nn_indices[0, 0])]  # Full 50-step from context
+    raw_demo_truncated = raw_demo_actions_full[:TARGET_ACTION_HORIZON]
+    if np.allclose(demo_actions, raw_demo_truncated, atol=1e-5):
+        check_pass(f"Truncated demo actions match first {TARGET_ACTION_HORIZON} steps of precomputed context")
+    else:
+        check_fail(f"Truncated demo actions DO NOT match precomputed context first {TARGET_ACTION_HORIZON} steps!")
+    
+    # Verify multiple samples to ensure consistency
+    num_verify = 5
+    all_correct = True
+    for vi in range(1, min(num_verify + 1, len(dataset))):
+        s = dataset[vi]
+        if s["query_actions"].shape[0] != TARGET_ACTION_HORIZON:
+            all_correct = False
+            break
+        if s["retrieved_0_actions"].shape[0] != TARGET_ACTION_HORIZON:
+            all_correct = False
+            break
+    if all_correct:
+        check_pass(f"Truncation verified across {num_verify} additional samples")
+    else:
+        check_fail(f"Truncation inconsistent across samples!")
+
+except Exception as e:
+    check_fail(f"Error loading dataset: {e}")
+    import traceback
+    traceback.print_exc()
+
+print()
+
+# ============================================================
+# CHECK 8: Verify config resolution
+# ============================================================
+print("=" * 60)
+print("CHECK 8: Verify training config resolution")
+print("=" * 60)
+
+try:
+    import etils.epath as epath
+    
+    # Check PI0_FAST_BASE_CHECKPOINT resolution
+    local_ckpt = epath.Path("pi0_fast_base_params")
+    if local_ckpt.exists():
+        check_pass(f"PI0_FAST_BASE_CHECKPOINT resolves to LOCAL: pi0_fast_base_params")
+    else:
+        check_fail(f"pi0_fast_base_params NOT found locally - will try S3 (requires network!)")
+    
+    # Check assets
+    assets_franka = os.path.join(WORK_DIR, "pi0_fast_base", "assets", "franka")
+    if os.path.exists(assets_franka):
+        check_pass(f"Assets dir (franka) exists")
+    else:
+        # Check alternative location
+        assets_alt = os.path.join(WORK_DIR, "assets")
+        if os.path.exists(assets_alt):
+            check_pass(f"Assets base dir exists: {assets_alt}")
+        else:
+            check_fail(f"No assets directory found")
+
+except Exception as e:
+    check_fail(f"Config check error: {e}")
+
+print()
+
+# ============================================================
+# SUMMARY
+# ============================================================
+print("=" * 60)
+print("VERIFICATION COMPLETE")
+print("=" * 60)
+total = passed + failed
+print(f"\n  ✓ Passed: {passed}/{total}")
+print(f"  ❌ Failed: {failed}/{total}")
+
+if failed == 0:
+    print("\n  🎉 ALL CHECKS PASSED! Ready to train.")
+else:
+    print(f"\n  ⚠ {failed} check(s) failed. Fix before training.")
+
+print(f"""
+To train RICL, submit:
+  cd {WORK_DIR}
+  sbatch slurm/train_ricl_libero.slurm
+""")