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

Wandb works, One output dir

Browse Source
This commit is contained in:
Cadene
2024-02-22 12:14:12 +00:00
parent ece89730e6
commit e3643d6146
11 changed files with 200 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
lerobot/scripts/eval.py
View File

@@ -21,19 +21,22 @@ def eval_policy(
save_video: bool = False,
video_dir: Path = None,
fps: int = 15,
env_step: int = None,
wandb=None,
):
rewards = []
if wandb is not None:
assert env_step is not None
sum_rewards = []
max_rewards = []
successes = []
for i in range(num_episodes):
ep_frames = []
def rendering_callback(env, td=None):
nonlocal ep_frames
frame = env.render()
ep_frames.append(frame)
ep_frames.append(env.render())
tensordict = env.reset()
if save_video:
if save_video or wandb:
# render first frame before rollout
rendering_callback(env)
@@ -41,35 +44,54 @@ def eval_policy(
rollout = env.rollout(
max_steps=max_steps,
policy=policy,
callback=rendering_callback if save_video else None,
callback=rendering_callback if save_video or wandb else None,
auto_reset=False,
tensordict=tensordict,
auto_cast_to_device=True,
)
# print(", ".join([f"{x:.3f}" for x in rollout["next", "reward"][:,0].tolist()]))
ep_reward = rollout["next", "reward"].sum()
ep_sum_reward = rollout["next", "reward"].sum()
ep_max_reward = rollout["next", "reward"].max()
ep_success = rollout["next", "success"].any()
rewards.append(ep_reward.item())
sum_rewards.append(ep_sum_reward.item())
max_rewards.append(ep_max_reward.item())
successes.append(ep_success.item())
if save_video:
video_dir.mkdir(parents=True, exist_ok=True)
# TODO(rcadene): make fps configurable
video_path = video_dir / f"eval_episode_{i}.mp4"
imageio.mimsave(video_path, np.stack(ep_frames), fps=fps)
if save_video or wandb:
stacked_frames = np.stack(ep_frames)
if save_video:
video_dir.mkdir(parents=True, exist_ok=True)
video_path = video_dir / f"eval_episode_{i}.mp4"
imageio.mimsave(video_path, stacked_frames, fps=fps)
first_episode = i == 0
if wandb and first_episode:
eval_video = wandb.Video(
stacked_frames.transpose(0, 3, 1, 2), fps=fps, format="mp4"
)
wandb.log({"eval_video": eval_video}, step=env_step)
metrics = {
"avg_reward": np.nanmean(rewards),
"avg_sum_reward": np.nanmean(sum_rewards),
"avg_max_reward": np.nanmean(max_rewards),
"pc_success": np.nanmean(successes) * 100,
}
return metrics
@hydra.main(version_base=None, config_name="default", config_path="../configs")
def eval(cfg: dict):
def eval_cli(cfg: dict):
eval(cfg, out_dir=hydra.core.hydra_config.HydraConfig.get().runtime.output_dir)
def eval(cfg: dict, out_dir=None):
if out_dir is None:
raise NotImplementedError()
assert torch.cuda.is_available()
set_seed(cfg.seed)
print(colored("Log dir:", "yellow", attrs=["bold"]), cfg.log_dir)
print(colored("Log dir:", "yellow", attrs=["bold"]), out_dir)
env = make_env(cfg)
@@ -95,13 +117,14 @@ def eval(cfg: dict):
metrics = eval_policy(
env,
policy=policy,
num_episodes=20,
save_video=True,
video_dir=Path(cfg.video_dir),
video_dir=Path(out_dir) / "eval",
fps=cfg.fps,
max_steps=cfg.episode_length,
num_episodes=cfg.eval_episodes,
)
print(metrics)
if __name__ == "__main__":
eval()
eval_cli()

75
lerobot/scripts/train.py
View File

@@ -20,24 +20,47 @@ from lerobot.scripts.eval import eval_policy
@hydra.main(version_base=None, config_name="default", config_path="../configs")
def train(cfg: dict):
def train_cli(cfg: dict):
train(
cfg,
out_dir=hydra.core.hydra_config.HydraConfig.get().run.dir,
job_name=hydra.core.hydra_config.HydraConfig.get().job.name,
)
def train_notebook(
out_dir=None, job_name=None, config_name="default", config_path="../configs"
):
from hydra import compose, initialize
hydra.core.global_hydra.GlobalHydra.instance().clear()
initialize(config_path=config_path)
cfg = compose(config_name=config_name)
train(cfg, out_dir=out_dir, job_name=job_name)
def train(cfg: dict, out_dir=None, job_name=None):
if out_dir is None:
raise NotImplementedError()
if job_name is None:
raise NotImplementedError()
assert torch.cuda.is_available()
set_seed(cfg.seed)
print(colored("Work dir:", "yellow", attrs=["bold"]), cfg.log_dir)
print(colored("Work dir:", "yellow", attrs=["bold"]), out_dir)
env = make_env(cfg)
policy = TDMPC(cfg)
if cfg.pretrained_model_path:
ckpt_path = (
"/home/rcadene/code/fowm/logs/xarm_lift/all/default/2/models/offline.pt"
)
if "offline" in cfg.pretrained_model_path:
policy.step = 25000
elif "final" in cfg.pretrained_model_path:
policy.step = 100000
else:
raise NotImplementedError()
policy.load(ckpt_path)
# TODO(rcadene): hack for old pretrained models from fowm
if "fowm" in cfg.pretrained_model_path:
if "offline" in cfg.pretrained_model_path:
policy.step = 25000
elif "final" in cfg.pretrained_model_path:
policy.step = 100000
else:
raise NotImplementedError()
policy.load(cfg.pretrained_model_path)
td_policy = TensorDictModule(
policy,
@@ -65,7 +88,7 @@ def train(cfg: dict):
sampler=online_sampler,
)
L = Logger(cfg.log_dir, cfg)
L = Logger(out_dir, job_name, cfg)
online_episode_idx = 0
start_time = time.time()
@@ -95,12 +118,14 @@ def train(cfg: dict):
)
online_buffer.extend(rollout)
ep_reward = rollout["next", "reward"].sum()
ep_sum_reward = rollout["next", "reward"].sum()
ep_max_reward = rollout["next", "reward"].max()
ep_success = rollout["next", "success"].any()
online_episode_idx += 1
rollout_metrics = {
"avg_reward": np.nanmean(ep_reward),
"avg_sum_reward": np.nanmean(ep_sum_reward),
"avg_max_reward": np.nanmean(ep_max_reward),
"pc_success": np.nanmean(ep_success) * 100,
}
num_updates = len(rollout) * cfg.utd
@@ -137,23 +162,23 @@ def train(cfg: dict):
env,
td_policy,
num_episodes=cfg.eval_episodes,
# TODO(rcadene): add step, env_step, L.video
env_step=env_step,
wandb=L._wandb,
)
common_metrics.update(eval_metrics)
L.log(common_metrics, category="eval")
last_log_step = env_step - env_step % cfg.eval_freq
# Save model periodically
# if cfg.save_model and env_step - last_save_step >= cfg.save_freq:
# L.save_model(policy, identifier=env_step)
# print(f"Model has been checkpointed at step {env_step}")
# last_save_step = env_step - env_step % cfg.save_freq
if cfg.save_model and env_step - last_save_step >= cfg.save_freq:
L.save_model(policy, identifier=env_step)
print(f"Model has been checkpointed at step {env_step}")
last_save_step = env_step - env_step % cfg.save_freq
# if cfg.save_model and is_offline and _step >= cfg.offline_steps:
# # save the model after offline training
# L.save_model(policy, identifier="offline")
if cfg.save_model and is_offline and _step >= cfg.offline_steps:
# save the model after offline training
L.save_model(policy, identifier="offline")
step = _step
@@ -177,4 +202,4 @@ def train(cfg: dict):
if __name__ == "__main__":
train()
train_cli()

16
lerobot/scripts/visualize_dataset.py
View File

@@ -15,7 +15,15 @@ from lerobot.common.datasets.factory import make_offline_buffer
@hydra.main(version_base=None, config_name="default", config_path="../configs")
def visualize_dataset(cfg: dict):
def visualize_dataset_cli(cfg: dict):
visualize_dataset(
cfg, out_dir=hydra.core.hydra_config.HydraConfig.get().runtime.output_dir
)
def visualize_dataset(cfg: dict, out_dir=None):
if out_dir is None:
raise NotImplementedError()
sampler = SliceSamplerWithoutReplacement(
num_slices=1,
@@ -40,10 +48,10 @@ def visualize_dataset(cfg: dict):
dim=0,
)
video_dir = Path(cfg.video_dir)
video_dir = Path(out_dir) / "visualize_dataset"
video_dir.mkdir(parents=True, exist_ok=True)
# TODO(rcadene): make fps configurable
video_path = video_dir / f"eval_episode_{ep_idx}.mp4"
video_path = video_dir / f"episode_{ep_idx}.mp4"
assert ep_frames.min().item() >= 0
assert ep_frames.max().item() > 1, "Not mendatory, but sanity check"
@@ -59,4 +67,4 @@ def visualize_dataset(cfg: dict):
if __name__ == "__main__":
visualize_dataset()
visualize_dataset_cli()