add scripts to sample trajectories

config/lowdim_pusht_dp_tf_20pc.yaml

@@ -0,0 +1,145 @@
+_target_: diffusion_policy.workspace.train_diffusion_transformer_lowdim_workspace.TrainDiffusionTransformerLowdimWorkspace
+action_dim: 2
+checkpoint:
+  save_last_ckpt: true
+  save_last_snapshot: false
+  topk:
+    format_str: epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt
+    k: 5
+    mode: max
+    monitor_key: test_mean_score
+dataloader:
+  batch_size: 256
+  num_workers: 1
+  persistent_workers: false
+  pin_memory: true
+  shuffle: true
+ema:
+  _target_: diffusion_policy.model.diffusion.ema_model.EMAModel
+  inv_gamma: 1.0
+  max_value: 0.9999
+  min_value: 0.0
+  power: 0.75
+  update_after_step: 0
+exp_name: default
+horizon: 10
+keypoint_visible_rate: 1.0
+logging:
+  group: null
+  id: null
+  mode: online
+  name: 2023.01.16-20.11.42_train_diffusion_transformer_lowdim_pusht_lowdim
+  project: diffusion_policy_debug
+  resume: true
+  tags:
+  - train_diffusion_transformer_lowdim
+  - pusht_lowdim
+  - default
+multi_run:
+  run_dir: data/outputs/2023.01.16/20.11.42_train_diffusion_transformer_lowdim_pusht_lowdim
+  wandb_name_base: 2023.01.16-20.11.42_train_diffusion_transformer_lowdim_pusht_lowdim
+n_action_steps: 8
+n_latency_steps: 0
+n_obs_steps: 2
+name: train_diffusion_transformer_lowdim
+obs_as_cond: true
+obs_dim: 20
+optimizer:
+  betas:
+  - 0.9
+  - 0.95
+  learning_rate: 0.0001
+  weight_decay: 0.001
+past_action_visible: false
+policy:
+  _target_: diffusion_policy.policy.diffusion_transformer_lowdim_policy.DiffusionTransformerLowdimPolicy
+  action_dim: 2
+  horizon: 10
+  model:
+    _target_: diffusion_policy.model.diffusion.transformer_for_diffusion.TransformerForDiffusion
+    causal_attn: true
+    cond_dim: 20
+    horizon: 10
+    input_dim: 2
+    n_cond_layers: 0
+    n_emb: 256
+    n_head: 4
+    n_layer: 8
+    n_obs_steps: 2
+    obs_as_cond: true
+    output_dim: 2
+    p_drop_attn: 0.3
+    p_drop_emb: 0.0
+    time_as_cond: true
+  n_action_steps: 8
+  n_obs_steps: 2
+  noise_scheduler:
+    _target_: diffusers.schedulers.scheduling_ddpm.DDPMScheduler
+    beta_end: 0.02
+    beta_schedule: squaredcos_cap_v2
+    beta_start: 0.0001
+    clip_sample: true
+    num_train_timesteps: 100
+    prediction_type: epsilon
+    variance_type: fixed_small
+  num_inference_steps: 100
+  obs_as_cond: true
+  obs_dim: 20
+  pred_action_steps_only: false
+pred_action_steps_only: false
+task:
+  action_dim: 2
+  dataset:
+    _target_: diffusion_policy.dataset.pusht_dataset.PushTLowdimDataset
+    horizon: 10
+    max_train_episodes: 90
+    pad_after: 7
+    pad_before: 1
+    seed: 42
+    val_ratio: 0.02
+    zarr_path: data/pusht/sampled_pusht_v7_replay_20pc.zarr
+  env_runner:
+    _target_: diffusion_policy.env_runner.pusht_keypoints_runner.PushTKeypointsRunner
+    agent_keypoints: false
+    fps: 10
+    keypoint_visible_rate: 1.0
+    legacy_test: true
+    max_steps: 300
+    n_action_steps: 8
+    n_envs: null
+    n_latency_steps: 0
+    n_obs_steps: 2
+    n_test: 50
+    n_test_vis: 4
+    n_train: 6
+    n_train_vis: 2
+    past_action: false
+    test_start_seed: 100000
+    train_start_seed: 0
+  keypoint_dim: 2
+  name: pusht_lowdim
+  obs_dim: 20
+task_name: pusht_lowdim
+training:
+  checkpoint_every: 50
+  debug: false
+  device: cuda:0
+  gradient_accumulate_every: 1
+  lr_scheduler: cosine
+  lr_warmup_steps: 1000
+  max_train_steps: null
+  max_val_steps: null
+  num_epochs: 8000
+  resume: true
+  rollout_every: 50
+  sample_every: 5
+  seed: 42
+  tqdm_interval_sec: 1.0
+  use_ema: true
+  val_every: 1
+val_dataloader:
+  batch_size: 256
+  num_workers: 1
+  persistent_workers: false
+  pin_memory: true
+  shuffle: false

eval.sh

@@ -1,4 +1,4 @@
 # python eval.py --checkpoint data/0550-test_mean_score=0.969.ckpt --output_dir data/pusht_eval_output --device cuda:0
 # python eval.py --checkpoint /mnt/diffusion_policy/data/outputs/2024.11.25/00.43.25_train_diffusion_unet_hybrid_pusht_image/checkpoints/epoch=0100-test_mean_score=0.804.ckpt --output_dir data/pusht_eval_output --device cuda:0
 # python eval.py --checkpoint /mnt/diffusion_policy/data/outputs/2024.11.25/02.30.25_train_diffusion_unet_hybrid_pusht_image/checkpoints/epoch=0350-test_mean_score=0.534.ckpt --output_dir data/pusht_eval_output --device cuda:0
-python eval.py --checkpoint /mnt/diffusion_policy/data/outputs/2024.11.25/04.20.56_train_diffusion_transformer_lowdim_pusht_lowdim/checkpoints/epoch=0550-test_mean_score=0.947.ckpt --output_dir data/pusht_eval_output --device cuda:0
+python eval.py --checkpoint /local2/xingcheng/diffusion_policy/data/outputs/2024.11.25/00.19.26_train_diffusion_transformer_lowdim_pusht_lowdim/checkpoints/epoch=0100-test_mean_score=0.237.ckpt --output_dir data/pusht_eval_output --device cuda:0

run.sh

@@ -1,4 +1,4 @@
-name=debug
+name=sampled_20pc
 task_name=pushT
 ## Image training config
 # Transformer
@@ -9,7 +9,7 @@ task_name=pushT
 
 ## LowDim training config
 # Transformer
-python train.py --config-dir=config --config-name=lowdim_pusht_diffusion_policy_tf.yaml training.seed=42 training.device=cuda:0 hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'
+python train.py --config-dir=config --config-name=lowdim_pusht_dp_tf_20pc training.seed=42 training.device=cuda:0 hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'
 
 # CNN
 # python train.py --config-dir=config --config-name=lowdim_pusht_diffusion_policy_cnn.yaml training.seed=42 training.device=cuda:0 hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'

sample_trajectory.py

@@ -0,0 +1,32 @@
+"""
+Usage:
+python sample_trajectory.py --data_path data/pusht/pusht_cchi_v7_replay.zarr -o data/sampled_pusht.zarr
+"""
+import click
+import numpy as np
+
+from diffusion_policy.common.replay_buffer import ReplayBuffer
+
+@click.command()
+@click.option('-dp', '--data_path', required=True, type=str)
+@click.option('-o', '--output', required=True, type=str)
+@click.option('-sr', '--sample_rate', default=0.2, type=float)
+@click.option('-sd', '--seed', default=42, type=int)
+def main(data_path, output, sample_rate, seed):
+    replay_buffer = ReplayBuffer.copy_from_path(data_path)
+    new_replay_buffer = ReplayBuffer.create_from_path(output, mode='a')
+
+    n_sampled_episodes = int(replay_buffer.n_episodes * sample_rate)
+
+    np.random.seed(seed)
+    sampled_indices = np.random.choice(replay_buffer.n_episodes, n_sampled_episodes, replace=False)
+    for idx in sampled_indices:
+        episode = replay_buffer.get_episode(idx)
+        new_replay_buffer.add_episode(episode, compressors='disk')
+
+    print(f"successfully sampled {n_sampled_episodes} episodes")
+    print(f"saved to {output}")
+
+
+if __name__ == "__main__":
+    main()