update vis

Author: zwbx

scripts/download_dataset.py

@@ -107,24 +107,6 @@ def download_task_files(task_name: str, save_dir: str, train_episodes: int, eval
 
     # Download train episodes
     if train_episodes > 0:
-        print(f"Downloading train episodes for {task_name}...")
-        for episode in range(train_episodes):
-            local_file = os.path.join(train_dir, f"episode{episode}", "low_dim_obs.pkl")
-            if os.path.exists(local_file):
-                print(f"  Skipping existing train episode {episode}")
-                continue
-            try:
-                hf_hub_download(
-                    repo_id=repo_id,
-                    repo_type="dataset",
-                    filename=f"train/{task_name}/variation0/episodes/episode{episode}/low_dim_obs.pkl",
-                    local_dir=save_dir,
-                    local_dir_use_symlinks=False,
-                    resume_download=True
-                )
-            except Exception as e:
-                print(f"  Warning: Failed to download train episode {episode}: {e}")
-                success = False
         # Download variation_descriptions.pkl for train
         variation_file = os.path.join(save_dir, "train", task_name, "variation0", "variation_descriptions.pkl")
         if os.path.exists(variation_file):
@@ -143,27 +125,29 @@ def download_task_files(task_name: str, save_dir: str, train_episodes: int, eval
             except Exception as e:
                 print(f"  Warning: Failed to download train variation_descriptions.pkl: {e}")
                 success = False
-    
-    # Download eval episodes
-    if eval_episodes > 0:
-        print(f"Downloading eval episodes for {task_name}...")
-        for episode in range(eval_episodes):
-            local_file = os.path.join(eval_dir, f"episode{episode}", "low_dim_obs.pkl")
+        # Download train episodes
+        print(f"Downloading train episodes for {task_name}...")
+        for episode in range(train_episodes):
+            local_file = os.path.join(train_dir, f"episode{episode}", "low_dim_obs.pkl")
             if os.path.exists(local_file):
-                print(f"  Skipping existing eval episode {episode}")
+                print(f"  Skipping existing train episode {episode}")
                 continue
             try:
                 hf_hub_download(
                     repo_id=repo_id,
                     repo_type="dataset",
-                    filename=f"eval/{task_name}/variation0/episodes/episode{episode}/low_dim_obs.pkl",
+                    filename=f"train/{task_name}/variation0/episodes/episode{episode}/low_dim_obs.pkl",
                     local_dir=save_dir,
                     local_dir_use_symlinks=False,
                     resume_download=True
                 )
             except Exception as e:
-                print(f"  Warning: Failed to download eval episode {episode}: {e}")
+                print(f"  Warning: Failed to download train episode {episode}: {e}")
                 success = False
+
+    
+    # Download eval episodes
+    if eval_episodes > 0:
         # Download variation_descriptions.pkl for eval
         variation_file = os.path.join(save_dir, "eval", task_name, "variation0", "variation_descriptions.pkl")
         if os.path.exists(variation_file):
@@ -182,6 +166,27 @@ def download_task_files(task_name: str, save_dir: str, train_episodes: int, eval
             except Exception as e:
                 print(f"  Warning: Failed to download eval variation_descriptions.pkl: {e}")
                 success = False
+
+        # Download eval episodes
+        print(f"Downloading eval episodes for {task_name}...")
+        for episode in range(eval_episodes):
+            local_file = os.path.join(eval_dir, f"episode{episode}", "low_dim_obs.pkl")
+            if os.path.exists(local_file):
+                print(f"  Skipping existing eval episode {episode}")
+                continue
+            try:
+                hf_hub_download(
+                    repo_id=repo_id,
+                    repo_type="dataset",
+                    filename=f"eval/{task_name}/variation0/episodes/episode{episode}/low_dim_obs.pkl",
+                    local_dir=save_dir,
+                    local_dir_use_symlinks=False,
+                    resume_download=True
+                )
+            except Exception as e:
+                print(f"  Warning: Failed to download eval episode {episode}: {e}")
+                success = False
+
 
     print(f"Completed downloading task: {task_name}")
     return success

src/envs/rlbench/rlbench_env.py

@@ -578,18 +578,7 @@ def get_action_space(self, cfg):
 
     def _load_demos(self, cfg,training=True):
         self.training = training
-
-        '''
-        # The "evaluation dataset" in Huggingface omits point cloud observations to reduce storage requirements.
-        # and the point cloud observation is neccessary for the visualization.
-        # As a result, we use the "training dataset" here for quick project startup.
-        # TODO: upload full eval dataset to huggingface
-        '''
         dataset_root_dir = cfg.dataset_root_train
-        # if training or cfg.debug:
-        #     dataset_root_dir = cfg.dataset_root_train
-        # else:
-        #     dataset_root_dir = cfg.dataset_root_eval
 
         obs_config = _make_obs_config(cfg)
         obs_config_demo = copy.deepcopy(obs_config)

src/workspace.py

@@ -155,9 +155,16 @@ def loop(self):
                 self.logger.log_metrics(updated_metrics, step=self._current_step)
                 print(f"Step {self._current_step}/{num_train_steps}, Loss: {updated_metrics.get('total_loss', 0):.4f}")
 
-            if  self._current_step % self.cfg.vis_every_steps == 0:
-                updated_metrics = self.vis()
-                self.logger.log_metrics(updated_metrics, step=self._current_step)
+            '''
+            # The dataset in Huggingface omits point cloud observations to reduce storage requirements.
+            # and the point cloud observation is neccessary for the visualization.
+            # As a result, the visualization is not available now or you can generate the full eval dataset manually.
+            # TODO: upload full eval dataset to huggingface
+            '''            
+
+            # if  self._current_step % self.cfg.vis_every_steps == 0:
+            #     updated_metrics = self.vis()
+            #     self.logger.log_metrics(updated_metrics, step=self._current_step)
 
             if self._current_step % self.cfg.eval_every_steps == 0:
                 updated_metrics = self.eval()