hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

Improve wandb logging and custom step tracking in logger

Browse Source 
- Modify logger to support multiple custom step keys
- Update logging method to handle custom step keys more flexibly

- Enhance logging of optimization step and frequency
Co-authored-by: michel-aractingi  <michel.aractingi@gmail.com>
This commit is contained in:
AdilZouitine
2025-02-17 10:08:49 +00:00
parent b7a0ffc3b8
commit 279e03b6c8
2 changed files with 24 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
lerobot/scripts/server/learner_server.py
View File

@@ -206,9 +206,9 @@ def start_learner_threads(
server_thread.start()
transition_thread.start()
param_push_thread.start()
# param_push_thread.start()
param_push_thread.join()
# param_push_thread.join()
transition_thread.join()
server_thread.join()
@@ -448,7 +448,9 @@ def add_actor_information_and_train(
policy.update_target_networks()
if optimization_step % cfg.training.log_freq == 0:
logger.log_dict(training_infos, step=optimization_step, mode="train")
training_infos["Optimization step"] = optimization_step
logger.log_dict(d=training_infos, mode="train", custom_step_key="Optimization step")
# logging.info(f"Training infos: {training_infos}")
time_for_one_optimization_step = time.time() - time_for_one_optimization_step
frequency_for_one_optimization_step = 1 / (time_for_one_optimization_step + 1e-9)
@@ -456,9 +458,12 @@ def add_actor_information_and_train(
logging.info(f"[LEARNER] Optimization frequency loop [Hz]: {frequency_for_one_optimization_step}")
logger.log_dict(
{"Optimization frequency loop [Hz]": frequency_for_one_optimization_step},
step=optimization_step,
{
"Optimization frequency loop [Hz]": frequency_for_one_optimization_step,
"Optimization step": optimization_step,
},
mode="train",
custom_step_key="Optimization step",
)
optimization_step += 1