diff --git a/lerobot/configs/policy/act_real_world.yaml b/lerobot/configs/policy/act_real.yaml similarity index 99% rename from lerobot/configs/policy/act_real_world.yaml rename to lerobot/configs/policy/act_real.yaml index c63f69c3..b786160e 100644 --- a/lerobot/configs/policy/act_real_world.yaml +++ b/lerobot/configs/policy/act_real.yaml @@ -24,7 +24,7 @@ override_dataset_stats: training: offline_steps: 80000 online_steps: 0 - eval_freq: 99999999999999 + eval_freq: -1 save_freq: 10000 log_freq: 100 save_model: true diff --git a/lerobot/configs/policy/diffusion_real.yaml b/lerobot/configs/policy/diffusion_real.yaml new file mode 100644 index 00000000..6ee291db --- /dev/null +++ b/lerobot/configs/policy/diffusion_real.yaml @@ -0,0 +1,114 @@ +# @package _global_ + +# Defaults for training for the PushT dataset as per https://github.com/real-stanford/diffusion_policy. +# Note: We do not track EMA model weights as we discovered it does not improve the results. See +# https://github.com/huggingface/lerobot/pull/134 for more details. + +seed: 100000 +dataset_repo_id: lerobot/pusht + +override_dataset_stats: + observation.images.cam_right_wrist: + # stats from imagenet, since we use a pretrained vision model + mean: [[[0.485]], [[0.456]], [[0.406]]] # (c,1,1) + std: [[[0.229]], [[0.224]], [[0.225]]] # (c,1,1) + observation.images.cam_left_wrist: + # stats from imagenet, since we use a pretrained vision model + mean: [[[0.485]], [[0.456]], [[0.406]]] # (c,1,1) + std: [[[0.229]], [[0.224]], [[0.225]]] # (c,1,1) + observation.images.cam_high: + # stats from imagenet, since we use a pretrained vision model + mean: [[[0.485]], [[0.456]], [[0.406]]] # (c,1,1) + std: [[[0.229]], [[0.224]], [[0.225]]] # (c,1,1) + observation.images.cam_low: + # stats from imagenet, since we use a pretrained vision model + mean: [[[0.485]], [[0.456]], [[0.406]]] # (c,1,1) + std: [[[0.229]], [[0.224]], [[0.225]]] # (c,1,1) + +training: + offline_steps: 200000 + online_steps: 0 + eval_freq: -1 + save_freq: 1000 + log_freq: 100 + save_model: true + + batch_size: 64 + grad_clip_norm: 10 + lr: 1.0e-4 + lr_scheduler: cosine + lr_warmup_steps: 500 + adam_betas: [0.95, 0.999] + adam_eps: 1.0e-8 + adam_weight_decay: 1.0e-6 + online_steps_between_rollouts: 1 + + delta_timestamps: + observation.images.cam_right_wrist: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1)]" + observation.images.cam_left_wrist: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1)]" + observation.images.cam_high: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1)]" + observation.images.cam_low: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1)]" + observation.state: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1)]" + action: "[i / ${fps} for i in range(1 - ${policy.n_obs_steps}, 1 - ${policy.n_obs_steps} + ${policy.horizon})]" + +eval: + n_episodes: 50 + batch_size: 50 + +policy: + name: diffusion + + # Input / output structure. + n_obs_steps: 2 + horizon: 16 + n_action_steps: 8 + + input_shapes: + # TODO(rcadene, alexander-soare): add variables for height and width from the dataset/env? + observation.images.cam_right_wrist: [3, 480, 640] + observation.images.cam_left_wrist: [3, 480, 640] + observation.images.cam_high: [3, 480, 640] + observation.images.cam_low: [3, 480, 640] + observation.state: ["${env.state_dim}"] + output_shapes: + action: ["${env.action_dim}"] + + # Normalization / Unnormalization + input_normalization_modes: + observation.images.cam_right_wrist: mean_std + observation.images.cam_left_wrist: mean_std + observation.images.cam_high: mean_std + observation.images.cam_low: mean_std + observation.state: mean_std + output_normalization_modes: + action: mean_std + + # Architecture / modeling. + # Vision backbone. + vision_backbone: resnet18 + crop_shape: null + crop_is_random: False + pretrained_backbone_weights: ResNet18_Weights.IMAGENET1K_V1 + use_group_norm: False + spatial_softmax_num_keypoints: 32 + # Unet. + down_dims: [512, 1024, 2048] + kernel_size: 5 + n_groups: 8 + diffusion_step_embed_dim: 128 + use_film_scale_modulation: True + # Noise scheduler. + noise_scheduler_type: DDPM + num_train_timesteps: 100 + beta_schedule: squaredcos_cap_v2 + beta_start: 0.0001 + beta_end: 0.02 + prediction_type: epsilon # epsilon / sample + clip_sample: True + clip_sample_range: 1.0 + + # Inference + num_inference_steps: 100 + + # Loss computation + do_mask_loss_for_padding: false diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py index c6313792..e7a26425 100644 --- a/lerobot/scripts/train.py +++ b/lerobot/scripts/train.py @@ -329,8 +329,12 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No logging.info("make_dataset") offline_dataset = make_dataset(cfg) - logging.info("make_env") - eval_env = make_env(cfg) + # Create environment used for evaluating checkpoints during training on simulation data. + # On real-world data, no need to create an environment as evaluations are done outside train.py, + # using the eval.py instead, with gym_dora environment and dora-rs. + if cfg.training.eval_freq > 0: + logging.info("make_env") + eval_env = make_env(cfg) logging.info("make_policy") policy = make_policy(hydra_cfg=cfg, dataset_stats=offline_dataset.stats) @@ -356,7 +360,7 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No # Note: this helper will be used in offline and online training loops. def evaluate_and_checkpoint_if_needed(step): - if step % cfg.training.eval_freq == 0: + if cfg.training.eval_freq > 0 and step % cfg.training.eval_freq == 0: logging.info(f"Eval policy at step {step}") eval_info = eval_policy( eval_env, @@ -417,6 +421,13 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No step += 1 + logging.info("End of offline training") + + if cfg.training.online_steps == 0: + if cfg.training.eval_freq > 0: + eval_env.close() + return + # create an env dedicated to online episodes collection from policy rollout online_training_env = make_env(cfg, n_envs=1) @@ -486,9 +497,10 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No step += 1 online_step += 1 + logging.info("End of online training") + eval_env.close() online_training_env.close() - logging.info("End of training") if __name__ == "__main__":