update stupid folder name from dexiang to the folder names used in paper

Author: ChonghaoSima

README.md

@@ -80,7 +80,7 @@ Non-edge components (e.g., Policy Training, Model Arithmetic) have been tested o
 
 ### Hardware
 
-For real-robot deployment (dual-arm setup, cameras, and table layout), see **[Hardware Setup & 3D Print Files](setup/README.md)**. That document covers supported platforms (Agilex Piper for FlattenFold / TeeShirtSort, ARX X5 for HangCloth), Intel RealSense D435i camera placement, 3D-printed grippers and mounts with usage notes, and inference host GPU (RTX 4090 in Ubuntu 20.04).
+For real-robot deployment (dual-arm setup, cameras, and table layout), see **[Hardware Setup & 3D Print Files](setup/README.md)**. That document covers supported platforms (Agilex Piper for Task_A / Task_B, ARX X5 for Task_C), Intel RealSense D435i camera placement, 3D-printed grippers and mounts with usage notes, and inference host GPU (RTX 4090 in Ubuntu 20.04).
 
 ## Installation
 
@@ -116,11 +116,11 @@ Download the Kai0 dataset so it is available under `./data` for training and eva
 python scripts/download_dataset.py
 ```
 
-This fetches the full dataset from [Hugging Face](https://huggingface.co/datasets/OpenDriveLab-org/Kai0) into `./data` (FlattenFold, HangCloth, TeeShirtSort). To download only specific tasks or use a custom path, see the [dataset docs](docs/dataset.md#step-1-download-the-dataset).
+This fetches the full dataset from [Hugging Face](https://huggingface.co/datasets/OpenDriveLab-org/Kai0) into `./data` (Task_A, Task_B, Task_C). To download only specific tasks or use a custom path, see the [dataset docs](docs/dataset.md#step-1-download-the-dataset).
 
 ### 2. Download checkpoints (optional, for testing)
 
-We provide **one best model per task** (FlattenFold, HangCloth, TeeShirtSort) in the [Kai0 repo on Hugging Face](https://huggingface.co/OpenDriveLab-org/Kai0/tree/main).
+We provide **one best model per task** (Task_A, Task_B, Task_C) in the [Kai0 repo on Hugging Face](https://huggingface.co/OpenDriveLab-org/Kai0/tree/main).
 
 From the repository root, you can download all best-model checkpoints to `./checkpoints` with:
 
@@ -131,7 +131,7 @@ python scripts/download_checkpoints.py
 To download only specific tasks or use a custom path, run:
 
 ```bash
-python scripts/download_checkpoints.py --tasks FlattenFold HangCloth --local-dir ./my_checkpoints
+python scripts/download_checkpoints.py --tasks Task_A Task_C --local-dir ./my_checkpoints
 ```
 
 After download, set `weight_loader` in the training config to the path of the corresponding checkpoint directory (see step 3 below). You can also use openpi’s pretrained π₀.5 checkpoint instead.
@@ -144,7 +144,7 @@ After the dataset is in `./data`, you can run **normal π₀.₅ full fine-tunin
 
 Edit [`src/openpi/training/config.py`](src/openpi/training/config.py) (around lines 1173–1226) for the task(s) you need:
 
-- **`repo_id`**: set to the **absolute path** to the dataset subset, e.g. `<path_to_repo_root>/data/FlattenFold/base`, `<path_to_repo_root>/data/TeeShirtSort/base`, or `<path_to_repo_root>/data/HangCloth/base`.
+- **`repo_id`**: set to the **absolute path** to the dataset subset, e.g. `<path_to_repo_root>/data/Task_A/base`, `<path_to_repo_root>/data/Task_B/base`, or `<path_to_repo_root>/data/Task_C/base`.
 - **`weight_loader`**: set to the path of your **π₀.₅ base checkpoint** — either the best model you downloaded in step 2 above, or openpi’s pretrained π₀.₅ checkpoint.
 
 Config names to use: e.g. `pi05_flatten_fold_normal`
@@ -300,7 +300,7 @@ For a ready-to-use script with environment setup (conda/venv activation, DDP con
 **Stage 2 — Advantage Estimation on New Data**: Use the trained estimator to label datasets with predicted advantage values.
 
 ```bash
-uv run python stage_advantage/annotation/eval.py Flatten-Fold KAI0 /path/to/dataset
+uv run python stage_advantage/annotation/eval.py Task-A KAI0 /path/to/dataset
 ```
 
 For a ready-to-use script with environment setup and status logging, see `stage_advantage/annotation/eval.sh`.

setup/README.md

@@ -6,15 +6,15 @@ Quick reference for deploying and debugging hardware for the supported task plat
 
 ## Table of Contents
 
-- [1. FlattenFold / TeeShirtSort (Agilex Piper)](#1-flattenfold--teeshirtsort-agilex-piper)
-- [2. HangCloth (ARX X5)](#2-hangcloth-arx-x5)
+- [1. Task_A / Task_B (Agilex Piper)](#1-task_a--task_b-agilex-piper)
+- [2. Task_C (ARX X5)](#2-task_c-arx-x5)
 - [3. Inference Host](#3-inference-host)
 
 ---
 
-## 1. FlattenFold / TeeShirtSort (Agilex Piper)
+## 1. Task_A / Task_B (Agilex Piper)
 
-**Directories:** `FlattenFold/`, `TeeShirtSort/`
+**Directories:** `Task_A/`, `Task_B/`
 
 ### 1.1 Components
 
@@ -24,7 +24,7 @@ Quick reference for deploying and debugging hardware for the supported task plat
 | Cameras | Intel RealSense D435 (triple-camera setup) |
 | Printed parts | Left/right wrist camera mounts, center camera mount, center camera base |
 
-### 1.2 FlattenFold Layout
+### 1.2 Task_A Layout
 
 | Parameter | Value |
 |-----------|-------|
@@ -37,7 +37,7 @@ Quick reference for deploying and debugging hardware for the supported task plat
 | Right primary arm → table front edge | 12 cm |
 | Left–right primary arm center distance | 39 cm |
 
-### 1.3 TeeShirtSort Layout (demoA-style)
+### 1.3 Task_B Layout (demoA-style)
 
 | Parameter | Value |
 |-----------|-------|
@@ -50,7 +50,7 @@ Quick reference for deploying and debugging hardware for the supported task plat
 | Right primary arm → table front edge | 11 cm |
 | Left–right primary arm center distance | 40 cm |
 
-### 1.4 3D Models — Usage (FlattenFold / TeeShirtSort)
+### 1.4 3D Models — Usage (Task_A / Task_B)
 
 #### Gripper (end-effector)
 
@@ -77,9 +77,9 @@ Quick reference for deploying and debugging hardware for the supported task plat
 
 ---
 
-## 2. HangCloth (ARX X5)
+## 2. Task_C (ARX X5)
 
-**Directory:** `HangCloth/`
+**Directory:** `Task_C/`
 
 ### 2.1 Components
 
@@ -102,7 +102,7 @@ Quick reference for deploying and debugging hardware for the supported task plat
 | Right primary arm → table front edge | 11 cm |
 | Left–right primary arm center distance | 53 cm |
 
-### 2.3 3D Models — Usage (HangCloth)
+### 2.3 3D Models — Usage (Task_C)
 
 #### Grippers (secondary arms)
 

stage_advantage/README.md

@@ -235,18 +235,18 @@ Examples:
 
 ```bash
 # KAI0 (two-timestep) on a dataset
-uv run python stage_advantage/annotation/eval.py Flatten-Fold KAI0 /path/to/dataset
+uv run python stage_advantage/annotation/eval.py Task-A KAI0 /path/to/dataset
 
 # PI06 (single-timestep)
-uv run python stage_advantage/annotation/eval.py Flatten-Fold PI06 /path/to/dataset
+uv run python stage_advantage/annotation/eval.py Task-A PI06 /path/to/dataset
 ```
 
-`<model_type>` is a key in `eval.py`'s `MODELS_CONFIG_MAP` (e.g. `Flatten-Fold`); `<model_name>` is `PI06` or `KAI0`; `<repo_id>` is the path to the LeRobot dataset. Results are written under `<repo_id>/data_<model_name>_<ckpt_steps>/`.
+`<model_type>` is a key in `eval.py`'s `MODELS_CONFIG_MAP` (e.g. `Task-A`); `<model_name>` is `PI06` or `KAI0`; `<repo_id>` is the path to the LeRobot dataset. Results are written under `<repo_id>/data_<model_name>_<ckpt_steps>/`.
 
 For a ready-to-use script with environment setup (conda/venv activation, environment variables) and status logging, see **`annotation/eval.sh`**:
 
 ```bash
-bash stage_advantage/annotation/eval.sh Flatten-Fold KAI0 /path/to/dataset
+bash stage_advantage/annotation/eval.sh Task-A KAI0 /path/to/dataset
 ```
 
 ### Evaluation Outputs
@@ -274,13 +274,13 @@ The output parquets can then be used in Stage 3 (AWBC) or fed back into Stage 0
 
 **Goal**: Train a policy using **Advantage-Weighted Behavior Cloning (AWBC)**. The advantage labels (from Stage 0 + Stage 2) are stored as `task_index` per frame and as prompt strings in `meta/tasks.jsonl`. By setting **`prompt_from_task=True`** in the data config, each sample’s prompt is taken from that mapping, so the policy is conditioned on the advantage-derived label (e.g. high vs low advantage) and effectively does advantage-weighted behavior cloning via the language channel.
 
-**Configs** (in `src/openpi/training/config.py`): `pi05_flatten_fold_awbc`, `pi05_tee_shirt_sort_awbc`, `pi05_hang_cloth_awbc`. Each uses `LerobotAgilexDataConfig` or `LerobotARXDataConfig` with `base_config=DataConfig(prompt_from_task=True)` and `repo_id` pointing to the **advantage** dataset (e.g. `.../data/FlattenFold/advantage`).
+**Configs** (in `src/openpi/training/config.py`): `pi05_flatten_fold_awbc`, `pi05_tee_shirt_sort_awbc`, `pi05_hang_cloth_awbc`. Each uses `LerobotAgilexDataConfig` or `LerobotARXDataConfig` with `base_config=DataConfig(prompt_from_task=True)` and `repo_id` pointing to the **advantage** dataset (e.g. `.../data/Task_A/advantage`).
 
 ### What the policy sees as prompt (training)
 
 The prompt is read from the dataset’s **`meta/tasks.jsonl`**: each frame’s `task_index` is mapped to a task string, and that string is passed to the policy as the language prompt. **`gt_label.py`** (Stage 0) writes these strings when it builds the advantage-labeled dataset.
 
-- **Binary mode** (typical): `task_index=0` → `"<task>, Advantage: negative"`, `task_index=1` → `"<task>, Advantage: positive"`. The `<task>` text is set in `gt_label.py` (e.g. `"fold the cloth"` for FlattenFold).
+- **Binary mode** (typical): `task_index=0` → `"<task>, Advantage: negative"`, `task_index=1` → `"<task>, Advantage: positive"`. The `<task>` text is set in `gt_label.py` (e.g. `"fold the cloth"` for Task_A).
 - **n_slices mode**: `task_index=i` → `"<task>, Advantage: {i}"`.
 
 So during AWBC training the model is conditioned on prompts that explicitly include the advantage label (e.g. `"fold the cloth, Advantage: positive"` or `"fold the cloth, Advantage: negative"`).
@@ -299,7 +299,7 @@ At **inference** time you must use the **same prompt format** as in training. To
 
 ### Before training
 
-1. **Produce the advantage dataset:** Run Stage 2 (eval) on your dataset so it has `data_PI06_100000/` or `data_KAI0_100000/`. Then run Stage 0 (e.g. `gt_labeling.sh`) with `DATA_PATH` = that repo and source subdirs `data_PI06_100000` / `data_KAI0_100000`; the script outputs a directory with `data/` (parquets with `task_index`), `meta/tasks.jsonl`, and `videos`. Use that directory as the advantage dataset (e.g. copy or link it to `./data/FlattenFold/advantage`).
+1. **Produce the advantage dataset:** Run Stage 2 (eval) on your dataset so it has `data_PI06_100000/` or `data_KAI0_100000/`. Then run Stage 0 (e.g. `gt_labeling.sh`) with `DATA_PATH` = that repo and source subdirs `data_PI06_100000` / `data_KAI0_100000`; the script outputs a directory with `data/` (parquets with `task_index`), `meta/tasks.jsonl`, and `videos`. Use that directory as the advantage dataset (e.g. copy or link it to `./data/Task_A/advantage`).
 2. In `config.py`, set **`repo_id`** to that advantage dataset path and **`weight_loader`** to your π₀.5 base checkpoint for the AWBC config(s) you use.
 3. **Compute norm stats:**  
    `uv run python scripts/compute_norm_states_fast.py --config-name pi05_flatten_fold_awbc`  

stage_advantage/annotation/README.md

@@ -16,7 +16,7 @@ uv run python scripts/train_pytorch.py ADVANTAGE_TORCH_KAI0_FLATTEN_FOLD --exp_n
 
 # Step 3: Evaluate the trained estimator on new data (PI06 or KAI0)
 # From repo root:
-uv run python stage_advantage/annotation/eval.py Flatten-Fold KAI0 /path/to/dataset
+uv run python stage_advantage/annotation/eval.py Task-A KAI0 /path/to/dataset
 
 # Step 4: Use the advantage-labeled data for AWBC (Stage 3)
 # After Stage 2, run gt_labeling.sh with DATA_PATH = eval repo (or gt_label.py --advantage-source absolute_advantage).

stage_advantage/awbc/README.md

@@ -8,9 +8,9 @@ All three are defined in `src/openpi/training/config.py`:
 
 | Config name | Task | Data config |
 |-------------|------|-------------|
-| `pi05_flatten_fold_awbc` | FlattenFold | `LerobotAgilexDataConfig`, `repo_id=.../data/FlattenFold/advantage` |
-| `pi05_tee_shirt_sort_awbc` | TeeShirtSort | `LerobotAgilexDataConfig`, `repo_id=.../data/TeeShirtSort/advantage` |
-| `pi05_hang_cloth_awbc` | HangCloth | `LerobotARXDataConfig`, `repo_id=.../data/HangCloth/advantage` |
+| `pi05_flatten_fold_awbc` | Task_A | `LerobotAgilexDataConfig`, `repo_id=.../data/Task_A/advantage` |
+| `pi05_tee_shirt_sort_awbc` | Task_B | `LerobotAgilexDataConfig`, `repo_id=.../data/Task_B/advantage` |
+| `pi05_hang_cloth_awbc` | Task_C | `LerobotARXDataConfig`, `repo_id=.../data/Task_C/advantage` |
 
 Each uses `base_config=DataConfig(prompt_from_task=True)` so that the dataset’s `task_index` column and `meta/tasks.jsonl` supply the prompt (advantage-derived label) per frame.
 
@@ -24,11 +24,11 @@ Each uses `base_config=DataConfig(prompt_from_task=True)` so that the dataset’
    To build your own advantage dataset instead:
    - Run **Stage 2** (eval) on your dataset → get `data_PI06_100000/` or `data_KAI0_100000/` with advantage columns.
    - Run **Stage 0** on that output: `gt_label.py --advantage-source absolute_advantage` (or `gt_labeling.sh` with `DATA_PATH` = the eval repo). The resulting directory (with `data/`, `meta/tasks.jsonl`, `videos/`) is your advantage dataset.
-   - Place or link it at e.g. `./data/FlattenFold/advantage` and set `repo_id` in config to that path.
+   - Place or link it at e.g. `./data/Task_A/advantage` and set `repo_id` in config to that path.
 
 2. **Config paths**  
    In `src/openpi/training/config.py`, for the AWBC config(s) you use:
-   - Set **`repo_id`** to the **absolute path** of the advantage dataset (e.g. `<path_to_repo_root>/data/FlattenFold/advantage`).
+   - Set **`repo_id`** to the **absolute path** of the advantage dataset (e.g. `<path_to_repo_root>/data/Task_A/advantage`).
    - Set **`weight_loader`** to your **π₀.5 base checkpoint** path.
 
 3. **Norm stats**