From 4c400b41a53238239a6c0986da0f483bf277012d Mon Sep 17 00:00:00 2001 From: Remi Cadene Date: Sun, 3 Mar 2024 13:22:09 +0000 Subject: [PATCH] Improve log msg in train.py --- lerobot/scripts/train.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py index a537835e..1c63fc97 100644 --- a/lerobot/scripts/train.py +++ b/lerobot/scripts/train.py @@ -170,6 +170,7 @@ def train(cfg: dict, out_dir=None, job_name=None): log_train_info(logger, train_info, step, cfg, offline_buffer, is_offline) if step > 0 and step % cfg.eval_freq == 0: + logging.info(f"Eval policy at step {step}") eval_info, first_video = eval_policy( env, td_policy, @@ -179,10 +180,12 @@ def train(cfg: dict, out_dir=None, job_name=None): log_eval_info(logger, eval_info, step, cfg, offline_buffer, is_offline) if cfg.wandb.enable: logger.log_video(first_video, step, mode="eval") + logging.info("Resume training") if step > 0 and cfg.save_model and step % cfg.save_freq == 0: - logging.info(f"Checkpoint model at step {step}") + logging.info(f"Checkpoint policy at step {step}") logger.save_model(policy, identifier=step) + logging.info("Resume training") step += 1 @@ -227,6 +230,7 @@ def train(cfg: dict, out_dir=None, job_name=None): log_train_info(logger, train_info, step, cfg, offline_buffer, is_offline) if step > 0 and step % cfg.eval_freq == 0: + logging.info(f"Eval policy at step {step}") eval_info, first_video = eval_policy( env, td_policy, @@ -236,10 +240,12 @@ def train(cfg: dict, out_dir=None, job_name=None): log_eval_info(logger, eval_info, step, cfg, offline_buffer, is_offline) if cfg.wandb.enable: logger.log_video(first_video, step, mode="eval") + logging.info("Resume training") if step > 0 and cfg.save_model and step % cfg.save_freq == 0: - logging.info(f"Checkpoint model at step {step}") + logging.info(f"Checkpoint policy at step {step}") logger.save_model(policy, identifier=step) + logging.info("Resume training") step += 1 online_step += 1