Minor fix (#77)

* minor fix

* lint

---------

Co-authored-by: liangyuxin <liang@ali.com>

README.md

@@ -61,7 +61,7 @@ pip install -e .
 Install wall_x:
 ```bash
 git submodule update --init --recursive
-MAX_JOBS=4 pip install --no-build-isolation --verbose .
+MAX_JOBS=4 pip install --no-build-isolation --verbose -e .
 ```
 
 ## Training

scripts/compute_norm_stats.py

@@ -18,14 +18,16 @@ def load_config(config_path):
     return config
 
 
-def load_lerobot_dataset(repo_id, action_horizon, args):
+def load_lerobot_dataset(repo_id, root, action_horizon, args):
     dataset_meta = LeRobotDatasetMetadata(repo_id)
     dataset = LeRobotDataset(
         repo_id,
+        root=root,
         delta_timestamps={
             key: [t / dataset_meta.fps for t in range(action_horizon)]
             for key in [KEY_MAPPINGS[repo_id]["action"]]
         },
+        video_backend="pyav",
     )
     num_batches = len(dataset) // args.batch_size
     generator = torch.Generator()
@@ -56,10 +58,11 @@ if __name__ == "__main__":
     config = load_config(path)
     lerobot_config = config["data"]["lerobot_config"]
     repo_id = lerobot_config.get("repo_id", None)
+    root = lerobot_config.get("root", None)
     assert repo_id is not None, "repo id is required"
     action_horizon = config["data"].get("action_horizon", 32)
 
-    data_loader, num_batches = load_lerobot_dataset(repo_id, action_horizon, args)
+    data_loader, num_batches = load_lerobot_dataset(repo_id, root, action_horizon, args)
 
     keys = ["state", "action"]
     stats = {key: normalize.RunningStats() for key in keys}

wall_x/data/utils.py

@@ -40,8 +40,8 @@ KEY_MAPPINGS = {
             "observation.images.leftImg": "left_wrist_view",
             "observation.images.rightImg": "right_wrist_view",
         },
-        "state": "state",
-        "action": "actions",
+        "state": "observation.state",
+        "action": "action",
     },
 }