Merge pull request #35 from Cadene/switch_between_train_and_eval

Switch between train and eval modes

lerobot/scripts/eval.py

@@ -32,6 +32,7 @@ def eval_policy(
     fps: int = 15,
     return_first_video: bool = False,
 ):
+    policy.eval()
     start = time.time()
     sum_rewards = []
     max_rewards = []

lerobot/scripts/train.py

@@ -181,6 +181,7 @@ def train(cfg: dict, out_dir=None, job_name=None):
         if offline_step == 0:
             logging.info("Start offline training on a fixed dataset")
         # TODO(rcadene): is it ok if step_t=0 = 0 and not 1 as previously done?
+        policy.train()
         train_info = policy.update(offline_buffer, step)
         if step % cfg.log_freq == 0:
             log_train_info(logger, train_info, step, cfg, offline_buffer, is_offline)