[FIX] update navdp training parameters

Author: wzcai99

internnav/model/basemodel/navdp/navdp_policy.py

@@ -7,7 +7,13 @@
 
 from internnav.configs.model.base_encoders import ModelCfg
 from internnav.configs.trainer.exp import ExpCfg
-from internnav.model.encoder.navdp_backbone import *
+from internnav.model.encoder.navdp_backbone import (
+    LearnablePositionalEncoding,
+    NavDP_ImageGoal_Backbone,
+    NavDP_PixelGoal_Backbone,
+    NavDP_RGBD_Backbone,
+    SinusoidalPosEmb,
+)
 
 
 class NavDPModelConfig(PretrainedConfig):
@@ -324,7 +330,7 @@ def predict_pointgoal_batch_action_vel(self, goal_point, input_images, input_dep
                 naction = self.noise_scheduler.step(model_output=noise_pred, timestep=k, sample=naction).prev_sample
 
             critic_values = self.predict_critic(naction, rgbd_embed)
-            all_trajectory = torch.cumsum(naction / 4.0, dim=1)
+            # all_trajectory = torch.cumsum(naction / 4.0, dim=1)
 
             negative_trajectory = torch.cumsum(naction / 4.0, dim=1)[(critic_values).argsort()[0:8]]
             positive_trajectory = torch.cumsum(naction / 4.0, dim=1)[(-critic_values).argsort()[0:8]]
@@ -343,12 +349,12 @@ def predict_nogoal_batch_action_vel(self, input_images, input_depths, sample_num
                 naction = self.noise_scheduler.step(model_output=noise_pred, timestep=k, sample=naction).prev_sample
 
             critic_values = self.predict_critic(naction, rgbd_embed)
-            all_trajectory = torch.cumsum(naction / 4.0, dim=1)
+            # all_trajectory = torch.cumsum(naction / 4.0, dim=1)
 
             negative_trajectory = torch.cumsum(naction / 4.0, dim=1)[(critic_values).argsort()[0:8]]
             positive_trajectory = torch.cumsum(naction / 4.0, dim=1)[(-critic_values).argsort()[0:8]]
             return negative_trajectory, positive_trajectory
 
 
 # if __name__ == "__main__":
-#     policy = NavDPNet(config=)
+#     policy = NavDPNet(config=)

scripts/train/configs/navdp.py

@@ -29,7 +29,7 @@
     ),
     il=IlCfg(
         epochs=1000,
-        batch_size=16,
+        batch_size=32,
         lr=1e-4,
         num_workers=8,
         weight_decay=1e-4,  # TODO
@@ -57,6 +57,7 @@
         prior_sample=False,
         memory_size=8,
         predict_size=24,
+        pixel_channel=4,
         temporal_depth=16,
         heads=8,
         token_dim=384,

scripts/train/train.py

@@ -60,20 +60,16 @@ def on_save(self, args, state, control, **kwargs):
 
 
 def _make_dir(config):
-    config.tensorboard_dir = config.tensorboard_dir % config.name
     config.tensorboard_dir = config.tensorboard_dir % config.name
     config.checkpoint_folder = config.checkpoint_folder % config.name
     config.log_dir = config.log_dir % config.name
     config.output_dir = config.output_dir % config.name
     if not os.path.exists(config.tensorboard_dir):
         os.makedirs(config.tensorboard_dir, exist_ok=True)
-        os.makedirs(config.tensorboard_dir, exist_ok=True)
     if not os.path.exists(config.checkpoint_folder):
         os.makedirs(config.checkpoint_folder, exist_ok=True)
-        os.makedirs(config.checkpoint_folder, exist_ok=True)
     if not os.path.exists(config.log_dir):
         os.makedirs(config.log_dir, exist_ok=True)
-        os.makedirs(config.log_dir, exist_ok=True)
 
 
 def main(config, model_class, model_config_class):
@@ -98,14 +94,12 @@ def main(config, model_class, model_config_class):
             world_size = int(os.getenv('WORLD_SIZE', '1'))
             rank = int(os.getenv('RANK', '0'))
 
-
             # Set CUDA device for each process
             device_id = local_rank
             torch.cuda.set_device(device_id)
             device = torch.device(f'cuda:{device_id}')
             print(f"World size: {world_size}, Local rank: {local_rank}, Global rank: {rank}")
 
-
             # Initialize distributed training environment
             if world_size > 1:
                 try:
@@ -116,7 +110,6 @@ def main(config, model_class, model_config_class):
                     print(f"Distributed initialization FAILED: {str(e)}")
                     world_size = 1
 
-            print("=" * 50)
             print("=" * 50)
             print("After distributed init:")
             print(f"LOCAL_RANK: {local_rank}")
@@ -146,13 +139,10 @@ def main(config, model_class, model_config_class):
                     print(f"Buffer {name} is on wrong device {buffer.device}, should be moved to {device}")
                     buffer.data = buffer.data.to(device)
 
-
             # If distributed training, wrap the model with DDP
             if world_size > 1:
                 model = torch.nn.parallel.DistributedDataParallel(
-                    model, device_ids=[local_rank], 
-                    output_device=local_rank,
-                    find_unused_parameters=True
+                    model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True
                 )
         # ------------ load logger ------------
         train_logger_filename = os.path.join(config.log_dir, 'train.log')
@@ -162,15 +152,10 @@ def main(config, model_class, model_config_class):
                 level=logging.INFO,
                 format_str='%(asctime)-15s %(message)s',
                 filename=train_logger_filename,
-                name='train',
-                level=logging.INFO,
-                format_str='%(asctime)-15s %(message)s',
-                filename=train_logger_filename,
             )
         else:
             # Other processes use console logging
             train_logger = MyLogger(name='train', level=logging.INFO, format_str='%(asctime)-15s %(message)s')
-            train_logger = MyLogger(name='train', level=logging.INFO, format_str='%(asctime)-15s %(message)s')
         transformers_logger = logging.getLogger("transformers")
         if transformers_logger.hasHandlers():
             transformers_logger.handlers = []
@@ -180,18 +165,6 @@ def main(config, model_class, model_config_class):
 
         # ------------ load dataset ------------
         if config.model_name == "navdp":
-            train_dataset_data = NavDP_Base_Datset(
-                config.il.root_dir,
-                config.il.dataset_navdp,
-                config.il.memory_size,
-                config.il.predict_size,
-                config.il.batch_size,
-                config.il.image_size,
-                config.il.scene_scale,
-                preload=config.il.preload,
-                random_digit=config.il.random_digit,
-                prior_sample=config.il.prior_sample,
-            )
             train_dataset_data = NavDP_Base_Datset(
                 config.il.root_dir,
                 config.il.dataset_navdp,
@@ -239,7 +212,6 @@ def main(config, model_class, model_config_class):
                 config.il.lerobot_features_dir,
                 dataset_data=train_dataset_data,
                 batch_size=config.il.batch_size,
-                batch_size=config.il.batch_size,
             )
             collate_fn = rdp_collate_fn(global_batch_size=global_batch_size)
         elif config.model_name == 'navdp':
@@ -255,7 +227,6 @@ def main(config, model_class, model_config_class):
             deepspeed='',
             gradient_checkpointing=False,
             bf16=False,  # fp16=False,
-            bf16=False,  # fp16=False,
             tf32=False,
             per_device_train_batch_size=config.il.batch_size,
             gradient_accumulation_steps=1,
@@ -267,7 +238,6 @@ def main(config, model_class, model_config_class):
             logging_steps=10.0,
             num_train_epochs=config.il.epochs,
             save_strategy='epoch',  # no
-            save_strategy='epoch',  # no
             save_steps=config.il.save_interval_epochs,
             save_total_limit=8,
             report_to=config.il.report_to,
@@ -279,7 +249,6 @@ def main(config, model_class, model_config_class):
             dataloader_drop_last=True,
             disable_tqdm=True,
             log_level="info",
-            log_level="info",
         )
 
         # Create the trainer
@@ -299,17 +268,14 @@ def main(config, model_class, model_config_class):
     except Exception as e:
         import traceback
 
-
         print(f"Unhandled exception: {str(e)}")
         print("Stack trace:")
         traceback.print_exc()
 
-
         # If distributed environment, ensure all processes exit
         if dist.is_initialized():
             dist.destroy_process_group()
 
-
         raise