tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Clean logging, Refactor

Browse Source
This commit is contained in:
Cadene
2024-02-29 23:13:06 +00:00
parent cb7b375526
commit 0b9027f05e
9 changed files with 229 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
lerobot/common/logger.py
View File

@@ -8,25 +8,6 @@ import pandas as pd
from omegaconf import OmegaConf
from termcolor import colored
CONSOLE_FORMAT = [
("episode", "E", "int"),
("step", "S", "int"),
("avg_sum_reward", "RS", "float"),
("avg_max_reward", "RM", "float"),
("pc_success", "SR", "float"),
("total_time", "T", "time"),
]
AGENT_METRICS = [
"consistency_loss",
"reward_loss",
"value_loss",
"total_loss",
"weighted_loss",
"pi_loss",
"grad_norm",
]
def make_dir(dir_path):
"""Create directory if it does not already exist."""
with contextlib.suppress(OSError):
@@ -80,10 +61,11 @@ class Logger:
"""Primary logger object. Logs either locally or using wandb."""
def __init__(self, log_dir, job_name, cfg):
self._log_dir = make_dir(Path(log_dir))
self._log_dir = Path(log_dir)
self._log_dir.mkdir(parents=True, exist_ok=True)
self._job_name = job_name
self._model_dir = make_dir(self._log_dir / "models")
self._buffer_dir = make_dir(self._log_dir / "buffers")
self._model_dir = self._log_dir / "models"
self._buffer_dir = self._log_dir / "buffers"
self._save_model = cfg.save_model
self._save_buffer = cfg.save_buffer
self._group = cfg_to_group(cfg)
@@ -121,10 +103,11 @@ class Logger:
print(colored("Logs will be synced with wandb.", "blue", attrs=["bold"]))
self._wandb = wandb
def save_model(self, agent, identifier):
def save_model(self, policy, identifier):
if self._save_model:
self._model_dir.mkdir(parents=True, exist_ok=True)
fp = self._model_dir / f"{str(identifier)}.pt"
agent.save(fp)
policy.save(fp)
if self._wandb:
artifact = self._wandb.Artifact(
self._group + "-" + str(self._seed) + "-" + str(identifier),
@@ -134,6 +117,7 @@ class Logger:
self._wandb.log_artifact(artifact)
def save_buffer(self, buffer, identifier):
self._buffer_dir.mkdir(parents=True, exist_ok=True)
fp = self._buffer_dir / f"{str(identifier)}.pkl"
buffer.save(fp)
if self._wandb:
@@ -153,31 +137,13 @@ class Logger:
self._wandb.finish()
print_run(self._cfg, self._eval[-1][-1])
def _format(self, key, value, ty):
if ty == "int":
return f'{colored(key + ":", "yellow")} {int(value):,}'
elif ty == "float":
return f'{colored(key + ":", "yellow")} {value:.01f}'
elif ty == "time":
value = str(datetime.timedelta(seconds=int(value)))
return f'{colored(key + ":", "yellow")} {value}'
else:
raise f"invalid log format type: {ty}"
def _print(self, d, category):
category = colored(category, "blue" if category == "train" else "green")
pieces = [f" {category:<14}"]
for k, disp_k, ty in CONSOLE_FORMAT:
pieces.append(f"{self._format(disp_k, d.get(k, 0), ty):<26}")
print(" ".join(pieces))
def log(self, d, category="train"):
assert category in {"train", "eval"}
def log_dict(self, d, step, mode="train"):
assert mode in {"train", "eval"}
if self._wandb is not None:
for k, v in d.items():
self._wandb.log({category + "/" + k: v}, step=d["step"])
if category == "eval":
keys = ["step", "avg_sum_reward", "avg_max_reward", "pc_success"]
self._eval.append(np.array([d[key] for key in keys]))
pd.DataFrame(np.array(self._eval)).to_csv(self._log_dir / "eval.log", header=keys, index=None)
self._print(d, category)
self._wandb.log({f"{mode}/{k}": v}, step=step)
def log_video(self, video, step, mode="train"):
assert mode in {"train", "eval"}
wandb_video = self._wandb.Video(video, fps=self.cfg.fps, format="mp4")
self._wandb.log({f"{mode}/video": wandb_video}, step=step)