Add files using upload-large-folder tool

README.md

@@ -7,11 +7,14 @@ iterations, with emphasis on RLBench2 `take tray out of oven`.
 
 - VLAarchTests benchmark code and generated public benchmark manifest
 - Patched AnyBimanual RLBench runtime used to execute the public oven benchmark
+- Patched PointFlowMatch + `diffusion_policy` source snapshot used to execute `take shoes out of box`
 - Official `katefgroup/3d_flowmatch_actor` PerAct2 checkpoint and public test data
+- Official PointFlowMatch `1717447341-indigo-quokka` checkpoint for `take_shoes_out_of_box`
 - Public AnyBimanual LF baseline weights and comparison logs
 - Verified benchmark reports:
   - oven subset run: `9/10`
   - oven full official run: `95/100 = 0.95`
+  - shoes GPU search: non-zero success verified before later simulator crash
   - hybrid public benchmark smoke outputs
 - DexGarmentLab benchmark-related validation scripts and validation logs
 
@@ -26,12 +29,20 @@ Official oven result artifacts:
 - `reports/3dfa_peract2_take_tray_out_of_oven_subset10/eval_after_official_ttm.json`
 - `reports/3dfa_peract2_take_tray_out_of_oven_full100/eval.json`
 
+Shoes result artifacts:
+
+- `reports/pointflowmatch_take_shoes_out_of_box_ep10_k50_gpu/summary.json`
+- `reports/pointflowmatch_take_shoes_out_of_box_ep10_k50_gpu/run.log`
+
 ## Important Code Paths
 
 - `code/VLAarchtests4/code/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/public_benchmark_package.py`
 - `code/VLAarchtests4/code/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_rlbench_hybrid_smoke.py`
 - `third_party/AnyBimanual/third_party/RLBench/rlbench/bimanual_tasks/bimanual_take_tray_out_of_oven.py`
 - `third_party/AnyBimanual/third_party/RLBench/rlbench/task_ttms/bimanual_take_tray_out_of_oven.ttm`
+- `third_party/PointFlowMatch/pfp/envs/rlbench_env.py`
+- `third_party/PointFlowMatch/pfp/policy/fm_policy.py`
+- `scripts/run_pointflowmatch_take_shoes_out_of_box.sh`
 
 ## External Dependencies Not Mirrored Here
 

docs/ENVIRONMENT_NOTES.md

@@ -45,3 +45,45 @@ xvfb-run -a -s "-screen 0 1400x900x24" python \
   --denoise_timesteps 5 \
   --denoise_model rectified_flow
 ```
+
+The `take shoes out of box` validation path was run from the bundled
+PointFlowMatch source on a Blackwell GPU machine after upgrading the RLBench
+environment to Torch `2.11.0+cu128` / torchvision `0.26.0+cu128` /
+torchaudio `2.11.0+cu128`. The bundled PointFlowMatch tree also contains two
+local compatibility fixes required for this workspace:
+
+- `pfp/envs/rlbench_env.py`
+  - adapts PointFlowMatch to the local RLBench camera naming and observation API
+  - broadens motion-planning failure recovery to handle simulator-side runtime failures
+- `pfp/policy/fm_policy.py`
+  - adds an inference-only fallback when legacy `composer` cannot import on modern Torch
+  - loads checkpoints with `weights_only=False` for PyTorch 2.6+
+
+The shoes evaluation command used here was:
+
+```bash
+scripts/run_pointflowmatch_take_shoes_out_of_box.sh 10 50
+```
+
+That wrapper expands to the equivalent raw command:
+
+```bash
+export PYTHONPATH=third_party/diffusion_policy:third_party/PointFlowMatch:${PYTHONPATH:-}
+export COPPELIASIM_ROOT=/path/to/CoppeliaSim
+export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:$COPPELIASIM_ROOT
+export QT_QPA_PLATFORM_PLUGIN_PATH=$COPPELIASIM_ROOT
+xvfb-run -a -s "-screen 0 1400x900x24" python \
+  third_party/PointFlowMatch/scripts/evaluate.py \
+  log_wandb=False \
+  env_runner.env_config.vis=False \
+  env_runner.num_episodes=10 \
+  env_runner.max_episode_length=200 \
+  policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka \
+  policy.num_k_infer=50
+```
+
+Result note for shoes:
+
+- `reports/pointflowmatch_take_shoes_out_of_box_ep10_k50_gpu/summary.json`
+  records a verified non-zero result before a later RLBench/PyRep crash in the
+  same longer rollout.

models/pointflowmatch_take_shoes_out_of_box/1717447341-indigo-quokka/1717447341-indigo-quokka/config.yaml

@@ -0,0 +1,80 @@
+model:
+  _target_: pfp.policy.fm_policy.FMPolicy
+  x_dim: 266
+  y_dim: 10
+  n_obs_steps: 2
+  n_pred_steps: 32
+  num_k_infer: 10
+  time_conditioning: true
+  norm_pcd_center:
+  - 0.4
+  - 0.0
+  - 1.4
+  augment_data: false
+  noise_type: gaussian
+  noise_scale: 1.0
+  loss_type: l2
+  flow_schedule: exp
+  exp_scale: 4.0
+  snr_sampler: uniform
+  obs_encoder:
+    _target_: pfp.backbones.pointnet.PointNetBackbone
+    embed_dim: 256
+    input_channels: 3
+    input_transform: false
+    use_group_norm: false
+  diffusion_net:
+    _target_: diffusion_policy.model.diffusion.conditional_unet1d.ConditionalUnet1D
+    input_dim: 10
+    global_cond_dim: 532
+    diffusion_step_embed_dim: 256
+    down_dims:
+    - 256
+    - 512
+    - 1024
+    kernel_size: 5
+    n_groups: 8
+    cond_predict_scale: true
+  loss_weights:
+    xyz: 10.0
+    rot6d: 10.0
+    grip: 1.0
+backbone:
+  _target_: pfp.backbones.pointnet.PointNetBackbone
+  embed_dim: 256
+  input_channels: 3
+  input_transform: false
+  use_group_norm: false
+seed: 1234
+epochs: 1500
+log_wandb: true
+task_name: take_shoes_out_of_box
+obs_features_dim: 256
+y_dim: 10
+x_dim: 266
+n_obs_steps: 2
+n_pred_steps: 32
+use_ema: true
+save_each_n_epochs: 500
+obs_mode: pcd
+run_name: null
+dataset:
+  n_obs_steps: 2
+  n_pred_steps: 32
+  subs_factor: 3
+  use_pc_color: false
+  n_points: 4096
+dataloader:
+  batch_size: 128
+  num_workers: 8
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 3.0e-05
+  betas:
+  - 0.95
+  - 0.999
+  eps: 1.0e-08
+  weight_decay: 1.0e-06
+lr_scheduler:
+  name: cosine
+  num_warmup_steps: 5000

reports/pointflowmatch_take_shoes_out_of_box_ep10_k50_gpu/run.log

@@ -0,0 +1,129 @@
+/workspace/envs/rlbench/lib/python3.10/site-packages/requests/__init__.py:109: RequestsDependencyWarning: urllib3 (1.26.13) or chardet (7.4.0.post2)/charset_normalizer (2.1.1) doesn't match a supported version!
+  warnings.warn(
+pytorch3d not installed
+WARNING: Rerun not installed. Visualization will not work.
+WARNING: Rerun not installed. Visualization will not work.
+seed: 5678
+log_wandb: false
+env_runner:
+  num_episodes: 10
+  max_episode_length: 200
+  verbose: true
+  env_config:
+    voxel_size: 0.01
+    headless: true
+    vis: false
+policy:
+  ckpt_name: 1717447341-indigo-quokka/1717447341-indigo-quokka
+  ckpt_episode: ep1500
+  num_k_infer: 50
+
+seed: 5678
+log_wandb: false
+env_runner:
+  num_episodes: 10
+  max_episode_length: 200
+  verbose: true
+  env_config:
+    voxel_size: 0.01
+    headless: true
+    vis: false
+    task_name: take_shoes_out_of_box
+    obs_mode: pcd
+    use_pc_color: false
+    n_points: 4096
+policy:
+  ckpt_name: 1717447341-indigo-quokka/1717447341-indigo-quokka
+  ckpt_episode: ep1500
+  num_k_infer: 50
+  _target_: pfp.policy.fm_policy.FMPolicy.load_from_checkpoint
+model:
+  _target_: pfp.policy.fm_policy.FMPolicy
+  x_dim: 266
+  y_dim: 10
+  n_obs_steps: 2
+  n_pred_steps: 32
+  num_k_infer: 10
+  time_conditioning: true
+  norm_pcd_center:
+  - 0.4
+  - 0.0
+  - 1.4
+  augment_data: false
+  noise_type: gaussian
+  noise_scale: 1.0
+  loss_type: l2
+  flow_schedule: exp
+  exp_scale: 4.0
+  snr_sampler: uniform
+  obs_encoder:
+    _target_: pfp.backbones.pointnet.PointNetBackbone
+    embed_dim: 256
+    input_channels: 3
+    input_transform: false
+    use_group_norm: false
+  diffusion_net:
+    _target_: diffusion_policy.model.diffusion.conditional_unet1d.ConditionalUnet1D
+    input_dim: 10
+    global_cond_dim: 532
+    diffusion_step_embed_dim: 256
+    down_dims:
+    - 256
+    - 512
+    - 1024
+    kernel_size: 5
+    n_groups: 8
+    cond_predict_scale: true
+  loss_weights:
+    xyz: 10.0
+    rot6d: 10.0
+    grip: 1.0
+
+output_dim 10
+[2026-04-03 00:44:27,184][diffusion_policy.model.diffusion.conditional_unet1d][INFO] - number of parameters: 7.285095e+07
+QStandardPaths: XDG_RUNTIME_DIR not set, defaulting to '/tmp/runtime-root'
+
  0%|          | 0/10 [00:00<?, ?it/s][2026-04-03 00:44:29,033][root][WARNING] - single robot
+
 10%|█         | 1/10 [02:31<22:41, 151.33s/it]'NoneType' object has no attribute 'step'
+'NoneType' object has no attribute 'step'
+'NoneType' object has no attribute 'step'
+'NoneType' object has no attribute 'step'
+'NoneType' object has no attribute 'step'
+'NoneType' object has no attribute 'step'
+Steps: 182
+Success: True
+[2026-04-03 00:47:00,210][root][WARNING] - single robot
+
 10%|█         | 1/10 [03:32<31:48, 212.03s/it]
+'NoneType' object has no attribute 'step'
+'NoneType' object has no attribute 'step'
+Error executing job with overrides: ['log_wandb=False', 'env_runner.env_config.vis=False', 'env_runner.num_episodes=10', 'env_runner.max_episode_length=200', 'policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka', 'policy.num_k_infer=50']
+Traceback (most recent call last):
+  File "/workspace/third_party/PointFlowMatch/scripts/evaluate.py", line 52, in main
+    _ = env_runner.run(policy)
+  File "/workspace/third_party/PointFlowMatch/pfp/envs/rlbench_runner.py", line 35, in run
+    reward, terminate = self.env.step(next_robot_state)
+  File "/workspace/third_party/PointFlowMatch/pfp/envs/rlbench_env.py", line 105, in step
+    reward, terminate = self._step_safe(action)
+  File "/workspace/third_party/PointFlowMatch/pfp/envs/rlbench_env.py", line 113, in _step_safe
+    _, reward, terminate = self.task.step(action)
+  File "/workspace/envs/rlbench/lib/python3.10/site-packages/rlbench/task_environment.py", line 109, in step
+    self._action_mode.action(self._scene, action)
+  File "/workspace/envs/rlbench/lib/python3.10/site-packages/rlbench/action_modes/action_mode.py", line 47, in action
+    self.arm_action_mode.action(scene, arm_action, ignore_collisions)
+  File "/workspace/envs/rlbench/lib/python3.10/site-packages/rlbench/action_modes/arm_action_modes.py", line 304, in action
+    success, terminate = scene.task.success()
+  File "/workspace/envs/rlbench/lib/python3.10/site-packages/rlbench/backend/task.py", line 303, in success
+    [cond.condition_met()[0] for cond in self._success_conditions])
+  File "/workspace/envs/rlbench/lib/python3.10/site-packages/rlbench/backend/task.py", line 303, in <listcomp>
+    [cond.condition_met()[0] for cond in self._success_conditions])
+  File "/workspace/envs/rlbench/lib/python3.10/site-packages/rlbench/backend/conditions.py", line 51, in condition_met
+    met = self._detector.is_detected(self._obj)
+  File "/workspace/third_party/PyRep/pyrep/objects/proximity_sensor.py", line 36, in is_detected
+    state, point = sim.simCheckProximitySensor(
+  File "/workspace/third_party/PyRep/pyrep/backend/sim.py", line 345, in simCheckProximitySensor
+    _check_return(state)
+  File "/workspace/third_party/PyRep/pyrep/backend/sim.py", line 27, in _check_return
+    raise RuntimeError(
+RuntimeError: The call failed on the V-REP side. Return value: -1
+
+Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
+QMutex: destroying locked mutex

reports/pointflowmatch_take_shoes_out_of_box_ep10_k50_gpu/summary.json

@@ -0,0 +1,16 @@
+{
+  "task_name": "take_shoes_out_of_box",
+  "model_family": "PointFlowMatch",
+  "checkpoint_name": "1717447341-indigo-quokka",
+  "checkpoint_source": "http://pointflowmatch.cs.uni-freiburg.de/download/1717447341-indigo-quokka.zip",
+  "num_k_infer": 50,
+  "requested_episodes": 10,
+  "completed_episodes_before_crash": 1,
+  "success_count_before_crash": 1,
+  "first_success_episode_index": 0,
+  "first_success_steps": 182,
+  "verified_non_zero_performance": true,
+  "run_status": "partial_success_then_crash",
+  "crash_reason": "RLBench/PyRep RuntimeError after first successful episode during continued rollout",
+  "report_log": "run.log"
+}

scripts/run_pointflowmatch_take_shoes_out_of_box.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+EPISODES="${1:-10}"
+NUM_K_INFER="${2:-50}"
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+PYTHON_BIN="${PYTHON_BIN:-python}"
+COPPELIASIM_ROOT="${COPPELIASIM_ROOT:?set COPPELIASIM_ROOT to your CoppeliaSim root}"
+REPORT_DIR="${REPORT_DIR:-$ROOT/reports/pointflowmatch_take_shoes_out_of_box_ep${EPISODES}_k${NUM_K_INFER}_gpu}"
+
+export PYTHONPATH="$ROOT/third_party/diffusion_policy:$ROOT/third_party/PointFlowMatch:${PYTHONPATH:-}"
+export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:$COPPELIASIM_ROOT"
+export QT_QPA_PLATFORM_PLUGIN_PATH="$COPPELIASIM_ROOT"
+
+mkdir -p "$REPORT_DIR"
+cd "$ROOT/third_party/PointFlowMatch"
+
+xvfb-run -a -s "-screen 0 1400x900x24" "$PYTHON_BIN" scripts/evaluate.py \
+  log_wandb=False \
+  env_runner.env_config.vis=False \
+  env_runner.num_episodes="$EPISODES" \
+  env_runner.max_episode_length=200 \
+  policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka \
+  policy.num_k_infer="$NUM_K_INFER" \
+  2>&1 | tee "$REPORT_DIR/run.log"

third_party/PointFlowMatch/outputs/2026-04-03/00-25-21/.hydra/config.yaml

@@ -0,0 +1,14 @@
+seed: 5678
+log_wandb: false
+env_runner:
+  num_episodes: 1
+  max_episode_length: 200
+  verbose: true
+  env_config:
+    voxel_size: 0.01
+    headless: true
+    vis: false
+policy:
+  ckpt_name: 1717447341-indigo-quokka/1717447341-indigo-quokka
+  ckpt_episode: ep1500
+  num_k_infer: 50

third_party/PointFlowMatch/outputs/2026-04-03/00-25-21/.hydra/hydra.yaml

@@ -0,0 +1,159 @@
+hydra:
+  run:
+    dir: outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}
+  sweep:
+    dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task:
+    - log_wandb=False
+    - env_runner.env_config.vis=False
+    - env_runner.num_episodes=1
+    - env_runner.max_episode_length=200
+    - policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka
+  job:
+    name: evaluate
+    chdir: null
+    override_dirname: env_runner.env_config.vis=False,env_runner.max_episode_length=200,env_runner.num_episodes=1,log_wandb=False,policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka
+    id: ???
+    num: ???
+    config_name: eval
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /workspace/third_party/PointFlowMatch
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /workspace/third_party/PointFlowMatch/conf
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /workspace/third_party/PointFlowMatch/outputs/2026-04-03/00-25-21
+    choices:
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

third_party/PointFlowMatch/outputs/2026-04-03/00-25-21/.hydra/overrides.yaml

@@ -0,0 +1,5 @@
+- log_wandb=False
+- env_runner.env_config.vis=False
+- env_runner.num_episodes=1
+- env_runner.max_episode_length=200
+- policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka

third_party/PointFlowMatch/outputs/2026-04-03/00-44-21/.hydra/config.yaml

@@ -0,0 +1,14 @@
+seed: 5678
+log_wandb: false
+env_runner:
+  num_episodes: 10
+  max_episode_length: 200
+  verbose: true
+  env_config:
+    voxel_size: 0.01
+    headless: true
+    vis: false
+policy:
+  ckpt_name: 1717447341-indigo-quokka/1717447341-indigo-quokka
+  ckpt_episode: ep1500
+  num_k_infer: 50

third_party/PointFlowMatch/outputs/2026-04-03/00-44-21/.hydra/hydra.yaml

@@ -0,0 +1,160 @@
+hydra:
+  run:
+    dir: outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}
+  sweep:
+    dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task:
+    - log_wandb=False
+    - env_runner.env_config.vis=False
+    - env_runner.num_episodes=10
+    - env_runner.max_episode_length=200
+    - policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka
+    - policy.num_k_infer=50
+  job:
+    name: evaluate
+    chdir: null
+    override_dirname: env_runner.env_config.vis=False,env_runner.max_episode_length=200,env_runner.num_episodes=10,log_wandb=False,policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka,policy.num_k_infer=50
+    id: ???
+    num: ???
+    config_name: eval
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /workspace/third_party/PointFlowMatch
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /workspace/third_party/PointFlowMatch/conf
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /workspace/third_party/PointFlowMatch/outputs/2026-04-03/00-44-21
+    choices:
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

third_party/PointFlowMatch/outputs/2026-04-03/00-44-21/.hydra/overrides.yaml

@@ -0,0 +1,6 @@
+- log_wandb=False
+- env_runner.env_config.vis=False
+- env_runner.num_episodes=10
+- env_runner.max_episode_length=200
+- policy.ckpt_name=1717447341-indigo-quokka/1717447341-indigo-quokka
+- policy.num_k_infer=50

third_party/PointFlowMatch/outputs/2026-04-03/00-44-21/evaluate.log

@@ -0,0 +1,3 @@
+[2026-04-03 00:44:27,184][diffusion_policy.model.diffusion.conditional_unet1d][INFO] - number of parameters: 7.285095e+07
+[2026-04-03 00:44:29,033][root][WARNING] - single robot
+[2026-04-03 00:47:00,210][root][WARNING] - single robot

third_party/diffusion_policy/.gitignore

@@ -0,0 +1,140 @@
+bin
+logs
+wandb
+outputs
+data
+data_local
+.vscode
+_wandb
+
+**/.DS_Store
+
+fuse.cfg
+
+*.ai
+
+# Generation results
+results/
+
+ray/auth.json
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

third_party/diffusion_policy/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Columbia Artificial Intelligence and Robotics Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

third_party/diffusion_policy/README.md

@@ -0,0 +1,437 @@
+# Diffusion Policy
+
+[[Project page]](https://diffusion-policy.cs.columbia.edu/)
+[[Paper]](https://diffusion-policy.cs.columbia.edu/#paper)
+[[Data]](https://diffusion-policy.cs.columbia.edu/data/)
+[[Colab (state)]](https://colab.research.google.com/drive/1gxdkgRVfM55zihY9TFLja97cSVZOZq2B?usp=sharing)
+[[Colab (vision)]](https://colab.research.google.com/drive/18GIHeOQ5DyjMN8iIRZL2EKZ0745NLIpg?usp=sharing)
+
+
+[Cheng Chi](http://cheng-chi.github.io/)<sup>1</sup>,
+[Siyuan Feng](https://www.cs.cmu.edu/~sfeng/)<sup>2</sup>,
+[Yilun Du](https://yilundu.github.io/)<sup>3</sup>,
+[Zhenjia Xu](https://www.zhenjiaxu.com/)<sup>1</sup>,
+[Eric Cousineau](https://www.eacousineau.com/)<sup>2</sup>,
+[Benjamin Burchfiel](http://www.benburchfiel.com/)<sup>2</sup>,
+[Shuran Song](https://www.cs.columbia.edu/~shurans/)<sup>1</sup>
+
+<sup>1</sup>Columbia University,
+<sup>2</sup>Toyota Research Institute,
+<sup>3</sup>MIT
+
+<img src="media/teaser.png" alt="drawing" width="100%"/>
+<img src="media/multimodal_sim.png" alt="drawing" width="100%"/>
+
+## 🛝 Try it out!
+Our self-contained Google Colab notebooks is the easiest way to play with Diffusion Policy. We provide separate notebooks for  [state-based environment](https://colab.research.google.com/drive/1gxdkgRVfM55zihY9TFLja97cSVZOZq2B?usp=sharing) and [vision-based environment](https://colab.research.google.com/drive/18GIHeOQ5DyjMN8iIRZL2EKZ0745NLIpg?usp=sharing).
+
+## 🧾 Checkout our experiment logs!
+For each experiment used to generate Table I,II and IV in the [paper](https://diffusion-policy.cs.columbia.edu/#paper), we provide:
+1. A `config.yaml` that contains all parameters needed to reproduce the experiment.
+2. Detailed training/eval `logs.json.txt` for every training step.
+3. Checkpoints for the best `epoch=*-test_mean_score=*.ckpt` and last `latest.ckpt` epoch of each run.
+
+Experiment logs are hosted on our website as nested directories in format:
+`https://diffusion-policy.cs.columbia.edu/data/experiments/<image|low_dim>/<task>/<method>/`
+
+Within each experiment directory you may find:
+```
+.
+├── config.yaml
+├── metrics
+│   └── logs.json.txt
+├── train_0
+│   ├── checkpoints
+│   │   ├── epoch=0300-test_mean_score=1.000.ckpt
+│   │   └── latest.ckpt
+│   └── logs.json.txt
+├── train_1
+│   ├── checkpoints
+│   │   ├── epoch=0250-test_mean_score=1.000.ckpt
+│   │   └── latest.ckpt
+│   └── logs.json.txt
+└── train_2
+    ├── checkpoints
+    │   ├── epoch=0250-test_mean_score=1.000.ckpt
+    │   └── latest.ckpt
+    └── logs.json.txt
+```
+The `metrics/logs.json.txt` file aggregates evaluation metrics from all 3 training runs every 50 epochs using `multirun_metrics.py`. The numbers reported in the paper correspond to `max` and `k_min_train_loss` aggregation keys.
+
+To download all files in a subdirectory, use:
+
+```console
+$ wget --recursive --no-parent --no-host-directories --relative --reject="index.html*" https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/square_ph/diffusion_policy_cnn/
+```
+
+## 🛠️ Installation
+### 🖥️ Simulation
+To reproduce our simulation benchmark results, install our conda environment on a Linux machine with Nvidia GPU. On Ubuntu 20.04 you need to install the following apt packages for mujoco:
+```console
+$ sudo apt install -y libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
+```
+
+We recommend [Mambaforge](https://github.com/conda-forge/miniforge#mambaforge) instead of the standard anaconda distribution for faster installation: 
+```console
+$ mamba env create -f conda_environment.yaml
+```
+
+but you can use conda as well: 
+```console
+$ conda env create -f conda_environment.yaml
+```
+
+The `conda_environment_macos.yaml` file is only for development on MacOS and does not have full support for benchmarks.
+
+### 🦾 Real Robot
+Hardware (for Push-T):
+* 1x [UR5-CB3](https://www.universal-robots.com/cb3) or [UR5e](https://www.universal-robots.com/products/ur5-robot/) ([RTDE Interface](https://www.universal-robots.com/articles/ur/interface-communication/real-time-data-exchange-rtde-guide/) is required)
+* 2x [RealSense D415](https://www.intelrealsense.com/depth-camera-d415/)
+* 1x [3Dconnexion SpaceMouse](https://3dconnexion.com/us/product/spacemouse-wireless/) (for teleop)
+* 1x [Millibar Robotics Manual Tool Changer](https://www.millibar.com/manual-tool-changer/) (only need robot side)
+* 1x 3D printed [End effector](https://cad.onshape.com/documents/a818888644a15afa6cc68ee5/w/2885b48b018cda84f425beca/e/3e8771c2124cee024edd2fed?renderMode=0&uiState=63ffcba6631ca919895e64e5)
+* 1x 3D printed [T-block](https://cad.onshape.com/documents/f1140134e38f6ed6902648d5/w/a78cf81827600e4ff4058d03/e/f35f57fb7589f72e05c76caf?renderMode=0&uiState=63ffcbc9af4a881b344898ee)
+* USB-C cables and screws for RealSense
+
+Software:
+* Ubuntu 20.04.3 (tested)
+* Mujoco dependencies: 
+`sudo apt install libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf`
+* [RealSense SDK](https://github.com/IntelRealSense/librealsense/blob/master/doc/distribution_linux.md)
+* Spacemouse dependencies: 
+`sudo apt install libspnav-dev spacenavd; sudo systemctl start spacenavd`
+* Conda environment `mamba env create -f conda_environment_real.yaml`
+
+## 🖥️ Reproducing Simulation Benchmark Results 
+### Download Training Data
+Under the repo root, create data subdirectory:
+```console
+[diffusion_policy]$ mkdir data && cd data
+```
+
+Download the corresponding zip file from [https://diffusion-policy.cs.columbia.edu/data/training/](https://diffusion-policy.cs.columbia.edu/data/training/)
+```console
+[data]$ wget https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip
+```
+
+Extract training data:
+```console
+[data]$ unzip pusht.zip && rm -f pusht.zip && cd ..
+```
+
+Grab config file for the corresponding experiment:
+```console
+[diffusion_policy]$ wget -O image_pusht_diffusion_policy_cnn.yaml https://diffusion-policy.cs.columbia.edu/data/experiments/image/pusht/diffusion_policy_cnn/config.yaml
+```
+
+### Running for a single seed
+Activate conda environment and login to [wandb](https://wandb.ai) (if you haven't already).
+```console
+[diffusion_policy]$ conda activate robodiff
+(robodiff)[diffusion_policy]$ wandb login
+```
+
+Launch training with seed 42 on GPU 0.
+```console
+(robodiff)[diffusion_policy]$ python train.py --config-dir=. --config-name=image_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}'
+```
+
+This will create a directory in format `data/outputs/yyyy.mm.dd/hh.mm.ss_<method_name>_<task_name>` where configs, logs and checkpoints are written to. The policy will be evaluated every 50 epochs with the success rate logged as `test/mean_score` on wandb, as well as videos for some rollouts.
+```console
+(robodiff)[diffusion_policy]$ tree data/outputs/2023.03.01/20.02.03_train_diffusion_unet_hybrid_pusht_image -I wandb
+data/outputs/2023.03.01/20.02.03_train_diffusion_unet_hybrid_pusht_image
+├── checkpoints
+│   ├── epoch=0000-test_mean_score=0.134.ckpt
+│   └── latest.ckpt
+├── .hydra
+│   ├── config.yaml
+│   ├── hydra.yaml
+│   └── overrides.yaml
+├── logs.json.txt
+├── media
+│   ├── 2k5u6wli.mp4
+│   ├── 2kvovxms.mp4
+│   ├── 2pxd9f6b.mp4
+│   ├── 2q5gjt5f.mp4
+│   ├── 2sawbf6m.mp4
+│   └── 538ubl79.mp4
+└── train.log
+
+3 directories, 13 files
+```
+
+### Running for multiple seeds
+Launch local ray cluster. For large scale experiments, you might want to setup an [AWS cluster with autoscaling](https://docs.ray.io/en/master/cluster/vms/user-guides/launching-clusters/aws.html). All other commands remain the same.
+```console
+(robodiff)[diffusion_policy]$ export CUDA_VISIBLE_DEVICES=0,1,2  # select GPUs to be managed by the ray cluster
+(robodiff)[diffusion_policy]$ ray start --head --num-gpus=3
+```
+
+Launch a ray client which will start 3 training workers (3 seeds) and 1 metrics monitor worker.
+```console
+(robodiff)[diffusion_policy]$ python ray_train_multirun.py --config-dir=. --config-name=image_pusht_diffusion_policy_cnn.yaml --seeds=42,43,44 --monitor_key=test/mean_score -- multi_run.run_dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}' multi_run.wandb_name_base='${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}'
+```
+
+In addition to the wandb log written by each training worker individually, the metrics monitor worker will log to wandb project `diffusion_policy_metrics` for the metrics aggregated from all 3 training runs. Local config, logs and checkpoints will be written to `data/outputs/yyyy.mm.dd/hh.mm.ss_<method_name>_<task_name>` in a directory structure identical to our [training logs](https://diffusion-policy.cs.columbia.edu/data/experiments/):
+```console
+(robodiff)[diffusion_policy]$ tree data/outputs/2023.03.01/22.13.58_train_diffusion_unet_hybrid_pusht_image -I 'wandb|media'
+data/outputs/2023.03.01/22.13.58_train_diffusion_unet_hybrid_pusht_image
+├── config.yaml
+├── metrics
+│   ├── logs.json.txt
+│   ├── metrics.json
+│   └── metrics.log
+├── train_0
+│   ├── checkpoints
+│   │   ├── epoch=0000-test_mean_score=0.174.ckpt
+│   │   └── latest.ckpt
+│   ├── logs.json.txt
+│   └── train.log
+├── train_1
+│   ├── checkpoints
+│   │   ├── epoch=0000-test_mean_score=0.131.ckpt
+│   │   └── latest.ckpt
+│   ├── logs.json.txt
+│   └── train.log
+└── train_2
+    ├── checkpoints
+    │   ├── epoch=0000-test_mean_score=0.105.ckpt
+    │   └── latest.ckpt
+    ├── logs.json.txt
+    └── train.log
+
+7 directories, 16 files
+```
+### 🆕 Evaluate Pre-trained Checkpoints
+Download a checkpoint from the published training log folders, such as [https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/pusht/diffusion_policy_cnn/train_0/checkpoints/epoch=0550-test_mean_score=0.969.ckpt](https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/pusht/diffusion_policy_cnn/train_0/checkpoints/epoch=0550-test_mean_score=0.969.ckpt).
+
+Run the evaluation script:
+```console
+(robodiff)[diffusion_policy]$ python eval.py --checkpoint data/0550-test_mean_score=0.969.ckpt --output_dir data/pusht_eval_output --device cuda:0
+```
+
+This will generate the following directory structure:
+```console
+(robodiff)[diffusion_policy]$ tree data/pusht_eval_output
+data/pusht_eval_output
+├── eval_log.json
+└── media
+    ├── 1fxtno84.mp4
+    ├── 224l7jqd.mp4
+    ├── 2fo4btlf.mp4
+    ├── 2in4cn7a.mp4
+    ├── 34b3o2qq.mp4
+    └── 3p7jqn32.mp4
+
+1 directory, 7 files
+```
+
+`eval_log.json` contains metrics that is logged to wandb during training:
+```console
+(robodiff)[diffusion_policy]$ cat data/pusht_eval_output/eval_log.json
+{
+  "test/mean_score": 0.9150393806777066,
+  "test/sim_max_reward_4300000": 1.0,
+  "test/sim_max_reward_4300001": 0.9872969750774386,
+...
+  "train/sim_video_1": "data/pusht_eval_output//media/2fo4btlf.mp4"
+}
+```
+
+## 🦾 Demo, Training and Eval on a Real Robot
+Make sure your UR5 robot is running and accepting command from its network interface (emergency stop button within reach at all time), your RealSense cameras plugged in to your workstation (tested with `realsense-viewer`) and your SpaceMouse connected with the `spacenavd` daemon running (verify with `systemctl status spacenavd`).
+
+Start the demonstration collection script. Press "C" to start recording. Use SpaceMouse to move the robot. Press "S" to stop recording. 
+```console
+(robodiff)[diffusion_policy]$ python demo_real_robot.py -o data/demo_pusht_real --robot_ip 192.168.0.204
+```
+
+This should result in a demonstration dataset in `data/demo_pusht_real` with in the same structure as our example [real Push-T training dataset](https://diffusion-policy.cs.columbia.edu/data/training/pusht_real.zip).
+
+To train a Diffusion Policy, launch training with config:
+```console
+(robodiff)[diffusion_policy]$ python train.py --config-name=train_diffusion_unet_real_image_workspace task.dataset_path=data/demo_pusht_real
+```
+Edit [`diffusion_policy/config/task/real_pusht_image.yaml`](./diffusion_policy/config/task/real_pusht_image.yaml) if your camera setup is different.
+
+Assuming the training has finished and you have a checkpoint at `data/outputs/blah/checkpoints/latest.ckpt`, launch the evaluation script with:
+```console
+python eval_real_robot.py -i data/outputs/blah/checkpoints/latest.ckpt -o data/eval_pusht_real --robot_ip 192.168.0.204
+```
+Press "C" to start evaluation (handing control over to the policy). Press "S" to stop the current episode.
+
+## 🗺️ Codebase Tutorial
+This codebase is structured under the requirement that:
+1. implementing `N` tasks and `M` methods will only require `O(N+M)` amount of code instead of `O(N*M)`
+2. while retaining maximum flexibility.
+
+To achieve this requirement, we 
+1. maintained a simple unified interface between tasks and methods and 
+2. made the implementation of the tasks and the methods independent of each other. 
+
+These design decisions come at the cost of code repetition between the tasks and the methods. However, we believe that the benefit of being able to add/modify task/methods without affecting the remainder and being able understand a task/method by reading the code linearly outweighs the cost of copying and pasting 😊.
+
+### The Split
+On the task side, we have:
+* `Dataset`: adapts a (third-party) dataset to the interface.
+* `EnvRunner`: executes a `Policy` that accepts the interface and produce logs and metrics.
+* `config/task/<task_name>.yaml`: contains all information needed to construct `Dataset` and `EnvRunner`.
+* (optional) `Env`: an `gym==0.21.0` compatible class that encapsulates the task environment.
+
+On the policy side, we have:
+* `Policy`: implements inference according to the interface and part of the training process.
+* `Workspace`: manages the life-cycle of training and evaluation (interleaved) of a method. 
+* `config/<workspace_name>.yaml`: contains all information needed to construct `Policy` and `Workspace`.
+
+### The Interface
+#### Low Dim
+A [`LowdimPolicy`](./diffusion_policy/policy/base_lowdim_policy.py) takes observation dictionary:
+- `"obs":` Tensor of shape `(B,To,Do)`
+
+and predicts action dictionary:
+- `"action": ` Tensor of shape `(B,Ta,Da)`
+
+A [`LowdimDataset`](./diffusion_policy/dataset/base_dataset.py) returns a sample of dictionary:
+- `"obs":` Tensor of shape `(To, Do)`
+- `"action":` Tensor of shape `(Ta, Da)`
+
+Its `get_normalizer` method returns a [`LinearNormalizer`](./diffusion_policy/model/common/normalizer.py) with keys `"obs","action"`.
+
+The `Policy` handles normalization on GPU with its copy of the `LinearNormalizer`. The parameters of the `LinearNormalizer` is saved as part of the `Policy`'s weights checkpoint.
+
+#### Image
+A [`ImagePolicy`](./diffusion_policy/policy/base_image_policy.py) takes observation dictionary:
+- `"key0":` Tensor of shape `(B,To,*)`
+- `"key1":` Tensor of shape e.g. `(B,To,H,W,3)` ([0,1] float32)
+
+and predicts action dictionary:
+- `"action": ` Tensor of shape `(B,Ta,Da)`
+
+A [`ImageDataset`](./diffusion_policy/dataset/base_dataset.py) returns a sample of dictionary:
+- `"obs":` Dict of
+    - `"key0":` Tensor of shape `(To, *)`
+    - `"key1":` Tensor fo shape `(To,H,W,3)`
+- `"action":` Tensor of shape `(Ta, Da)`
+
+Its `get_normalizer` method returns a [`LinearNormalizer`](./diffusion_policy/model/common/normalizer.py) with keys `"key0","key1","action"`.
+
+#### Example
+```
+To = 3
+Ta = 4
+T = 6
+|o|o|o|
+| | |a|a|a|a|
+|o|o|
+| |a|a|a|a|a|
+| | | | |a|a|
+```
+Terminology in the paper: `varname` in the codebase
+- Observation Horizon: `To|n_obs_steps`
+- Action Horizon: `Ta|n_action_steps`
+- Prediction Horizon: `T|horizon`
+
+The classical (e.g. MDP) single step observation/action formulation is included as a special case where `To=1` and `Ta=1`.
+
+## 🔩 Key Components
+### `Workspace`
+A `Workspace` object encapsulates all states and code needed to run an experiment. 
+* Inherits from [`BaseWorkspace`](./diffusion_policy/workspace/base_workspace.py).
+* A single `OmegaConf` config object generated by `hydra` should contain all information needed to construct the Workspace object and running experiments. This config correspond to `config/<workspace_name>.yaml` + hydra overrides.
+* The `run` method contains the entire pipeline for the experiment.
+* Checkpoints happen at the `Workspace` level. All training states implemented as object attributes are automatically saved by the `save_checkpoint` method.
+* All other states for the experiment should be implemented as local variables in the `run` method.
+
+The entrypoint for training is `train.py` which uses `@hydra.main` decorator. Read [hydra](https://hydra.cc/)'s official documentation for command line arguments and config overrides. For example, the argument `task=<task_name>` will replace the `task` subtree of the config with the content of `config/task/<task_name>.yaml`, thereby selecting the task to run for this experiment.
+
+### `Dataset`
+A `Dataset` object:
+* Inherits from `torch.utils.data.Dataset`.
+* Returns a sample conforming to [the interface](#the-interface) depending on whether the task has Low Dim or Image observations.
+* Has a method `get_normalizer` that returns a `LinearNormalizer` conforming to [the interface](#the-interface).
+
+Normalization is a very common source of bugs during project development. It is sometimes helpful to print out the specific `scale` and `bias` vectors used for each key in the `LinearNormalizer`.
+
+Most of our implementations of `Dataset` uses a combination of [`ReplayBuffer`](#replaybuffer) and [`SequenceSampler`](./diffusion_policy/common/sampler.py) to generate samples. Correctly handling padding at the beginning and the end of each demonstration episode according to `To` and `Ta` is important for good performance. Please read our [`SequenceSampler`](./diffusion_policy/common/sampler.py) before implementing your own sampling method.
+
+### `Policy`
+A `Policy` object:
+* Inherits from `BaseLowdimPolicy` or `BaseImagePolicy`.
+* Has a method `predict_action` that given observation dict, predicts actions conforming to [the interface](#the-interface).
+* Has a method `set_normalizer` that takes in a `LinearNormalizer` and handles observation/action normalization internally in the policy.
+* (optional) Might has a method `compute_loss` that takes in a batch and returns the loss to be optimized.
+* (optional) Usually each `Policy` class correspond to a `Workspace` class due to the differences of training and evaluation process between methods.
+
+### `EnvRunner`
+A `EnvRunner` object abstracts away the subtle differences between different task environments.
+* Has a method `run` that takes a `Policy` object for evaluation, and returns a dict of logs and metrics. Each value should be compatible with `wandb.log`. 
+
+To maximize evaluation speed, we usually vectorize environments using our modification of [`gym.vector.AsyncVectorEnv`](./diffusion_policy/gym_util/async_vector_env.py) which runs each individual environment in a separate process (workaround python GIL). 
+
+⚠️ Since subprocesses are launched using `fork` on linux, you need to be specially careful for environments that creates its OpenGL context during initialization (e.g. robosuite) which, once inherited by the child process memory space, often causes obscure bugs like segmentation fault. As a workaround, you can provide a `dummy_env_fn` that constructs an environment without initializing OpenGL.
+
+### `ReplayBuffer`
+The [`ReplayBuffer`](./diffusion_policy/common/replay_buffer.py) is a key data structure for storing a demonstration dataset both in-memory and on-disk with chunking and compression. It makes heavy use of the [`zarr`](https://zarr.readthedocs.io/en/stable/index.html) format but also has a `numpy` backend for lower access overhead.
+
+On disk, it can be stored as a nested directory (e.g. `data/pusht_cchi_v7_replay.zarr`) or a zip file (e.g. `data/robomimic/datasets/square/mh/image_abs.hdf5.zarr.zip`).
+
+Due to the relative small size of our datasets, it's often possible to store the entire image-based dataset in RAM with [`Jpeg2000` compression](./diffusion_policy/codecs/imagecodecs_numcodecs.py) which eliminates disk IO during training at the expense increasing of CPU workload.
+
+Example:
+```
+data/pusht_cchi_v7_replay.zarr
+ ├── data
+ │   ├── action (25650, 2) float32
+ │   ├── img (25650, 96, 96, 3) float32
+ │   ├── keypoint (25650, 9, 2) float32
+ │   ├── n_contacts (25650, 1) float32
+ │   └── state (25650, 5) float32
+ └── meta
+     └── episode_ends (206,) int64
+```
+
+Each array in `data` stores one data field from all episodes concatenated along the first dimension (time). The `meta/episode_ends` array stores the end index for each episode along the fist dimension.
+
+### `SharedMemoryRingBuffer`
+The [`SharedMemoryRingBuffer`](./diffusion_policy/shared_memory/shared_memory_ring_buffer.py) is a lock-free FILO data structure used extensively in our [real robot implementation](./diffusion_policy/real_world) to utilize multiple CPU cores while avoiding pickle serialization and locking overhead for `multiprocessing.Queue`. 
+
+As an example, we would like to get the most recent `To` frames from 5 RealSense cameras. We launch 1 realsense SDK/pipeline per process using [`SingleRealsense`](./diffusion_policy/real_world/single_realsense.py), each continuously writes the captured images into a `SharedMemoryRingBuffer` shared with the main process. We can very quickly get the last `To` frames in the main process due to the FILO nature of `SharedMemoryRingBuffer`.
+
+We also implemented [`SharedMemoryQueue`](./diffusion_policy/shared_memory/shared_memory_queue.py) for FIFO, which is used in [`RTDEInterpolationController`](./diffusion_policy/real_world/rtde_interpolation_controller.py).
+
+### `RealEnv`
+In contrast to [OpenAI Gym](https://gymnasium.farama.org/), our polices interact with the environment asynchronously. In [`RealEnv`](./diffusion_policy/real_world/real_env.py), the `step` method in `gym` is split into two methods: `get_obs` and `exec_actions`. 
+
+The `get_obs` method returns the latest observation from `SharedMemoryRingBuffer` as well as their corresponding timestamps. This method can be call at any time during an evaluation episode.
+
+The `exec_actions` method accepts a sequence of actions and timestamps for the expected time of execution for each step. Once called, the actions are simply enqueued to the `RTDEInterpolationController`, and the method returns without blocking for execution.
+
+## 🩹 Adding a Task
+Read and imitate:
+* `diffusion_policy/dataset/pusht_image_dataset.py`
+* `diffusion_policy/env_runner/pusht_image_runner.py`
+* `diffusion_policy/config/task/pusht_image.yaml`
+
+Make sure that `shape_meta` correspond to input and output shapes for your task. Make sure `env_runner._target_` and `dataset._target_` point to the new classes you have added. When training, add `task=<your_task_name>` to `train.py`'s arguments.
+
+## 🩹 Adding a Method
+Read and imitate:
+* `diffusion_policy/workspace/train_diffusion_unet_image_workspace.py`
+* `diffusion_policy/policy/diffusion_unet_image_policy.py`
+* `diffusion_policy/config/train_diffusion_unet_image_workspace.yaml`
+
+Make sure your workspace yaml's `_target_` points to the new workspace class you created.
+
+## 🏷️ License
+This repository is released under the MIT license. See [LICENSE](LICENSE) for additional details.
+
+## 🙏 Acknowledgement
+* Our [`ConditionalUnet1D`](./diffusion_policy/model/diffusion/conditional_unet1d.py) implementation is adapted from [Planning with Diffusion](https://github.com/jannerm/diffuser).
+* Our [`TransformerForDiffusion`](./diffusion_policy/model/diffusion/transformer_for_diffusion.py) implementation is adapted from [MinGPT](https://github.com/karpathy/minGPT).
+* The [BET](./diffusion_policy/model/bet) baseline is adapted from [its original repo](https://github.com/notmahi/bet).
+* The [IBC](./diffusion_policy/policy/ibc_dfo_lowdim_policy.py) baseline is adapted from [Kevin Zakka's reimplementation](https://github.com/kevinzakka/ibc).
+* The [Robomimic](https://github.com/ARISE-Initiative/robomimic) tasks and [`ObservationEncoder`](https://github.com/ARISE-Initiative/robomimic/blob/master/robomimic/models/obs_nets.py) are used extensively in this project.
+* The [Push-T](./diffusion_policy/env/pusht) task is adapted from [IBC](https://github.com/google-research/ibc).
+* The [Block Pushing](./diffusion_policy/env/block_pushing) task is adapted from [BET](https://github.com/notmahi/bet) and [IBC](https://github.com/google-research/ibc).
+* The [Kitchen](./diffusion_policy/env/kitchen) task is adapted from [BET](https://github.com/notmahi/bet) and [Relay Policy Learning](https://github.com/google-research/relay-policy-learning).
+* Our [shared_memory](./diffusion_policy/shared_memory) data structures are heavily inspired by [shared-ndarray2](https://gitlab.com/osu-nrsg/shared-ndarray2).

third_party/diffusion_policy/conda_environment.yaml

@@ -0,0 +1,65 @@
+name: robodiff
+channels:
+  - pytorch
+  - pytorch3d
+  - nvidia
+  - conda-forge
+dependencies:
+  - python=3.9
+  - pip=22.2.2
+  - cudatoolkit=11.6
+  - pytorch=1.12.1
+  - torchvision=0.13.1
+  - pytorch3d=0.7.0
+  - numpy=1.23.3
+  - numba==0.56.4
+  - scipy==1.9.1
+  - py-opencv=4.6.0
+  - cffi=1.15.1
+  - ipykernel=6.16
+  - matplotlib=3.6.1
+  - zarr=2.12.0
+  - numcodecs=0.10.2
+  - h5py=3.7.0
+  - hydra-core=1.2.0
+  - einops=0.4.1
+  - tqdm=4.64.1
+  - dill=0.3.5.1
+  - scikit-video=1.1.11
+  - scikit-image=0.19.3
+  - gym=0.21.0
+  - pymunk=6.2.1
+  - wandb=0.13.3
+  - threadpoolctl=3.1.0
+  - shapely=1.8.4
+  - cython=0.29.32
+  - imageio=2.22.0
+  - imageio-ffmpeg=0.4.7
+  - termcolor=2.0.1
+  - tensorboard=2.10.1
+  - tensorboardx=2.5.1
+  - psutil=5.9.2
+  - click=8.0.4
+  - boto3=1.24.96
+  - accelerate=0.13.2
+  - datasets=2.6.1
+  - diffusers=0.11.1
+  - av=10.0.0
+  - cmake=3.24.3
+  # trick to avoid cpu affinity issue described in https://github.com/pytorch/pytorch/issues/99625
+  - llvm-openmp=14
+  # trick to force reinstall imagecodecs via pip
+  - imagecodecs==2022.8.8
+  - pip:
+    - ray[default,tune]==2.2.0
+    # requires mujoco py dependencies libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
+    - free-mujoco-py==2.1.6
+    - pygame==2.1.2
+    - pybullet-svl==3.1.6.4
+    - robosuite @ https://github.com/cheng-chi/robosuite/archive/277ab9588ad7a4f4b55cf75508b44aa67ec171f0.tar.gz
+    - robomimic==0.2.0
+    - pytorchvideo==0.1.5
+    # pip package required for jpeg-xl
+    - imagecodecs==2022.9.26
+    - r3m @ https://github.com/facebookresearch/r3m/archive/b2334e726887fa0206962d7984c69c5fb09cceab.tar.gz
+    - dm-control==1.0.9

third_party/diffusion_policy/conda_environment_macos.yaml

@@ -0,0 +1,55 @@
+name: robodiff
+channels:
+  - pytorch
+  - conda-forge
+dependencies:
+  - python=3.9
+  - pip=22.2.2
+  - pytorch=1.12.1
+  - torchvision=0.13.1
+  - numpy=1.23.3
+  - numba==0.56.4
+  - scipy==1.9.1
+  - py-opencv=4.6.0
+  - cffi=1.15.1
+  - ipykernel=6.16
+  - matplotlib=3.6.1
+  - zarr=2.12.0
+  - numcodecs=0.10.2
+  - h5py=3.7.0
+  - hydra-core=1.2.0
+  - einops=0.4.1
+  - tqdm=4.64.1
+  - dill=0.3.5.1
+  - scikit-video=1.1.11
+  - scikit-image=0.19.3
+  - gym=0.21.0
+  - pymunk=6.2.1
+  - wandb=0.13.3
+  - threadpoolctl=3.1.0
+  - shapely=1.8.4
+  - cython=0.29.32
+  - imageio=2.22.0
+  - imageio-ffmpeg=0.4.7
+  - termcolor=2.0.1
+  - tensorboard=2.10.1
+  - tensorboardx=2.5.1
+  - psutil=5.9.2
+  - click=8.0.4
+  - boto3=1.24.96
+  - accelerate=0.13.2
+  - datasets=2.6.1
+  - diffusers=0.11.1
+  - av=10.0.0
+  - cmake=3.24.3
+  # trick to force reinstall imagecodecs via pip
+  - imagecodecs==2022.8.8
+  - pip:
+    - ray[default,tune]==2.2.0
+    - pygame==2.1.2
+    - robomimic==0.2.0
+    - pytorchvideo==0.1.5
+    - atomics==1.0.2
+    # No support for jpeg-xl for MacOS
+    - imagecodecs==2022.9.26
+    - r3m @ https://github.com/facebookresearch/r3m/archive/b2334e726887fa0206962d7984c69c5fb09cceab.tar.gz

third_party/diffusion_policy/conda_environment_real.yaml

@@ -0,0 +1,73 @@
+name: robodiff
+channels:
+  - pytorch
+  - pytorch3d
+  - nvidia
+  - conda-forge
+dependencies:
+  - python=3.9
+  - pip=22.2.2
+  - cudatoolkit=11.6
+  - pytorch=1.12.1
+  - torchvision=0.13.1
+  - pytorch3d=0.7.0
+  - numpy=1.23.3
+  - numba==0.56.4
+  - scipy==1.9.1
+  - py-opencv=4.6.0
+  - cffi=1.15.1
+  - ipykernel=6.16
+  - matplotlib=3.6.1
+  - zarr=2.12.0
+  - numcodecs=0.10.2
+  - h5py=3.7.0
+  - hydra-core=1.2.0
+  - einops=0.4.1
+  - tqdm=4.64.1
+  - dill=0.3.5.1
+  - scikit-video=1.1.11
+  - scikit-image=0.19.3
+  - gym=0.21.0
+  - pymunk=6.2.1
+  - wandb=0.13.3
+  - threadpoolctl=3.1.0
+  - shapely=1.8.4
+  - cython=0.29.32
+  - imageio=2.22.0
+  - imageio-ffmpeg=0.4.7
+  - termcolor=2.0.1
+  - tensorboard=2.10.1
+  - tensorboardx=2.5.1
+  - psutil=5.9.2
+  - click=8.0.4
+  - boto3=1.24.96
+  - accelerate=0.13.2
+  - datasets=2.6.1
+  - diffusers=0.11.1
+  - av=10.0.0
+  - cmake=3.24.3
+  # trick to avoid cpu affinity issue described in https://github.com/pytorch/pytorch/issues/99625
+  - llvm-openmp=14
+  # trick to force reinstall imagecodecs via pip
+  - imagecodecs==2022.8.8
+  - pip:
+    - ray[default,tune]==2.2.0
+    # requires mujoco py dependencies libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
+    - free-mujoco-py==2.1.6
+    - pygame==2.1.2
+    - pybullet-svl==3.1.6.4
+    - robosuite @ https://github.com/cheng-chi/robosuite/archive/277ab9588ad7a4f4b55cf75508b44aa67ec171f0.tar.gz
+    - robomimic==0.2.0
+    - pytorchvideo==0.1.5
+    # requires librealsense https://github.com/IntelRealSense/librealsense/blob/master/doc/distribution_linux.md
+    - pyrealsense2==2.51.1.4348
+    # requires apt install libspnav-dev spacenavd; systemctl start spacenavd
+    - spnav @ https://github.com/cheng-chi/spnav/archive/c1c938ebe3cc542db4685e0d13850ff1abfdb943.tar.gz
+    - ur-rtde==1.5.5
+    - atomics==1.0.2
+    # pip package required for jpeg-xl
+    - imagecodecs==2022.9.26
+    - r3m @ https://github.com/facebookresearch/r3m/archive/b2334e726887fa0206962d7984c69c5fb09cceab.tar.gz
+    - dm-control==1.0.9
+    - pynput==1.7.6
+  

third_party/diffusion_policy/demo_pusht.py

@@ -0,0 +1,120 @@
+import numpy as np
+import click
+from diffusion_policy.common.replay_buffer import ReplayBuffer
+from diffusion_policy.env.pusht.pusht_keypoints_env import PushTKeypointsEnv
+import pygame
+
+@click.command()
+@click.option('-o', '--output', required=True)
+@click.option('-rs', '--render_size', default=96, type=int)
+@click.option('-hz', '--control_hz', default=10, type=int)
+def main(output, render_size, control_hz):
+    """
+    Collect demonstration for the Push-T task.
+    
+    Usage: python demo_pusht.py -o data/pusht_demo.zarr
+    
+    This script is compatible with both Linux and MacOS.
+    Hover mouse close to the blue circle to start.
+    Push the T block into the green area. 
+    The episode will automatically terminate if the task is succeeded.
+    Press "Q" to exit.
+    Press "R" to retry.
+    Hold "Space" to pause.
+    """
+    
+    # create replay buffer in read-write mode
+    replay_buffer = ReplayBuffer.create_from_path(output, mode='a')
+
+    # create PushT env with keypoints
+    kp_kwargs = PushTKeypointsEnv.genenerate_keypoint_manager_params()
+    env = PushTKeypointsEnv(render_size=render_size, render_action=False, **kp_kwargs)
+    agent = env.teleop_agent()
+    clock = pygame.time.Clock()
+    
+    # episode-level while loop
+    while True:
+        episode = list()
+        # record in seed order, starting with 0
+        seed = replay_buffer.n_episodes
+        print(f'starting seed {seed}')
+
+        # set seed for env
+        env.seed(seed)
+        
+        # reset env and get observations (including info and render for recording)
+        obs = env.reset()
+        info = env._get_info()
+        img = env.render(mode='human')
+        
+        # loop state
+        retry = False
+        pause = False
+        done = False
+        plan_idx = 0
+        pygame.display.set_caption(f'plan_idx:{plan_idx}')
+        # step-level while loop
+        while not done:
+            # process keypress events
+            for event in pygame.event.get():
+                if event.type == pygame.KEYDOWN:
+                    if event.key == pygame.K_SPACE:
+                        # hold Space to pause
+                        plan_idx += 1
+                        pygame.display.set_caption(f'plan_idx:{plan_idx}')
+                        pause = True
+                    elif event.key == pygame.K_r:
+                        # press "R" to retry
+                        retry=True
+                    elif event.key == pygame.K_q:
+                        # press "Q" to exit
+                        exit(0)
+                if event.type == pygame.KEYUP:
+                    if event.key == pygame.K_SPACE:
+                        pause = False
+
+            # handle control flow
+            if retry:
+                break
+            if pause:
+                continue
+            
+            # get action from mouse
+            # None if mouse is not close to the agent
+            act = agent.act(obs)
+            if not act is None:
+                # teleop started
+                # state dim 2+3
+                state = np.concatenate([info['pos_agent'], info['block_pose']])
+                # discard unused information such as visibility mask and agent pos
+                # for compatibility
+                keypoint = obs.reshape(2,-1)[0].reshape(-1,2)[:9]
+                data = {
+                    'img': img,
+                    'state': np.float32(state),
+                    'keypoint': np.float32(keypoint),
+                    'action': np.float32(act),
+                    'n_contacts': np.float32([info['n_contacts']])
+                }
+                episode.append(data)
+                
+            # step env and render
+            obs, reward, done, info = env.step(act)
+            img = env.render(mode='human')
+            
+            # regulate control frequency
+            clock.tick(control_hz)
+        if not retry:
+            # save episode buffer to replay buffer (on disk)
+            data_dict = dict()
+            for key in episode[0].keys():
+                data_dict[key] = np.stack(
+                    [x[key] for x in episode])
+            replay_buffer.add_episode(data_dict, compressors='disk')
+            print(f'saved seed {seed}')
+        else:
+            print(f'retry seed {seed}')
+
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/demo_real_robot.py

@@ -0,0 +1,160 @@
+"""
+Usage:
+(robodiff)$ python demo_real_robot.py -o <demo_save_dir> --robot_ip <ip_of_ur5>
+
+Robot movement:
+Move your SpaceMouse to move the robot EEF (locked in xy plane).
+Press SpaceMouse right button to unlock z axis.
+Press SpaceMouse left button to enable rotation axes.
+
+Recording control:
+Click the opencv window (make sure it's in focus).
+Press "C" to start recording.
+Press "S" to stop recording.
+Press "Q" to exit program.
+Press "Backspace" to delete the previously recorded episode.
+"""
+
+# %%
+import time
+from multiprocessing.managers import SharedMemoryManager
+import click
+import cv2
+import numpy as np
+import scipy.spatial.transform as st
+from diffusion_policy.real_world.real_env import RealEnv
+from diffusion_policy.real_world.spacemouse_shared_memory import Spacemouse
+from diffusion_policy.common.precise_sleep import precise_wait
+from diffusion_policy.real_world.keystroke_counter import (
+    KeystrokeCounter, Key, KeyCode
+)
+
+@click.command()
+@click.option('--output', '-o', required=True, help="Directory to save demonstration dataset.")
+@click.option('--robot_ip', '-ri', required=True, help="UR5's IP address e.g. 192.168.0.204")
+@click.option('--vis_camera_idx', default=0, type=int, help="Which RealSense camera to visualize.")
+@click.option('--init_joints', '-j', is_flag=True, default=False, help="Whether to initialize robot joint configuration in the beginning.")
+@click.option('--frequency', '-f', default=10, type=float, help="Control frequency in Hz.")
+@click.option('--command_latency', '-cl', default=0.01, type=float, help="Latency between receiving SapceMouse command to executing on Robot in Sec.")
+def main(output, robot_ip, vis_camera_idx, init_joints, frequency, command_latency):
+    dt = 1/frequency
+    with SharedMemoryManager() as shm_manager:
+        with KeystrokeCounter() as key_counter, \
+            Spacemouse(shm_manager=shm_manager) as sm, \
+            RealEnv(
+                output_dir=output, 
+                robot_ip=robot_ip, 
+                # recording resolution
+                obs_image_resolution=(1280,720),
+                frequency=frequency,
+                init_joints=init_joints,
+                enable_multi_cam_vis=True,
+                record_raw_video=True,
+                # number of threads per camera view for video recording (H.264)
+                thread_per_video=3,
+                # video recording quality, lower is better (but slower).
+                video_crf=21,
+                shm_manager=shm_manager
+            ) as env:
+            cv2.setNumThreads(1)
+
+            # realsense exposure
+            env.realsense.set_exposure(exposure=120, gain=0)
+            # realsense white balance
+            env.realsense.set_white_balance(white_balance=5900)
+
+            time.sleep(1.0)
+            print('Ready!')
+            state = env.get_robot_state()
+            target_pose = state['TargetTCPPose']
+            t_start = time.monotonic()
+            iter_idx = 0
+            stop = False
+            is_recording = False
+            while not stop:
+                # calculate timing
+                t_cycle_end = t_start + (iter_idx + 1) * dt
+                t_sample = t_cycle_end - command_latency
+                t_command_target = t_cycle_end + dt
+
+                # pump obs
+                obs = env.get_obs()
+
+                # handle key presses
+                press_events = key_counter.get_press_events()
+                for key_stroke in press_events:
+                    if key_stroke == KeyCode(char='q'):
+                        # Exit program
+                        stop = True
+                    elif key_stroke == KeyCode(char='c'):
+                        # Start recording
+                        env.start_episode(t_start + (iter_idx + 2) * dt - time.monotonic() + time.time())
+                        key_counter.clear()
+                        is_recording = True
+                        print('Recording!')
+                    elif key_stroke == KeyCode(char='s'):
+                        # Stop recording
+                        env.end_episode()
+                        key_counter.clear()
+                        is_recording = False
+                        print('Stopped.')
+                    elif key_stroke == Key.backspace:
+                        # Delete the most recent recorded episode
+                        if click.confirm('Are you sure to drop an episode?'):
+                            env.drop_episode()
+                            key_counter.clear()
+                            is_recording = False
+                        # delete
+                stage = key_counter[Key.space]
+
+                # visualize
+                vis_img = obs[f'camera_{vis_camera_idx}'][-1,:,:,::-1].copy()
+                episode_id = env.replay_buffer.n_episodes
+                text = f'Episode: {episode_id}, Stage: {stage}'
+                if is_recording:
+                    text += ', Recording!'
+                cv2.putText(
+                    vis_img,
+                    text,
+                    (10,30),
+                    fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                    fontScale=1,
+                    thickness=2,
+                    color=(255,255,255)
+                )
+
+                cv2.imshow('default', vis_img)
+                cv2.pollKey()
+
+                precise_wait(t_sample)
+                # get teleop command
+                sm_state = sm.get_motion_state_transformed()
+                # print(sm_state)
+                dpos = sm_state[:3] * (env.max_pos_speed / frequency)
+                drot_xyz = sm_state[3:] * (env.max_rot_speed / frequency)
+                
+                if not sm.is_button_pressed(0):
+                    # translation mode
+                    drot_xyz[:] = 0
+                else:
+                    dpos[:] = 0
+                if not sm.is_button_pressed(1):
+                    # 2D translation mode
+                    dpos[2] = 0    
+
+                drot = st.Rotation.from_euler('xyz', drot_xyz)
+                target_pose[:3] += dpos
+                target_pose[3:] = (drot * st.Rotation.from_rotvec(
+                    target_pose[3:])).as_rotvec()
+
+                # execute teleop command
+                env.exec_actions(
+                    actions=[target_pose], 
+                    timestamps=[t_command_target-time.monotonic()+time.time()],
+                    stages=[stage])
+                precise_wait(t_cycle_end)
+                iter_idx += 1
+
+# %%
+if __name__ == '__main__':
+    main()

third_party/diffusion_policy/diffusion_policy.egg-info/PKG-INFO

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: diffusion_policy
+Version: 0.0.0
+License-File: LICENSE
+Dynamic: license-file

third_party/diffusion_policy/diffusion_policy.egg-info/SOURCES.txt

@@ -0,0 +1,19 @@
+LICENSE
+README.md
+setup.py
+diffusion_policy.egg-info/PKG-INFO
+diffusion_policy.egg-info/SOURCES.txt
+diffusion_policy.egg-info/dependency_links.txt
+diffusion_policy.egg-info/top_level.txt
+tests/test_block_pushing.py
+tests/test_cv2_util.py
+tests/test_multi_realsense.py
+tests/test_pose_trajectory_interpolator.py
+tests/test_precise_sleep.py
+tests/test_replay_buffer.py
+tests/test_ring_buffer.py
+tests/test_robomimic_image_runner.py
+tests/test_robomimic_lowdim_runner.py
+tests/test_shared_queue.py
+tests/test_single_realsense.py
+tests/test_timestamp_accumulator.py

third_party/diffusion_policy/diffusion_policy.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

third_party/diffusion_policy/diffusion_policy.egg-info/top_level.txt

@@ -0,0 +1 @@
+

third_party/diffusion_policy/diffusion_policy/workspace/train_diffusion_unet_lowdim_workspace.py

@@ -0,0 +1,306 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+    os.chdir(ROOT_DIR)
+
+import os
+import hydra
+import torch
+from omegaconf import OmegaConf
+import pathlib
+from torch.utils.data import DataLoader
+import copy
+import numpy as np
+import random
+import wandb
+import tqdm
+import shutil
+
+from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+from diffusion_policy.policy.diffusion_unet_lowdim_policy import DiffusionUnetLowdimPolicy
+from diffusion_policy.dataset.base_dataset import BaseLowdimDataset
+from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner
+from diffusion_policy.common.checkpoint_util import TopKCheckpointManager
+from diffusion_policy.common.json_logger import JsonLogger
+from diffusion_policy.model.common.lr_scheduler import get_scheduler
+from diffusers.training_utils import EMAModel
+
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+# %%
+class TrainDiffusionUnetLowdimWorkspace(BaseWorkspace):
+    include_keys = ['global_step', 'epoch']
+
+    def __init__(self, cfg: OmegaConf, output_dir=None):
+        super().__init__(cfg, output_dir=output_dir)
+
+        # set seed
+        seed = cfg.training.seed
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+        # configure model
+        self.model: DiffusionUnetLowdimPolicy
+        self.model = hydra.utils.instantiate(cfg.policy)
+
+        self.ema_model: DiffusionUnetLowdimPolicy = None
+        if cfg.training.use_ema:
+            self.ema_model = copy.deepcopy(self.model)
+
+        # configure training state
+        self.optimizer = hydra.utils.instantiate(
+            cfg.optimizer, params=self.model.parameters())
+
+        self.global_step = 0
+        self.epoch = 0
+
+    def run(self):
+        cfg = copy.deepcopy(self.cfg)
+
+        # resume training
+        if cfg.training.resume:
+            lastest_ckpt_path = self.get_checkpoint_path()
+            if lastest_ckpt_path.is_file():
+                print(f"Resuming from checkpoint {lastest_ckpt_path}")
+                self.load_checkpoint(path=lastest_ckpt_path)
+
+        # configure dataset
+        dataset: BaseLowdimDataset
+        dataset = hydra.utils.instantiate(cfg.task.dataset)
+        assert isinstance(dataset, BaseLowdimDataset)
+        train_dataloader = DataLoader(dataset, **cfg.dataloader)
+        normalizer = dataset.get_normalizer()
+
+        # configure validation dataset
+        val_dataset = dataset.get_validation_dataset()
+        val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader)
+
+        self.model.set_normalizer(normalizer)
+        if cfg.training.use_ema:
+            self.ema_model.set_normalizer(normalizer)
+
+        # configure lr scheduler
+        lr_scheduler = get_scheduler(
+            cfg.training.lr_scheduler,
+            optimizer=self.optimizer,
+            num_warmup_steps=cfg.training.lr_warmup_steps,
+            num_training_steps=(
+                len(train_dataloader) * cfg.training.num_epochs) \
+                    // cfg.training.gradient_accumulate_every,
+            # pytorch assumes stepping LRScheduler every epoch
+            # however huggingface diffusers steps it every batch
+            last_epoch=self.global_step-1
+        )
+
+        # configure ema
+        ema: EMAModel = None
+        if cfg.training.use_ema:
+            ema = hydra.utils.instantiate(
+                cfg.ema,
+                model=self.ema_model)
+
+        # configure env runner
+        env_runner: BaseLowdimRunner
+        env_runner = hydra.utils.instantiate(
+            cfg.task.env_runner,
+            output_dir=self.output_dir)
+        assert isinstance(env_runner, BaseLowdimRunner)
+
+        # configure logging
+        wandb_run = wandb.init(
+            dir=str(self.output_dir),
+            config=OmegaConf.to_container(cfg, resolve=True),
+            **cfg.logging
+        )
+        wandb.config.update(
+            {
+                "output_dir": self.output_dir,
+            }
+        )
+
+        # configure checkpoint
+        topk_manager = TopKCheckpointManager(
+            save_dir=os.path.join(self.output_dir, 'checkpoints'),
+            **cfg.checkpoint.topk
+        )
+
+        # device transfer
+        device = torch.device(cfg.training.device)
+        self.model.to(device)
+        if self.ema_model is not None:
+            self.ema_model.to(device)
+        optimizer_to(self.optimizer, device)
+
+        # save batch for sampling
+        train_sampling_batch = None
+
+        if cfg.training.debug:
+            cfg.training.num_epochs = 2
+            cfg.training.max_train_steps = 3
+            cfg.training.max_val_steps = 3
+            cfg.training.rollout_every = 1
+            cfg.training.checkpoint_every = 1
+            cfg.training.val_every = 1
+            cfg.training.sample_every = 1
+
+        # training loop
+        log_path = os.path.join(self.output_dir, 'logs.json.txt')
+        with JsonLogger(log_path) as json_logger:
+            for local_epoch_idx in range(cfg.training.num_epochs):
+                step_log = dict()
+                # ========= train for this epoch ==========
+                train_losses = list()
+                with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", 
+                        leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                    for batch_idx, batch in enumerate(tepoch):
+                        # device transfer
+                        batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                        if train_sampling_batch is None:
+                            train_sampling_batch = batch
+
+                        # compute loss
+                        raw_loss = self.model.compute_loss(batch)
+                        loss = raw_loss / cfg.training.gradient_accumulate_every
+                        loss.backward()
+
+                        # step optimizer
+                        if self.global_step % cfg.training.gradient_accumulate_every == 0:
+                            self.optimizer.step()
+                            self.optimizer.zero_grad()
+                            lr_scheduler.step()
+                        
+                        # update ema
+                        if cfg.training.use_ema:
+                            ema.step(self.model)
+
+                        # logging
+                        raw_loss_cpu = raw_loss.item()
+                        tepoch.set_postfix(loss=raw_loss_cpu, refresh=False)
+                        train_losses.append(raw_loss_cpu)
+                        step_log = {
+                            'train_loss': raw_loss_cpu,
+                            'global_step': self.global_step,
+                            'epoch': self.epoch,
+                            'lr': lr_scheduler.get_last_lr()[0]
+                        }
+
+                        is_last_batch = (batch_idx == (len(train_dataloader)-1))
+                        if not is_last_batch:
+                            # log of last step is combined with validation and rollout
+                            wandb_run.log(step_log, step=self.global_step)
+                            json_logger.log(step_log)
+                            self.global_step += 1
+
+                        if (cfg.training.max_train_steps is not None) \
+                            and batch_idx >= (cfg.training.max_train_steps-1):
+                            break
+                
+                # at the end of each epoch
+                # replace train_loss with epoch average
+                train_loss = np.mean(train_losses)
+                step_log['train_loss'] = train_loss
+
+                # ========= eval for this epoch ==========
+                policy = self.model
+                if cfg.training.use_ema:
+                    policy = self.ema_model
+                policy.eval()
+
+                # run rollout
+                if (self.epoch % cfg.training.rollout_every) == 0:
+                    runner_log = env_runner.run(policy)
+                    # log all
+                    step_log.update(runner_log)
+
+                # run validation
+                if (self.epoch % cfg.training.val_every) == 0:
+                    with torch.no_grad():
+                        val_losses = list()
+                        with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", 
+                                leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                            for batch_idx, batch in enumerate(tepoch):
+                                batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                                loss = self.model.compute_loss(batch)
+                                val_losses.append(loss)
+                                if (cfg.training.max_val_steps is not None) \
+                                    and batch_idx >= (cfg.training.max_val_steps-1):
+                                    break
+                        if len(val_losses) > 0:
+                            val_loss = torch.mean(torch.tensor(val_losses)).item()
+                            # log epoch average validation loss
+                            step_log['val_loss'] = val_loss
+
+                # run diffusion sampling on a training batch
+                if (self.epoch % cfg.training.sample_every) == 0:
+                    with torch.no_grad():
+                        # sample trajectory from training set, and evaluate difference
+                        batch = train_sampling_batch
+                        obs_dict = {'obs': batch['obs']}
+                        gt_action = batch['action']
+                        
+                        result = policy.predict_action(obs_dict)
+                        if cfg.pred_action_steps_only:
+                            pred_action = result['action']
+                            start = cfg.n_obs_steps - 1
+                            end = start + cfg.n_action_steps
+                            gt_action = gt_action[:,start:end]
+                        else:
+                            pred_action = result['action_pred']
+                        mse = torch.nn.functional.mse_loss(pred_action, gt_action)
+                        # log
+                        step_log['train_action_mse_error'] = mse.item()
+                        # release RAM
+                        del batch
+                        del obs_dict
+                        del gt_action
+                        del result
+                        del pred_action
+                        del mse
+                
+                # checkpoint
+                if (self.epoch % cfg.training.checkpoint_every) == 0:
+                    # checkpointing
+                    if cfg.checkpoint.save_last_ckpt:
+                        self.save_checkpoint()
+                    if cfg.checkpoint.save_last_snapshot:
+                        self.save_snapshot()
+
+                    # sanitize metric names
+                    metric_dict = dict()
+                    for key, value in step_log.items():
+                        new_key = key.replace('/', '_')
+                        metric_dict[new_key] = value
+                    
+                    # We can't copy the last checkpoint here
+                    # since save_checkpoint uses threads.
+                    # therefore at this point the file might have been empty!
+                    topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict)
+
+                    if topk_ckpt_path is not None:
+                        self.save_checkpoint(path=topk_ckpt_path)
+                # ========= eval end for this epoch ==========
+                policy.train()
+
+                # end of epoch
+                # log of last step is combined with validation and rollout
+                wandb_run.log(step_log, step=self.global_step)
+                json_logger.log(step_log)
+                self.global_step += 1
+                self.epoch += 1
+
+@hydra.main(
+    version_base=None,
+    config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), 
+    config_name=pathlib.Path(__file__).stem)
+def main(cfg):
+    workspace = TrainDiffusionUnetLowdimWorkspace(cfg)
+    workspace.run()
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_hybrid_workspace.py

@@ -0,0 +1,283 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+    os.chdir(ROOT_DIR)
+
+import os
+import hydra
+import torch
+from omegaconf import OmegaConf
+import pathlib
+from torch.utils.data import DataLoader
+import copy
+import random
+import wandb
+import tqdm
+import numpy as np
+import shutil
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+from diffusion_policy.policy.ibc_dfo_hybrid_image_policy import IbcDfoHybridImagePolicy
+from diffusion_policy.dataset.base_dataset import BaseImageDataset
+from diffusion_policy.env_runner.base_image_runner import BaseImageRunner
+from diffusion_policy.common.checkpoint_util import TopKCheckpointManager
+from diffusion_policy.common.json_logger import JsonLogger
+from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to
+from diffusion_policy.model.diffusion.ema_model import EMAModel
+from diffusion_policy.model.common.lr_scheduler import get_scheduler
+
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+class TrainIbcDfoHybridWorkspace(BaseWorkspace):
+    include_keys = ['global_step', 'epoch']
+
+    def __init__(self, cfg: OmegaConf, output_dir=None):
+        super().__init__(cfg, output_dir=output_dir)
+
+        # set seed
+        seed = cfg.training.seed
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+        # configure model
+        self.model: IbcDfoHybridImagePolicy= hydra.utils.instantiate(cfg.policy)
+
+        # configure training state
+        self.optimizer = hydra.utils.instantiate(
+            cfg.optimizer, params=self.model.parameters())
+
+        # configure training state
+        self.global_step = 0
+        self.epoch = 0
+
+    def run(self):
+        cfg = copy.deepcopy(self.cfg)
+
+        # resume training
+        if cfg.training.resume:
+            lastest_ckpt_path = self.get_checkpoint_path()
+            if lastest_ckpt_path.is_file():
+                print(f"Resuming from checkpoint {lastest_ckpt_path}")
+                self.load_checkpoint(path=lastest_ckpt_path)
+
+        # configure dataset
+        dataset: BaseImageDataset
+        dataset = hydra.utils.instantiate(cfg.task.dataset)
+        assert isinstance(dataset, BaseImageDataset)
+        train_dataloader = DataLoader(dataset, **cfg.dataloader)
+        normalizer = dataset.get_normalizer()
+
+        # configure validation dataset
+        val_dataset = dataset.get_validation_dataset()
+        val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader)
+
+        self.model.set_normalizer(normalizer)
+
+        # configure lr scheduler
+        lr_scheduler = get_scheduler(
+            cfg.training.lr_scheduler,
+            optimizer=self.optimizer,
+            num_warmup_steps=cfg.training.lr_warmup_steps,
+            num_training_steps=(
+                len(train_dataloader) * cfg.training.num_epochs) \
+                    // cfg.training.gradient_accumulate_every,
+            # pytorch assumes stepping LRScheduler every epoch
+            # however huggingface diffusers steps it every batch
+            last_epoch=self.global_step-1
+        )
+
+        # configure env
+        env_runner: BaseImageRunner
+        env_runner = hydra.utils.instantiate(
+            cfg.task.env_runner,
+            output_dir=self.output_dir)
+        assert isinstance(env_runner, BaseImageRunner)
+
+        # configure logging
+        wandb_run = wandb.init(
+            dir=str(self.output_dir),
+            config=OmegaConf.to_container(cfg, resolve=True),
+            **cfg.logging
+        )
+        wandb.config.update(
+            {
+                "output_dir": self.output_dir,
+            }
+        )
+
+        # configure checkpoint
+        topk_manager = TopKCheckpointManager(
+            save_dir=os.path.join(self.output_dir, 'checkpoints'),
+            **cfg.checkpoint.topk
+        )
+
+        # device transfer
+        device = torch.device(cfg.training.device)
+        self.model.to(device)
+        optimizer_to(self.optimizer, device)
+
+        # save batch for sampling
+        train_sampling_batch = None
+
+        if cfg.training.debug:
+            cfg.training.num_epochs = 2
+            cfg.training.max_train_steps = 3
+            cfg.training.max_val_steps = 3
+            cfg.training.rollout_every = 1
+            cfg.training.checkpoint_every = 1
+            cfg.training.val_every = 1
+            cfg.training.sample_every = 1
+
+        # training loop
+        log_path = os.path.join(self.output_dir, 'logs.json.txt')
+        with JsonLogger(log_path) as json_logger:
+            for local_epoch_idx in range(cfg.training.num_epochs):
+                step_log = dict()
+                # ========= train for this epoch ==========
+                train_losses = list()
+                with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", 
+                        leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                    for batch_idx, batch in enumerate(tepoch):
+                        # device transfer
+                        batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                        if train_sampling_batch is None:
+                            train_sampling_batch = batch
+
+                        # compute loss
+                        raw_loss = self.model.compute_loss(batch)
+                        loss = raw_loss / cfg.training.gradient_accumulate_every
+                        loss.backward()
+
+                        # step optimizer
+                        if self.global_step % cfg.training.gradient_accumulate_every == 0:
+                            self.optimizer.step()
+                            self.optimizer.zero_grad()
+                            lr_scheduler.step()
+
+                        # logging
+                        raw_loss_cpu = raw_loss.item()
+                        tepoch.set_postfix(loss=raw_loss_cpu, refresh=False)
+                        train_losses.append(raw_loss_cpu)
+                        step_log = {
+                            'train_loss': raw_loss_cpu,
+                            'global_step': self.global_step,
+                            'epoch': self.epoch,
+                            'lr': lr_scheduler.get_last_lr()[0]
+                        }
+
+                        is_last_batch = (batch_idx == (len(train_dataloader)-1))
+                        if not is_last_batch:
+                            # log of last step is combined with validation and rollout
+                            wandb_run.log(step_log, step=self.global_step)
+                            json_logger.log(step_log)
+                            self.global_step += 1
+
+                        if (cfg.training.max_train_steps is not None) \
+                            and batch_idx >= (cfg.training.max_train_steps-1):
+                            break
+
+                # at the end of each epoch
+                # replace train_loss with epoch average
+                train_loss = np.mean(train_losses)
+                step_log['train_loss'] = train_loss
+
+                # ========= eval for this epoch ==========
+                policy = self.model
+                policy.eval()
+
+                # run rollout
+                if (self.epoch % cfg.training.rollout_every) == 0:
+                    runner_log = env_runner.run(policy)
+                    # log all
+                    step_log.update(runner_log)
+
+                # run validation
+                if (self.epoch % cfg.training.val_every) == 0:
+                    with torch.no_grad():
+                        val_losses = list()
+                        with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", 
+                                leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                            for batch_idx, batch in enumerate(tepoch):
+                                batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                                loss = self.model.compute_loss(batch)
+                                val_losses.append(loss)
+                                if (cfg.training.max_val_steps is not None) \
+                                    and batch_idx >= (cfg.training.max_val_steps-1):
+                                    break
+                        if len(val_losses) > 0:
+                            val_loss = torch.mean(torch.tensor(val_losses)).item()
+                            # log epoch average validation loss
+                            step_log['val_loss'] = val_loss
+
+                # run diffusion sampling on a training batch
+                if (self.epoch % cfg.training.sample_every) == 0:
+                    with torch.no_grad():
+                        # sample trajectory from training set, and evaluate difference
+                        batch = train_sampling_batch
+                        n_samples = cfg.training.sample_max_batch
+                        batch = dict_apply(train_sampling_batch, 
+                            lambda x: x.to(device, non_blocking=True))
+                        obs_dict = dict_apply(batch['obs'], lambda x: x[:n_samples])
+                        gt_action = batch['action']
+                        
+                        result = policy.predict_action(obs_dict)
+                        pred_action = result['action']
+                        start = cfg.n_obs_steps - 1
+                        end = start + cfg.n_action_steps
+                        gt_action = gt_action[:,start:end]
+                        mse = torch.nn.functional.mse_loss(pred_action, gt_action)
+                        # log
+                        step_log['train_action_mse_error'] = mse.item()
+                        # release RAM
+                        del batch
+                        del obs_dict
+                        del gt_action
+                        del result
+                        del pred_action
+                        del mse
+                
+                # checkpoint
+                if (self.epoch % cfg.training.checkpoint_every) == 0:
+                    # checkpointing
+                    if cfg.checkpoint.save_last_ckpt:
+                        self.save_checkpoint()
+                    if cfg.checkpoint.save_last_snapshot:
+                        self.save_snapshot()
+
+                    # sanitize metric names
+                    metric_dict = dict()
+                    for key, value in step_log.items():
+                        new_key = key.replace('/', '_')
+                        metric_dict[new_key] = value
+                    
+                    # We can't copy the last checkpoint here
+                    # since save_checkpoint uses threads.
+                    # therefore at this point the file might have been empty!
+                    topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict)
+
+                    if topk_ckpt_path is not None:
+                        self.save_checkpoint(path=topk_ckpt_path)
+                # ========= eval end for this epoch ==========
+                policy.train()
+
+                # end of epoch
+                # log of last step is combined with validation and rollout
+                wandb_run.log(step_log, step=self.global_step)
+                json_logger.log(step_log)
+                self.global_step += 1
+                self.epoch += 1
+
+@hydra.main(
+    version_base=None,
+    config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), 
+    config_name=pathlib.Path(__file__).stem)
+def main(cfg):
+    workspace = TrainIbcDfoHybridWorkspace(cfg)
+    workspace.run()
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/diffusion_policy/workspace/train_ibc_dfo_lowdim_workspace.py

@@ -0,0 +1,282 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+    os.chdir(ROOT_DIR)
+
+import os
+import hydra
+import torch
+from omegaconf import OmegaConf
+import pathlib
+from torch.utils.data import DataLoader
+import copy
+import numpy as np
+import random
+import wandb
+import tqdm
+import shutil
+
+from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+from diffusion_policy.policy.ibc_dfo_lowdim_policy import IbcDfoLowdimPolicy
+from diffusion_policy.dataset.base_dataset import BaseLowdimDataset
+from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner
+from diffusion_policy.common.checkpoint_util import TopKCheckpointManager
+from diffusion_policy.common.json_logger import JsonLogger
+from diffusion_policy.model.common.lr_scheduler import get_scheduler
+
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+# %%
+class TrainIbcDfoLowdimWorkspace(BaseWorkspace):
+    include_keys = ['global_step', 'epoch']
+
+    def __init__(self, cfg: OmegaConf, output_dir=None):
+        super().__init__(cfg, output_dir=output_dir)
+
+        # set seed
+        seed = cfg.training.seed
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+        # configure model
+        self.model: IbcDfoLowdimPolicy
+        self.model = hydra.utils.instantiate(cfg.policy)
+
+        # configure training state
+        self.optimizer = hydra.utils.instantiate(
+            cfg.optimizer, params=self.model.parameters())
+
+        self.global_step = 0
+        self.epoch = 0
+
+    def run(self):
+        cfg = copy.deepcopy(self.cfg)
+
+        # resume training
+        if cfg.training.resume:
+            lastest_ckpt_path = self.get_checkpoint_path()
+            if lastest_ckpt_path.is_file():
+                print(f"Resuming from checkpoint {lastest_ckpt_path}")
+                self.load_checkpoint(path=lastest_ckpt_path)
+
+        # configure dataset
+        dataset: BaseLowdimDataset
+        dataset = hydra.utils.instantiate(cfg.task.dataset)
+        assert isinstance(dataset, BaseLowdimDataset)
+        train_dataloader = DataLoader(dataset, **cfg.dataloader)
+        normalizer = dataset.get_normalizer()
+
+        # configure validation dataset
+        val_dataset = dataset.get_validation_dataset()
+        val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader)
+
+        self.model.set_normalizer(normalizer)
+
+        # configure lr scheduler
+        lr_scheduler = get_scheduler(
+            cfg.training.lr_scheduler,
+            optimizer=self.optimizer,
+            num_warmup_steps=cfg.training.lr_warmup_steps,
+            num_training_steps=(
+                len(train_dataloader) * cfg.training.num_epochs) \
+                    // cfg.training.gradient_accumulate_every,
+            # pytorch assumes stepping LRScheduler every epoch
+            # however huggingface diffusers steps it every batch
+            last_epoch=self.global_step-1
+        )
+
+        # configure env runner
+        env_runner: BaseLowdimRunner
+        env_runner = hydra.utils.instantiate(
+            cfg.task.env_runner,
+            output_dir=self.output_dir)
+        assert isinstance(env_runner, BaseLowdimRunner)
+
+        # configure logging
+        wandb_run = wandb.init(
+            dir=str(self.output_dir),
+            config=OmegaConf.to_container(cfg, resolve=True),
+            **cfg.logging
+        )
+        wandb.config.update(
+            {
+                "output_dir": self.output_dir,
+            }
+        )
+
+        # configure checkpoint
+        topk_manager = TopKCheckpointManager(
+            save_dir=os.path.join(self.output_dir, 'checkpoints'),
+            **cfg.checkpoint.topk
+        )
+
+        # device transfer
+        device = torch.device(cfg.training.device)
+        self.model.to(device)
+        optimizer_to(self.optimizer, device)
+
+        # save batch for sampling
+        train_sampling_batch = None
+
+        if cfg.training.debug:
+            cfg.training.num_epochs = 2
+            cfg.training.max_train_steps = 3
+            cfg.training.max_val_steps = 3
+            cfg.training.rollout_every = 1
+            cfg.training.checkpoint_every = 1
+            cfg.training.val_every = 1
+            cfg.training.sample_every = 1
+
+        # training loop
+        log_path = os.path.join(self.output_dir, 'logs.json.txt')
+        with JsonLogger(log_path) as json_logger:
+            for local_epoch_idx in range(cfg.training.num_epochs):
+                step_log = dict()
+                # ========= train for this epoch ==========
+                train_losses = list()
+                with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", 
+                        leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                    for batch_idx, batch in enumerate(tepoch):
+                        # device transfer
+                        batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                        if train_sampling_batch is None:
+                            train_sampling_batch = batch
+
+                        # compute loss
+                        raw_loss = self.model.compute_loss(batch)
+                        loss = raw_loss / cfg.training.gradient_accumulate_every
+                        loss.backward()
+
+                        # step optimizer
+                        if self.global_step % cfg.training.gradient_accumulate_every == 0:
+                            self.optimizer.step()
+                            self.optimizer.zero_grad()
+                            lr_scheduler.step()
+
+                        # logging
+                        raw_loss_cpu = raw_loss.item()
+                        tepoch.set_postfix(loss=raw_loss_cpu, refresh=False)
+                        train_losses.append(raw_loss_cpu)
+                        step_log = {
+                            'train_loss': raw_loss_cpu,
+                            'global_step': self.global_step,
+                            'epoch': self.epoch,
+                            'lr': lr_scheduler.get_last_lr()[0]
+                        }
+
+                        is_last_batch = (batch_idx == (len(train_dataloader)-1))
+                        if not is_last_batch:
+                            # log of last step is combined with validation and rollout
+                            wandb_run.log(step_log, step=self.global_step)
+                            json_logger.log(step_log)
+                            self.global_step += 1
+
+                        if (cfg.training.max_train_steps is not None) \
+                            and batch_idx >= (cfg.training.max_train_steps-1):
+                            break
+                
+                # at the end of each epoch
+                # replace train_loss with epoch average
+                train_loss = np.mean(train_losses)
+                step_log['train_loss'] = train_loss
+
+                # ========= eval for this epoch ==========
+                policy = self.model
+                policy.eval()
+
+                # run rollout
+                if (self.epoch % cfg.training.rollout_every) == 0:
+                    runner_log = env_runner.run(policy)
+                    # log all
+                    step_log.update(runner_log)
+
+                # run validation
+                if (self.epoch % cfg.training.val_every) == 0:
+                    with torch.no_grad():
+                        val_losses = list()
+                        with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", 
+                                leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                            for batch_idx, batch in enumerate(tepoch):
+                                batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                                loss = self.model.compute_loss(batch)
+                                val_losses.append(loss)
+                                if (cfg.training.max_val_steps is not None) \
+                                    and batch_idx >= (cfg.training.max_val_steps-1):
+                                    break
+                        if len(val_losses) > 0:
+                            val_loss = torch.mean(torch.tensor(val_losses)).item()
+                            # log epoch average validation loss
+                            step_log['val_loss'] = val_loss
+
+                # run diffusion sampling on a training batch
+                if (self.epoch % cfg.training.sample_every) == 0:
+                    with torch.no_grad():
+                        # sample trajectory from training set, and evaluate difference
+                        batch = train_sampling_batch
+                        n_samples = cfg.training.sample_max_batch
+                        obs_dict = {'obs': batch['obs'][:n_samples]}
+                        gt_action = batch['action'][:n_samples]
+
+                        result = policy.predict_action(obs_dict)
+                        pred_action = result['action']
+                        start = cfg.n_obs_steps - 1
+                        end = start + cfg.n_action_steps
+                        gt_action = gt_action[:,start:end]
+                        mse = torch.nn.functional.mse_loss(pred_action, gt_action)
+                        # log
+                        step_log['train_action_mse_error'] = mse.item()
+                        # release RAM
+                        del batch
+                        del obs_dict
+                        del gt_action
+                        del result
+                        del pred_action
+                        del mse
+                
+                # checkpoint
+                if (self.epoch % cfg.training.checkpoint_every) == 0:
+                    # checkpointing
+                    if cfg.checkpoint.save_last_ckpt:
+                        self.save_checkpoint()
+                    if cfg.checkpoint.save_last_snapshot:
+                        self.save_snapshot()
+
+                    # sanitize metric names
+                    metric_dict = dict()
+                    for key, value in step_log.items():
+                        new_key = key.replace('/', '_')
+                        metric_dict[new_key] = value
+                    
+                    # We can't copy the last checkpoint here
+                    # since save_checkpoint uses threads.
+                    # therefore at this point the file might have been empty!
+                    topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict)
+
+                    if topk_ckpt_path is not None:
+                        self.save_checkpoint(path=topk_ckpt_path)
+                # ========= eval end for this epoch ==========
+                policy.train()
+
+                # end of epoch
+                # log of last step is combined with validation and rollout
+                wandb_run.log(step_log, step=self.global_step)
+                json_logger.log(step_log)
+                self.global_step += 1
+                self.epoch += 1
+
+@hydra.main(
+    version_base=None,
+    config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), 
+    config_name=pathlib.Path(__file__).stem)
+def main(cfg):
+    workspace = TrainIbcDfoLowdimWorkspace(cfg)
+    workspace.run()
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/diffusion_policy/workspace/train_robomimic_image_workspace.py

@@ -0,0 +1,254 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+    os.chdir(ROOT_DIR)
+
+import os
+import hydra
+import torch
+from omegaconf import OmegaConf
+import pathlib
+from torch.utils.data import DataLoader
+import copy
+import random
+import wandb
+import tqdm
+import numpy as np
+import shutil
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+from diffusion_policy.policy.robomimic_image_policy import RobomimicImagePolicy
+from diffusion_policy.dataset.base_dataset import BaseImageDataset
+from diffusion_policy.env_runner.base_image_runner import BaseImageRunner
+from diffusion_policy.common.checkpoint_util import TopKCheckpointManager
+from diffusion_policy.common.json_logger import JsonLogger
+from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to
+
+
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+class TrainRobomimicImageWorkspace(BaseWorkspace):
+    include_keys = ['global_step', 'epoch']
+
+    def __init__(self, cfg: OmegaConf, output_dir=None):
+        super().__init__(cfg, output_dir=output_dir)
+
+        # set seed
+        seed = cfg.training.seed
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+        # configure model
+        self.model: RobomimicImagePolicy = hydra.utils.instantiate(cfg.policy)
+
+        # configure training state
+        self.global_step = 0
+        self.epoch = 0
+
+    def run(self):
+        cfg = copy.deepcopy(self.cfg)
+
+        # resume training
+        if cfg.training.resume:
+            lastest_ckpt_path = self.get_checkpoint_path()
+            if lastest_ckpt_path.is_file():
+                print(f"Resuming from checkpoint {lastest_ckpt_path}")
+                self.load_checkpoint(path=lastest_ckpt_path)
+
+        # configure dataset
+        dataset: BaseImageDataset
+        dataset = hydra.utils.instantiate(cfg.task.dataset)
+        assert isinstance(dataset, BaseImageDataset)
+        train_dataloader = DataLoader(dataset, **cfg.dataloader)
+        normalizer = dataset.get_normalizer()
+
+        # configure validation dataset
+        val_dataset = dataset.get_validation_dataset()
+        val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader)
+
+        self.model.set_normalizer(normalizer)
+
+        # configure env
+        env_runner: BaseImageRunner
+        env_runner = hydra.utils.instantiate(
+            cfg.task.env_runner,
+            output_dir=self.output_dir)
+        assert isinstance(env_runner, BaseImageRunner)
+
+        # configure logging
+        wandb_run = wandb.init(
+            dir=str(self.output_dir),
+            config=OmegaConf.to_container(cfg, resolve=True),
+            **cfg.logging
+        )
+        wandb.config.update(
+            {
+                "output_dir": self.output_dir,
+            }
+        )
+
+        # configure checkpoint
+        topk_manager = TopKCheckpointManager(
+            save_dir=os.path.join(self.output_dir, 'checkpoints'),
+            **cfg.checkpoint.topk
+        )
+
+        # device transfer
+        device = torch.device(cfg.training.device)
+        self.model.to(device)
+
+        # save batch for sampling
+        train_sampling_batch = None
+
+        if cfg.training.debug:
+            cfg.training.num_epochs = 2
+            cfg.training.max_train_steps = 3
+            cfg.training.max_val_steps = 3
+            cfg.training.rollout_every = 1
+            cfg.training.checkpoint_every = 1
+            cfg.training.val_every = 1
+            cfg.training.sample_every = 1
+
+        # training loop
+        log_path = os.path.join(self.output_dir, 'logs.json.txt')
+        with JsonLogger(log_path) as json_logger:
+            for local_epoch_idx in range(cfg.training.num_epochs):
+                step_log = dict()
+                # ========= train for this epoch ==========
+                train_losses = list()
+                with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", 
+                        leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                    for batch_idx, batch in enumerate(tepoch):
+                        # device transfer
+                        batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                        if train_sampling_batch is None:
+                            train_sampling_batch = batch
+
+                        info = self.model.train_on_batch(batch, epoch=self.epoch)
+
+                        # logging 
+                        loss_cpu = info['losses']['action_loss'].item()
+                        tepoch.set_postfix(loss=loss_cpu, refresh=False)
+                        train_losses.append(loss_cpu)
+                        step_log = {
+                            'train_loss': loss_cpu,
+                            'global_step': self.global_step,
+                            'epoch': self.epoch
+                        }
+
+                        is_last_batch = (batch_idx == (len(train_dataloader)-1))
+                        if not is_last_batch:
+                            # log of last step is combined with validation and rollout
+                            wandb_run.log(step_log, step=self.global_step)
+                            json_logger.log(step_log)
+                            self.global_step += 1
+
+                        if (cfg.training.max_train_steps is not None) \
+                            and batch_idx >= (cfg.training.max_train_steps-1):
+                            break
+
+                # at the end of each epoch
+                # replace train_loss with epoch average
+                train_loss = np.mean(train_losses)
+                step_log['train_loss'] = train_loss
+
+                # ========= eval for this epoch ==========
+                self.model.eval()
+
+                # run rollout
+                if (self.epoch % cfg.training.rollout_every) == 0:
+                    runner_log = env_runner.run(self.model)
+                    # log all
+                    step_log.update(runner_log)
+
+                # run validation
+                if (self.epoch % cfg.training.val_every) == 0:
+                    with torch.no_grad():
+                        val_losses = list()
+                        with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", 
+                                leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                            for batch_idx, batch in enumerate(tepoch):
+                                batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                                info = self.model.train_on_batch(batch, epoch=self.epoch, validate=True)
+                                loss = info['losses']['action_loss']
+                                val_losses.append(loss)
+                                if (cfg.training.max_val_steps is not None) \
+                                    and batch_idx >= (cfg.training.max_val_steps-1):
+                                    break
+                        if len(val_losses) > 0:
+                            val_loss = torch.mean(torch.tensor(val_losses)).item()
+                            # log epoch average validation loss
+                            step_log['val_loss'] = val_loss
+
+                # run diffusion sampling on a training batch
+                if (self.epoch % cfg.training.sample_every) == 0:
+                    with torch.no_grad():
+                        # sample trajectory from training set, and evaluate difference
+                        batch = dict_apply(train_sampling_batch, lambda x: x.to(device, non_blocking=True))
+                        obs_dict = batch['obs']
+                        gt_action = batch['action']
+                        T = gt_action.shape[1]
+
+                        pred_actions = list()
+                        self.model.reset()
+                        for i in range(T):
+                            result = self.model.predict_action(
+                                dict_apply(obs_dict, lambda x: x[:,[i]])
+                            )
+                            pred_actions.append(result['action'])
+                        pred_actions = torch.cat(pred_actions, dim=1)
+                        mse = torch.nn.functional.mse_loss(pred_actions, gt_action)
+                        step_log['train_action_mse_error'] = mse.item()
+                        del batch
+                        del obs_dict
+                        del gt_action
+                        del result
+                        del pred_actions
+                        del mse
+
+                # checkpoint
+                if (self.epoch % cfg.training.checkpoint_every) == 0:
+                    # checkpointing
+                    if cfg.checkpoint.save_last_ckpt:
+                        self.save_checkpoint()
+                    if cfg.checkpoint.save_last_snapshot:
+                        self.save_snapshot()
+
+                    # sanitize metric names
+                    metric_dict = dict()
+                    for key, value in step_log.items():
+                        new_key = key.replace('/', '_')
+                        metric_dict[new_key] = value
+                    
+                    # We can't copy the last checkpoint here
+                    # since save_checkpoint uses threads.
+                    # therefore at this point the file might have been empty!
+                    topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict)
+
+                    if topk_ckpt_path is not None:
+                        self.save_checkpoint(path=topk_ckpt_path)
+                # ========= eval end for this epoch ==========
+                self.model.train()
+
+                # end of epoch
+                # log of last step is combined with validation and rollout
+                wandb_run.log(step_log, step=self.global_step)
+                json_logger.log(step_log)
+                self.global_step += 1
+                self.epoch += 1
+
+
+@hydra.main(
+    version_base=None,
+    config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), 
+    config_name=pathlib.Path(__file__).stem)
+def main(cfg):
+    workspace = TrainRobomimicImageWorkspace(cfg)
+    workspace.run()
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/diffusion_policy/workspace/train_robomimic_lowdim_workspace.py

@@ -0,0 +1,221 @@
+if __name__ == "__main__":
+    import sys
+    import os
+    import pathlib
+
+    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
+    sys.path.append(ROOT_DIR)
+    os.chdir(ROOT_DIR)
+
+import os
+import hydra
+import torch
+from omegaconf import OmegaConf
+import pathlib
+from torch.utils.data import DataLoader
+import copy
+import random
+import wandb
+import tqdm
+import numpy as np
+import shutil
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+from diffusion_policy.policy.robomimic_lowdim_policy import RobomimicLowdimPolicy
+from diffusion_policy.dataset.base_dataset import BaseLowdimDataset
+from diffusion_policy.env_runner.base_lowdim_runner import BaseLowdimRunner
+from diffusion_policy.common.checkpoint_util import TopKCheckpointManager
+from diffusion_policy.common.json_logger import JsonLogger
+from diffusion_policy.common.pytorch_util import dict_apply, optimizer_to
+
+
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+class TrainRobomimicLowdimWorkspace(BaseWorkspace):
+    include_keys = ['global_step', 'epoch']
+
+    def __init__(self, cfg: OmegaConf):
+        super().__init__(cfg)
+
+        # set seed
+        seed = cfg.training.seed
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+        # configure model
+        self.model: RobomimicLowdimPolicy = hydra.utils.instantiate(cfg.policy)
+
+        # configure training state
+        self.global_step = 0
+        self.epoch = 0
+
+    def run(self):
+        cfg = copy.deepcopy(self.cfg)
+
+        # resume training
+        if cfg.training.resume:
+            lastest_ckpt_path = self.get_checkpoint_path()
+            if lastest_ckpt_path.is_file():
+                print(f"Resuming from checkpoint {lastest_ckpt_path}")
+                self.load_checkpoint(path=lastest_ckpt_path)
+
+        # configure dataset
+        dataset: BaseLowdimDataset
+        dataset = hydra.utils.instantiate(cfg.task.dataset)
+        assert isinstance(dataset, BaseLowdimDataset)
+        train_dataloader = DataLoader(dataset, **cfg.dataloader)
+        normalizer = dataset.get_normalizer()
+
+        # configure validation dataset
+        val_dataset = dataset.get_validation_dataset()
+        val_dataloader = DataLoader(val_dataset, **cfg.val_dataloader)
+
+        self.model.set_normalizer(normalizer)
+
+        # configure env
+        env_runner: BaseLowdimRunner
+        env_runner = hydra.utils.instantiate(
+            cfg.task.env_runner,
+            output_dir=self.output_dir)
+        assert isinstance(env_runner, BaseLowdimRunner)
+
+        # configure logging
+        wandb_run = wandb.init(
+            dir=str(self.output_dir),
+            config=OmegaConf.to_container(cfg, resolve=True),
+            **cfg.logging
+        )
+        wandb.config.update(
+            {
+                "output_dir": self.output_dir,
+            }
+        )
+
+        # configure checkpoint
+        topk_manager = TopKCheckpointManager(
+            save_dir=os.path.join(self.output_dir, 'checkpoints'),
+            **cfg.checkpoint.topk
+        )
+
+        # device transfer
+        device = torch.device(cfg.training.device)
+        self.model.to(device)
+
+        if cfg.training.debug:
+            cfg.training.num_epochs = 2
+            cfg.training.max_train_steps = 3
+            cfg.training.max_val_steps = 3
+            cfg.training.rollout_every = 1
+            cfg.training.checkpoint_every = 1
+            cfg.training.val_every = 1
+
+        # training loop
+        log_path = os.path.join(self.output_dir, 'logs.json.txt')
+        with JsonLogger(log_path) as json_logger:
+            for local_epoch_idx in range(cfg.training.num_epochs):
+                step_log = dict()
+                # ========= train for this epoch ==========
+                train_losses = list()
+                with tqdm.tqdm(train_dataloader, desc=f"Training epoch {self.epoch}", 
+                        leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                    for batch_idx, batch in enumerate(tepoch):
+                        # device transfer
+                        batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                        info = self.model.train_on_batch(batch, epoch=self.epoch)
+
+                        # logging 
+                        loss_cpu = info['losses']['action_loss'].item()
+                        tepoch.set_postfix(loss=loss_cpu, refresh=False)
+                        train_losses.append(loss_cpu)
+                        step_log = {
+                            'train_loss': loss_cpu,
+                            'global_step': self.global_step,
+                            'epoch': self.epoch
+                        }
+
+                        is_last_batch = (batch_idx == (len(train_dataloader)-1))
+                        if not is_last_batch:
+                            # log of last step is combined with validation and rollout
+                            wandb_run.log(step_log, step=self.global_step)
+                            json_logger.log(step_log)
+                            self.global_step += 1
+
+                        if (cfg.training.max_train_steps is not None) \
+                            and batch_idx >= (cfg.training.max_train_steps-1):
+                            break
+
+                # at the end of each epoch
+                # replace train_loss with epoch average
+                train_loss = np.mean(train_losses)
+                step_log['train_loss'] = train_loss
+
+                # ========= eval for this epoch ==========
+                self.model.eval()
+
+                # run rollout
+                if (self.epoch % cfg.training.rollout_every) == 0:
+                    runner_log = env_runner.run(self.model)
+                    # log all
+                    step_log.update(runner_log)
+
+                # run validation
+                if (self.epoch % cfg.training.val_every) == 0:
+                    with torch.no_grad():
+                        val_losses = list()
+                        with tqdm.tqdm(val_dataloader, desc=f"Validation epoch {self.epoch}", 
+                                leave=False, mininterval=cfg.training.tqdm_interval_sec) as tepoch:
+                            for batch_idx, batch in enumerate(tepoch):
+                                batch = dict_apply(batch, lambda x: x.to(device, non_blocking=True))
+                                info = self.model.train_on_batch(batch, epoch=self.epoch, validate=True)
+                                loss = info['losses']['action_loss']
+                                val_losses.append(loss)
+                                if (cfg.training.max_val_steps is not None) \
+                                    and batch_idx >= (cfg.training.max_val_steps-1):
+                                    break
+                        if len(val_losses) > 0:
+                            val_loss = torch.mean(torch.tensor(val_losses)).item()
+                            # log epoch average validation loss
+                            step_log['val_loss'] = val_loss
+
+                # checkpoint
+                if (self.epoch % cfg.training.checkpoint_every) == 0:
+                    # checkpointing
+                    if cfg.checkpoint.save_last_ckpt:
+                        self.save_checkpoint()
+                    if cfg.checkpoint.save_last_snapshot:
+                        self.save_snapshot()
+
+                    # sanitize metric names
+                    metric_dict = dict()
+                    for key, value in step_log.items():
+                        new_key = key.replace('/', '_')
+                        metric_dict[new_key] = value
+                    
+                    # We can't copy the last checkpoint here
+                    # since save_checkpoint uses threads.
+                    # therefore at this point the file might have been empty!
+                    topk_ckpt_path = topk_manager.get_ckpt_path(metric_dict)
+
+                    if topk_ckpt_path is not None:
+                        self.save_checkpoint(path=topk_ckpt_path)
+                # ========= eval end for this epoch ==========
+                self.model.train()
+
+                # end of epoch
+                # log of last step is combined with validation and rollout
+                wandb_run.log(step_log, step=self.global_step)
+                json_logger.log(step_log)
+                self.global_step += 1
+                self.epoch += 1
+
+
+@hydra.main(
+    version_base=None,
+    config_path=str(pathlib.Path(__file__).parent.parent.joinpath("config")), 
+    config_name=pathlib.Path(__file__).stem)
+def main(cfg):
+    workspace = TrainRobomimicLowdimWorkspace(cfg)
+    workspace.run()
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/eval.py

@@ -0,0 +1,64 @@
+"""
+Usage:
+python eval.py --checkpoint data/image/pusht/diffusion_policy_cnn/train_0/checkpoints/latest.ckpt -o data/pusht_eval_output
+"""
+
+import sys
+# use line-buffering for both stdout and stderr
+sys.stdout = open(sys.stdout.fileno(), mode='w', buffering=1)
+sys.stderr = open(sys.stderr.fileno(), mode='w', buffering=1)
+
+import os
+import pathlib
+import click
+import hydra
+import torch
+import dill
+import wandb
+import json
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+
+@click.command()
+@click.option('-c', '--checkpoint', required=True)
+@click.option('-o', '--output_dir', required=True)
+@click.option('-d', '--device', default='cuda:0')
+def main(checkpoint, output_dir, device):
+    if os.path.exists(output_dir):
+        click.confirm(f"Output path {output_dir} already exists! Overwrite?", abort=True)
+    pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)
+    
+    # load checkpoint
+    payload = torch.load(open(checkpoint, 'rb'), pickle_module=dill)
+    cfg = payload['cfg']
+    cls = hydra.utils.get_class(cfg._target_)
+    workspace = cls(cfg, output_dir=output_dir)
+    workspace: BaseWorkspace
+    workspace.load_payload(payload, exclude_keys=None, include_keys=None)
+    
+    # get policy from workspace
+    policy = workspace.model
+    if cfg.training.use_ema:
+        policy = workspace.ema_model
+    
+    device = torch.device(device)
+    policy.to(device)
+    policy.eval()
+    
+    # run eval
+    env_runner = hydra.utils.instantiate(
+        cfg.task.env_runner,
+        output_dir=output_dir)
+    runner_log = env_runner.run(policy)
+    
+    # dump log to json
+    json_log = dict()
+    for key, value in runner_log.items():
+        if isinstance(value, wandb.sdk.data_types.video.Video):
+            json_log[key] = value._path
+        else:
+            json_log[key] = value
+    out_path = os.path.join(output_dir, 'eval_log.json')
+    json.dump(json_log, open(out_path, 'w'), indent=2, sort_keys=True)
+
+if __name__ == '__main__':
+    main()

third_party/diffusion_policy/eval_real_robot.py

@@ -0,0 +1,418 @@
+"""
+Usage:
+(robodiff)$ python eval_real_robot.py -i <ckpt_path> -o <save_dir> --robot_ip <ip_of_ur5>
+
+================ Human in control ==============
+Robot movement:
+Move your SpaceMouse to move the robot EEF (locked in xy plane).
+Press SpaceMouse right button to unlock z axis.
+Press SpaceMouse left button to enable rotation axes.
+
+Recording control:
+Click the opencv window (make sure it's in focus).
+Press "C" to start evaluation (hand control over to policy).
+Press "Q" to exit program.
+
+================ Policy in control ==============
+Make sure you can hit the robot hardware emergency-stop button quickly! 
+
+Recording control:
+Press "S" to stop evaluation and gain control back.
+"""
+
+# %%
+import time
+from multiprocessing.managers import SharedMemoryManager
+import click
+import cv2
+import numpy as np
+import torch
+import dill
+import hydra
+import pathlib
+import skvideo.io
+from omegaconf import OmegaConf
+import scipy.spatial.transform as st
+from diffusion_policy.real_world.real_env import RealEnv
+from diffusion_policy.real_world.spacemouse_shared_memory import Spacemouse
+from diffusion_policy.common.precise_sleep import precise_wait
+from diffusion_policy.real_world.real_inference_util import (
+    get_real_obs_resolution, 
+    get_real_obs_dict)
+from diffusion_policy.common.pytorch_util import dict_apply
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+from diffusion_policy.policy.base_image_policy import BaseImagePolicy
+from diffusion_policy.common.cv2_util import get_image_transform
+
+
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+@click.command()
+@click.option('--input', '-i', required=True, help='Path to checkpoint')
+@click.option('--output', '-o', required=True, help='Directory to save recording')
+@click.option('--robot_ip', '-ri', required=True, help="UR5's IP address e.g. 192.168.0.204")
+@click.option('--match_dataset', '-m', default=None, help='Dataset used to overlay and adjust initial condition')
+@click.option('--match_episode', '-me', default=None, type=int, help='Match specific episode from the match dataset')
+@click.option('--vis_camera_idx', default=0, type=int, help="Which RealSense camera to visualize.")
+@click.option('--init_joints', '-j', is_flag=True, default=False, help="Whether to initialize robot joint configuration in the beginning.")
+@click.option('--steps_per_inference', '-si', default=6, type=int, help="Action horizon for inference.")
+@click.option('--max_duration', '-md', default=60, help='Max duration for each epoch in seconds.')
+@click.option('--frequency', '-f', default=10, type=float, help="Control frequency in Hz.")
+@click.option('--command_latency', '-cl', default=0.01, type=float, help="Latency between receiving SapceMouse command to executing on Robot in Sec.")
+def main(input, output, robot_ip, match_dataset, match_episode,
+    vis_camera_idx, init_joints, 
+    steps_per_inference, max_duration,
+    frequency, command_latency):
+    # load match_dataset
+    match_camera_idx = 0
+    episode_first_frame_map = dict()
+    if match_dataset is not None:
+        match_dir = pathlib.Path(match_dataset)
+        match_video_dir = match_dir.joinpath('videos')
+        for vid_dir in match_video_dir.glob("*/"):
+            episode_idx = int(vid_dir.stem)
+            match_video_path = vid_dir.joinpath(f'{match_camera_idx}.mp4')
+            if match_video_path.exists():
+                frames = skvideo.io.vread(
+                    str(match_video_path), num_frames=1)
+                episode_first_frame_map[episode_idx] = frames[0]
+    print(f"Loaded initial frame for {len(episode_first_frame_map)} episodes")
+    
+    # load checkpoint
+    ckpt_path = input
+    payload = torch.load(open(ckpt_path, 'rb'), pickle_module=dill)
+    cfg = payload['cfg']
+    cls = hydra.utils.get_class(cfg._target_)
+    workspace = cls(cfg)
+    workspace: BaseWorkspace
+    workspace.load_payload(payload, exclude_keys=None, include_keys=None)
+
+    # hacks for method-specific setup.
+    action_offset = 0
+    delta_action = False
+    if 'diffusion' in cfg.name:
+        # diffusion model
+        policy: BaseImagePolicy
+        policy = workspace.model
+        if cfg.training.use_ema:
+            policy = workspace.ema_model
+
+        device = torch.device('cuda')
+        policy.eval().to(device)
+
+        # set inference params
+        policy.num_inference_steps = 16 # DDIM inference iterations
+        policy.n_action_steps = policy.horizon - policy.n_obs_steps + 1
+
+    elif 'robomimic' in cfg.name:
+        # BCRNN model
+        policy: BaseImagePolicy
+        policy = workspace.model
+
+        device = torch.device('cuda')
+        policy.eval().to(device)
+
+        # BCRNN always has action horizon of 1
+        steps_per_inference = 1
+        action_offset = cfg.n_latency_steps
+        delta_action = cfg.task.dataset.get('delta_action', False)
+
+    elif 'ibc' in cfg.name:
+        policy: BaseImagePolicy
+        policy = workspace.model
+        policy.pred_n_iter = 5
+        policy.pred_n_samples = 4096
+
+        device = torch.device('cuda')
+        policy.eval().to(device)
+        steps_per_inference = 1
+        action_offset = 1
+        delta_action = cfg.task.dataset.get('delta_action', False)
+    else:
+        raise RuntimeError("Unsupported policy type: ", cfg.name)
+
+    # setup experiment
+    dt = 1/frequency
+
+    obs_res = get_real_obs_resolution(cfg.task.shape_meta)
+    n_obs_steps = cfg.n_obs_steps
+    print("n_obs_steps: ", n_obs_steps)
+    print("steps_per_inference:", steps_per_inference)
+    print("action_offset:", action_offset)
+
+    with SharedMemoryManager() as shm_manager:
+        with Spacemouse(shm_manager=shm_manager) as sm, RealEnv(
+            output_dir=output, 
+            robot_ip=robot_ip, 
+            frequency=frequency,
+            n_obs_steps=n_obs_steps,
+            obs_image_resolution=obs_res,
+            obs_float32=True,
+            init_joints=init_joints,
+            enable_multi_cam_vis=True,
+            record_raw_video=True,
+            # number of threads per camera view for video recording (H.264)
+            thread_per_video=3,
+            # video recording quality, lower is better (but slower).
+            video_crf=21,
+            shm_manager=shm_manager) as env:
+            cv2.setNumThreads(1)
+
+            # Should be the same as demo
+            # realsense exposure
+            env.realsense.set_exposure(exposure=120, gain=0)
+            # realsense white balance
+            env.realsense.set_white_balance(white_balance=5900)
+
+            print("Waiting for realsense")
+            time.sleep(1.0)
+
+            print("Warming up policy inference")
+            obs = env.get_obs()
+            with torch.no_grad():
+                policy.reset()
+                obs_dict_np = get_real_obs_dict(
+                    env_obs=obs, shape_meta=cfg.task.shape_meta)
+                obs_dict = dict_apply(obs_dict_np, 
+                    lambda x: torch.from_numpy(x).unsqueeze(0).to(device))
+                result = policy.predict_action(obs_dict)
+                action = result['action'][0].detach().to('cpu').numpy()
+                assert action.shape[-1] == 2
+                del result
+
+            print('Ready!')
+            while True:
+                # ========= human control loop ==========
+                print("Human in control!")
+                state = env.get_robot_state()
+                target_pose = state['TargetTCPPose']
+                t_start = time.monotonic()
+                iter_idx = 0
+                while True:
+                    # calculate timing
+                    t_cycle_end = t_start + (iter_idx + 1) * dt
+                    t_sample = t_cycle_end - command_latency
+                    t_command_target = t_cycle_end + dt
+
+                    # pump obs
+                    obs = env.get_obs()
+
+                    # visualize
+                    episode_id = env.replay_buffer.n_episodes
+                    vis_img = obs[f'camera_{vis_camera_idx}'][-1]
+                    match_episode_id = episode_id
+                    if match_episode is not None:
+                        match_episode_id = match_episode
+                    if match_episode_id in episode_first_frame_map:
+                        match_img = episode_first_frame_map[match_episode_id]
+                        ih, iw, _ = match_img.shape
+                        oh, ow, _ = vis_img.shape
+                        tf = get_image_transform(
+                            input_res=(iw, ih), 
+                            output_res=(ow, oh), 
+                            bgr_to_rgb=False)
+                        match_img = tf(match_img).astype(np.float32) / 255
+                        vis_img = np.minimum(vis_img, match_img)
+
+                    text = f'Episode: {episode_id}'
+                    cv2.putText(
+                        vis_img,
+                        text,
+                        (10,20),
+                        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                        fontScale=0.5,
+                        thickness=1,
+                        color=(255,255,255)
+                    )
+                    cv2.imshow('default', vis_img[...,::-1])
+                    key_stroke = cv2.pollKey()
+                    if key_stroke == ord('q'):
+                        # Exit program
+                        env.end_episode()
+                        exit(0)
+                    elif key_stroke == ord('c'):
+                        # Exit human control loop
+                        # hand control over to the policy
+                        break
+
+                    precise_wait(t_sample)
+                    # get teleop command
+                    sm_state = sm.get_motion_state_transformed()
+                    # print(sm_state)
+                    dpos = sm_state[:3] * (env.max_pos_speed / frequency)
+                    drot_xyz = sm_state[3:] * (env.max_rot_speed / frequency)
+  
+                    if not sm.is_button_pressed(0):
+                        # translation mode
+                        drot_xyz[:] = 0
+                    else:
+                        dpos[:] = 0
+                    if not sm.is_button_pressed(1):
+                        # 2D translation mode
+                        dpos[2] = 0    
+
+                    drot = st.Rotation.from_euler('xyz', drot_xyz)
+                    target_pose[:3] += dpos
+                    target_pose[3:] = (drot * st.Rotation.from_rotvec(
+                        target_pose[3:])).as_rotvec()
+                    # clip target pose
+                    target_pose[:2] = np.clip(target_pose[:2], [0.25, -0.45], [0.77, 0.40])
+
+                    # execute teleop command
+                    env.exec_actions(
+                        actions=[target_pose], 
+                        timestamps=[t_command_target-time.monotonic()+time.time()])
+                    precise_wait(t_cycle_end)
+                    iter_idx += 1
+                
+                # ========== policy control loop ==============
+                try:
+                    # start episode
+                    policy.reset()
+                    start_delay = 1.0
+                    eval_t_start = time.time() + start_delay
+                    t_start = time.monotonic() + start_delay
+                    env.start_episode(eval_t_start)
+                    # wait for 1/30 sec to get the closest frame actually
+                    # reduces overall latency
+                    frame_latency = 1/30
+                    precise_wait(eval_t_start - frame_latency, time_func=time.time)
+                    print("Started!")
+                    iter_idx = 0
+                    term_area_start_timestamp = float('inf')
+                    perv_target_pose = None
+                    while True:
+                        # calculate timing
+                        t_cycle_end = t_start + (iter_idx + steps_per_inference) * dt
+
+                        # get obs
+                        print('get_obs')
+                        obs = env.get_obs()
+                        obs_timestamps = obs['timestamp']
+                        print(f'Obs latency {time.time() - obs_timestamps[-1]}')
+
+                        # run inference
+                        with torch.no_grad():
+                            s = time.time()
+                            obs_dict_np = get_real_obs_dict(
+                                env_obs=obs, shape_meta=cfg.task.shape_meta)
+                            obs_dict = dict_apply(obs_dict_np, 
+                                lambda x: torch.from_numpy(x).unsqueeze(0).to(device))
+                            result = policy.predict_action(obs_dict)
+                            # this action starts from the first obs step
+                            action = result['action'][0].detach().to('cpu').numpy()
+                            print('Inference latency:', time.time() - s)
+                        
+                        # convert policy action to env actions
+                        if delta_action:
+                            assert len(action) == 1
+                            if perv_target_pose is None:
+                                perv_target_pose = obs['robot_eef_pose'][-1]
+                            this_target_pose = perv_target_pose.copy()
+                            this_target_pose[[0,1]] += action[-1]
+                            perv_target_pose = this_target_pose
+                            this_target_poses = np.expand_dims(this_target_pose, axis=0)
+                        else:
+                            this_target_poses = np.zeros((len(action), len(target_pose)), dtype=np.float64)
+                            this_target_poses[:] = target_pose
+                            this_target_poses[:,[0,1]] = action
+
+                        # deal with timing
+                        # the same step actions are always the target for
+                        action_timestamps = (np.arange(len(action), dtype=np.float64) + action_offset
+                            ) * dt + obs_timestamps[-1]
+                        action_exec_latency = 0.01
+                        curr_time = time.time()
+                        is_new = action_timestamps > (curr_time + action_exec_latency)
+                        if np.sum(is_new) == 0:
+                            # exceeded time budget, still do something
+                            this_target_poses = this_target_poses[[-1]]
+                            # schedule on next available step
+                            next_step_idx = int(np.ceil((curr_time - eval_t_start) / dt))
+                            action_timestamp = eval_t_start + (next_step_idx) * dt
+                            print('Over budget', action_timestamp - curr_time)
+                            action_timestamps = np.array([action_timestamp])
+                        else:
+                            this_target_poses = this_target_poses[is_new]
+                            action_timestamps = action_timestamps[is_new]
+
+                        # clip actions
+                        this_target_poses[:,:2] = np.clip(
+                            this_target_poses[:,:2], [0.25, -0.45], [0.77, 0.40])
+
+                        # execute actions
+                        env.exec_actions(
+                            actions=this_target_poses,
+                            timestamps=action_timestamps
+                        )
+                        print(f"Submitted {len(this_target_poses)} steps of actions.")
+
+                        # visualize
+                        episode_id = env.replay_buffer.n_episodes
+                        vis_img = obs[f'camera_{vis_camera_idx}'][-1]
+                        text = 'Episode: {}, Time: {:.1f}'.format(
+                            episode_id, time.monotonic() - t_start
+                        )
+                        cv2.putText(
+                            vis_img,
+                            text,
+                            (10,20),
+                            fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                            fontScale=0.5,
+                            thickness=1,
+                            color=(255,255,255)
+                        )
+                        cv2.imshow('default', vis_img[...,::-1])
+
+
+                        key_stroke = cv2.pollKey()
+                        if key_stroke == ord('s'):
+                            # Stop episode
+                            # Hand control back to human
+                            env.end_episode()
+                            print('Stopped.')
+                            break
+
+                        # auto termination
+                        terminate = False
+                        if time.monotonic() - t_start > max_duration:
+                            terminate = True
+                            print('Terminated by the timeout!')
+
+                        term_pose = np.array([ 3.40948500e-01,  2.17721816e-01,  4.59076878e-02,  2.22014183e+00, -2.22184883e+00, -4.07186655e-04])
+                        curr_pose = obs['robot_eef_pose'][-1]
+                        dist = np.linalg.norm((curr_pose - term_pose)[:2], axis=-1)
+                        if dist < 0.03:
+                            # in termination area
+                            curr_timestamp = obs['timestamp'][-1]
+                            if term_area_start_timestamp > curr_timestamp:
+                                term_area_start_timestamp = curr_timestamp
+                            else:
+                                term_area_time = curr_timestamp - term_area_start_timestamp
+                                if term_area_time > 0.5:
+                                    terminate = True
+                                    print('Terminated by the policy!')
+                        else:
+                            # out of the area
+                            term_area_start_timestamp = float('inf')
+
+                        if terminate:
+                            env.end_episode()
+                            break
+
+                        # wait for execution
+                        precise_wait(t_cycle_end - frame_latency)
+                        iter_idx += steps_per_inference
+
+                except KeyboardInterrupt:
+                    print("Interrupted!")
+                    # stop robot.
+                    env.end_episode()
+                
+                print("Stopped.")
+
+
+
+# %%
+if __name__ == '__main__':
+    main()

third_party/diffusion_policy/multirun_metrics.py

@@ -0,0 +1,267 @@
+from typing import List, Optional
+import pathlib
+import pandas as pd
+import numpy as np
+import numba
+import click
+import time
+import collections
+import json
+import wandb
+import yaml
+import numbers
+import scipy.ndimage as sn
+from diffusion_policy.common.json_logger import read_json_log, JsonLogger
+import logging
+
+@numba.jit(nopython=True)
+def get_indexed_window_average(
+        arr: np.ndarray, idxs: np.ndarray, window_size: int):
+    result = np.zeros(idxs.shape, dtype=arr.dtype)
+    length = arr.shape[0]
+    for i in range(len(idxs)):
+        idx = idxs[i]
+        start = max(idx - window_size, 0)
+        end = min(start + window_size, length)
+        result[i] = np.mean(arr[start:end])
+    return result
+
+
+def compute_metrics(log_df: pd.DataFrame, key: str, 
+        end_step: Optional[int]=None,
+        k_min_loss: int=10,
+        k_around_max: int=10,
+        max_k_window: int=10,
+        replace_slash: int=True,
+        ):
+    if key not in log_df:
+        return dict()
+
+    # prepare data
+    if end_step is not None:
+        log_df = log_df.iloc[:end_step]
+    is_key = ~pd.isnull(log_df[key])
+    is_key_idxs = is_key.index[is_key].to_numpy()
+    if len(is_key_idxs) == 0:
+        return dict()
+
+    key_data = log_df[key][is_key].to_numpy()
+    # after adding validation to workspace
+    # rollout happens at the last step of each epoch
+    # where the reported train_loss and val_loss
+    # are already the average for that epoch
+    train_loss = log_df['train_loss'][is_key].to_numpy()
+    val_loss = log_df['val_loss'][is_key].to_numpy()
+
+    result = dict()
+
+    log_key = key
+    if replace_slash:
+        log_key = key.replace('/', '_')
+    # max
+    max_value = np.max(key_data)
+    result['max/'+log_key] = max_value
+
+    # k_around_max
+    max_idx = np.argmax(key_data)
+    end = min(max_idx + k_around_max // 2, len(key_data))
+    start = max(end - k_around_max, 0)
+    k_around_max_value = np.mean(key_data[start:end])
+    result['k_around_max/'+log_key] = k_around_max_value
+
+    # max_k_window
+    k_window_value = sn.uniform_filter1d(key_data, size=max_k_window, axis=0, mode='nearest')
+    max_k_window_value = np.max(k_window_value)
+    result['max_k_window/'+log_key] = max_k_window_value
+
+    # min_train_loss
+    min_idx = np.argmin(train_loss)
+    min_train_loss_value = key_data[min_idx]
+    result['min_train_loss/'+log_key] = min_train_loss_value
+
+    # min_val_loss
+    min_idx = np.argmin(val_loss)
+    min_val_loss_value = key_data[min_idx]
+    result['min_val_loss/'+log_key] = min_val_loss_value
+
+    # k_min_train_loss
+    min_loss_idxs = np.argsort(train_loss)[:k_min_loss]
+    k_min_train_loss_value = np.mean(key_data[min_loss_idxs])
+    result['k_min_train_loss/'+log_key] = k_min_train_loss_value
+
+    # k_min_val_loss
+    min_loss_idxs = np.argsort(val_loss)[:k_min_loss]
+    k_min_val_loss_value = np.mean(key_data[min_loss_idxs])
+    result['k_min_val_loss/'+log_key] = k_min_val_loss_value
+
+    # last
+    result['last/'+log_key] = key_data[-1]
+
+    # global step for visualization
+    result['metric_global_step/'+log_key] = is_key_idxs[-1]
+    return result
+
+
+def compute_metrics_agg(
+        log_dfs: List[pd.DataFrame], 
+        key: str, end_step:int, 
+        **kwargs):
+    
+    # compute metrics
+    results = collections.defaultdict(list)
+    for log_df in log_dfs:
+        result = compute_metrics(log_df, key=key, end_step=end_step, **kwargs)
+        for k, v in result.items():
+            results[k].append(v)
+    # agg
+    agg_result = dict()
+    for k, v in results.items():
+        value = np.mean(v)
+        if k.startswith('metric_global_step'):
+            value = int(value)
+        agg_result[k] = value
+    return agg_result
+
+
+@click.command()
+@click.option('--input', '-i', required=True, help='Root logging dir, contains train_* dirs')
+@click.option('--key', '-k', multiple=True, default=['test/mean_score'])
+@click.option('--interval', default=10, type=float)
+@click.option('--replace_slash', default=True, type=bool)
+@click.option('--index_key', '-ik', default='global_step')
+@click.option('--use_wandb', '-w', is_flag=True, default=False)
+@click.option('--project', default=None)
+@click.option('--name', default=None)
+@click.option('--id', default=None)
+@click.option('--group', default=None)
+def main(
+    input,
+    key,
+    interval,
+    replace_slash,
+    index_key,
+    use_wandb,
+    # wandb args
+    project,
+    name,
+    id,
+    group):
+    root_dir = pathlib.Path(input)
+    assert root_dir.is_dir()
+    metrics_dir = root_dir.joinpath('metrics')
+    metrics_dir.mkdir(exist_ok=True)
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(message)s",
+        handlers=[
+            logging.FileHandler(str(metrics_dir.joinpath("metrics.log"))),
+            logging.StreamHandler()
+        ]
+    )
+    
+    train_dirs = list(root_dir.glob('train_*'))
+    log_files = [x.joinpath('logs.json.txt') for x in train_dirs]
+    logging.info("Monitor waiting for log files!")
+    while True:
+        # wait for files to show up
+        files_exist = True
+        for log_file in log_files:
+            if not log_file.is_file():
+                files_exist = False
+        if files_exist:
+            break
+        time.sleep(1.0)
+    logging.info("All log files ready!")
+
+    # init path
+    metric_log_path = metrics_dir.joinpath('logs.json.txt')
+    metric_path = metrics_dir.joinpath('metrics.json')
+    config_path = root_dir.joinpath('config.yaml')
+
+    # load config
+    config = yaml.safe_load(config_path.open('r'))
+
+    # init wandb
+    wandb_run = None
+    if use_wandb:
+        wandb_kwargs = config['logging']
+        if project is not None:
+            wandb_kwargs['project'] = project
+        if id is not None:
+            wandb_kwargs['id'] = id
+        if name is not None:
+            wandb_kwargs['name'] = name
+        if group is not None:
+            wandb_kwargs['group'] = group
+        wandb_kwargs['resume'] = True
+        wandb_run = wandb.init(
+            dir=str(metrics_dir),
+            config=config,
+            # auto-resume run, automatically load id
+            # as long as using the same dir.
+            # https://docs.wandb.ai/guides/track/advanced/resuming#resuming-guidance
+            **wandb_kwargs
+        )
+        wandb.config.update(
+            {
+                "output_dir": str(root_dir),
+            }
+        )
+
+    with JsonLogger(metric_log_path) as json_logger:
+        last_log = json_logger.get_last_log()
+        while True:
+            # read json files
+            log_dfs = [read_json_log(str(x), required_keys=key) for x in log_files]
+
+            # previously logged data point
+            last_log_idx = -1
+            if last_log is not None:
+                last_log_idx = log_dfs[0].index[log_dfs[0][index_key] <= last_log[index_key]][-1]
+            
+            start_idx = last_log_idx + 1
+            # last idx where we have a data point from all logs
+            end_idx = min(*[len(x) for x in log_dfs])
+
+            # log every position
+            for this_idx in range(start_idx, end_idx):
+                # compute metrics
+                all_metrics = dict()
+                global_step = log_dfs[0]['global_step'][this_idx]
+                epoch = log_dfs[0]['epoch'][this_idx]
+                all_metrics['global_step'] = global_step
+                all_metrics['epoch'] = epoch
+                for k in key:
+                    metrics = compute_metrics_agg(
+                        log_dfs=log_dfs, key=k, end_step=this_idx+1,
+                        replace_slash=replace_slash)
+                    all_metrics.update(metrics)
+
+                # sanitize metrics
+                old_metrics = all_metrics
+                all_metrics = dict()
+                for k, v in old_metrics.items():
+                    if isinstance(v, numbers.Integral):
+                        all_metrics[k] = int(v)
+                    elif isinstance(v, numbers.Number):
+                        all_metrics[k] = float(v)
+                
+                has_update = all_metrics != last_log
+                if has_update:
+                    last_log = all_metrics
+                    json_logger.log(all_metrics)
+
+                    with metric_path.open('w') as f:
+                        json.dump(all_metrics, f, sort_keys=True, indent=2)
+
+                    if wandb_run is not None:
+                        wandb_run.log(all_metrics, step=all_metrics[index_key])
+
+                    logging.info(f"Metrics logged at step {all_metrics[index_key]}")
+            
+            time.sleep(interval)
+
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/pyrightconfig.json

@@ -0,0 +1,7 @@
+{
+    "exclude": [
+        "data/**",
+        "data_local/**",
+        "outputs/**"
+    ]
+}

third_party/diffusion_policy/ray_exec.py

@@ -0,0 +1,121 @@
+"""
+Usage:
+Training:
+python train.py --config-name=train_diffusion_lowdim_workspace -- logger.mode=online
+"""
+import os
+import ray
+import click
+
+def worker_fn(command_args, data_src=None, unbuffer_python=False, use_shell=False):
+    import os
+    import subprocess
+    import signal
+    import time
+
+    # setup data symlink
+    if data_src is not None:
+        cwd = os.getcwd()
+        src = os.path.expanduser(data_src)
+        dst = os.path.join(cwd, 'data')
+        try:
+            os.symlink(src=src, dst=dst)
+        except FileExistsError:
+            # it's fine if it already exists
+            pass
+
+    # run command
+    process_env = os.environ.copy()
+    if unbuffer_python:
+        # disable stdout/stderr buffering for subprocess (if python)
+        # to remove latency between print statement and receiving printed result
+        process_env['PYTHONUNBUFFERED'] = 'TRUE'
+    
+    # ray worker masks out Ctrl-C signal (ie SIGINT)
+    # here we unblock this signal for the child process
+    def preexec_function():
+        import signal
+        signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGINT})
+    
+    if use_shell:
+        command_args = ' '.join(command_args)
+
+    # stdout passthrough to ray worker, which is then passed to ray driver
+    process = subprocess.Popen(
+        args=command_args, 
+        env=process_env,
+        preexec_fn=preexec_function,
+        shell=use_shell)
+
+    while process.poll() is None:
+        try:
+            # sleep to ensure that monitor thread can acquire gil
+            # and raise KeyboardInterrupt here. 
+            time.sleep(0.01)
+        except KeyboardInterrupt:
+            process.send_signal(signal.SIGINT)
+            print('SIGINT sent to subprocess')
+        except Exception as e:
+            process.terminate()
+            raise e
+
+    if process.returncode not in (0, -2):
+        print("Failed execution!")
+        raise RuntimeError("Failed execution.")
+    return process.returncode
+
+
+@click.command()
+@click.option('--ray_address', '-ra', default='auto')
+@click.option('--num_cpus', '-nc', default=7, type=float)
+@click.option('--num_gpus', '-ng', default=1, type=float)
+@click.option('--max_retries', '-mr', default=0, type=int)
+@click.option('--data_src', '-d', default='./data', type=str)
+@click.option('--unbuffer_python', '-u', is_flag=True, default=False)
+@click.argument('command_args', nargs=-1, type=str)
+def main(ray_address, 
+    num_cpus, num_gpus, max_retries, 
+    data_src, unbuffer_python, 
+    command_args):
+    # expand path
+    if data_src is not None:
+        data_src = os.path.abspath(os.path.expanduser(data_src))
+
+    # init ray
+    root_dir = os.path.dirname(__file__)
+    runtime_env = {
+        'working_dir': root_dir,
+        'excludes': ['.git']
+    }
+    ray.init(
+        address=ray_address, 
+        runtime_env=runtime_env
+    )
+    # remote worker func
+    worker_ray = ray.remote(worker_fn).options(
+        num_cpus=num_cpus, 
+        num_gpus=num_gpus,
+        max_retries=max_retries,
+        # resources=resources,
+        retry_exceptions=True
+        )
+    # run
+    task_ref = worker_ray.remote(command_args, data_src, unbuffer_python)
+
+    try:
+        # normal case
+        result = ray.get(task_ref)
+        print('Return code: ', result)
+    except KeyboardInterrupt:
+        # a KeyboardInterrupt will be raised in worker
+        ray.cancel(task_ref, force=False)
+        result = ray.get(task_ref)
+        print('Return code: ', result)
+    except Exception as e:
+        # worker will be terminated
+        ray.cancel(task_ref, force=True)
+        raise e
+    
+
+if __name__ == '__main__':
+    main()

third_party/diffusion_policy/ray_train_multirun.py

@@ -0,0 +1,271 @@
+"""
+Start local ray cluster
+(robodiff)$ export CUDA_VISIBLE_DEVICES=0,1,2 # select GPUs to be managed by the ray cluster
+(robodiff)$ ray start --head --num-gpus=3
+
+Training:
+python ray_train_multirun.py --config-name=train_diffusion_unet_lowdim_workspace --seeds=42,43,44 --monitor_key=test/mean_score -- logger.mode=online training.eval_first=True
+"""
+import os
+import ray
+import click
+import hydra
+import yaml
+import wandb
+import pathlib
+import collections
+from pprint import pprint
+from omegaconf import OmegaConf
+from ray_exec import worker_fn
+from ray.util.placement_group import (
+    placement_group,
+)
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+@click.command()
+@click.option('--config-name', '-cn', required=True, type=str)
+@click.option('--config-dir', '-cd', default=None, type=str)
+@click.option('--seeds', '-s', default='42,43,44', type=str)
+@click.option('--monitor_key', '-k', multiple=True, default=['test/mean_score'])
+@click.option('--ray_address', '-ra', default='auto')
+@click.option('--num_cpus', '-nc', default=7, type=float)
+@click.option('--num_gpus', '-ng', default=1, type=float)
+@click.option('--max_retries', '-mr', default=0, type=int)
+@click.option('--monitor_max_retires', default=3, type=int)
+@click.option('--data_src', '-d', default='./data', type=str)
+@click.option('--unbuffer_python', '-u', is_flag=True, default=False)
+@click.option('--single_node', '-sn', is_flag=True, default=False, help='run all experiments on a single machine')
+@click.argument('command_args', nargs=-1, type=str)
+def main(config_name, config_dir, seeds, monitor_key, ray_address, 
+    num_cpus, num_gpus, max_retries, monitor_max_retires,
+    data_src, unbuffer_python, 
+    single_node, command_args):
+    # parse args
+    seeds = [int(x) for x in seeds.split(',')]
+    # expand path
+    if data_src is not None:
+        data_src = os.path.abspath(os.path.expanduser(data_src))
+
+    # initialize hydra
+    if config_dir is None:
+        config_path_abs = pathlib.Path(__file__).parent.joinpath(
+            'diffusion_policy','config')
+        config_path_rel = str(config_path_abs.relative_to(pathlib.Path.cwd()))
+    else:
+        config_path_rel = config_dir
+
+    run_command_args = list()
+    monitor_command_args = list()
+    with hydra.initialize(
+        version_base=None, 
+        config_path=config_path_rel):
+
+        # generate raw config
+        cfg = hydra.compose(
+            config_name=config_name, 
+            overrides=command_args)
+        OmegaConf.resolve(cfg)
+    
+        # manually create output dir
+        output_dir = pathlib.Path(cfg.multi_run.run_dir)
+        output_dir.mkdir(parents=True, exist_ok=False)
+        config_path = output_dir.joinpath('config.yaml')
+        print(output_dir)
+
+        # save current config
+        yaml.dump(OmegaConf.to_container(cfg, resolve=True), 
+            config_path.open('w'), default_flow_style=False)
+
+        # wandb
+        wandb_group_id = wandb.util.generate_id()
+        name_base = cfg.multi_run.wandb_name_base
+
+        # create monitor command args
+        monitor_command_args = [
+            'python',
+            'multirun_metrics.py',
+            '--input', str(output_dir),
+            '--use_wandb',
+            '--project', 'diffusion_policy_metrics',
+            '--group', wandb_group_id
+        ]
+        for k in monitor_key:
+            monitor_command_args.extend([
+                '--key', k
+            ])
+
+        # generate command args
+        run_command_args = list()
+        for i, seed in enumerate(seeds):
+            test_start_seed = (seed + 1) * 100000
+            this_output_dir = output_dir.joinpath(f'train_{i}')
+            this_output_dir.mkdir()
+            wandb_name = name_base + f'_train_{i}'
+            wandb_run_id = wandb_group_id + f'_train_{i}'
+
+            this_command_args = [
+                'python',
+                'train.py',
+                '--config-name='+config_name,
+                '--config-dir='+config_path_rel
+            ]
+
+            this_command_args.extend(command_args)
+            this_command_args.extend([
+                f'training.seed={seed}',
+                f'task.env_runner.test_start_seed={test_start_seed}',
+                f'logging.name={wandb_name}',
+                f'logging.id={wandb_run_id}',
+                f'logging.group={wandb_group_id}',
+                f'hydra.run.dir={this_output_dir}'
+            ])
+            run_command_args.append(this_command_args)
+
+    # init ray
+    root_dir = os.path.dirname(__file__)
+    runtime_env = {
+        'working_dir': root_dir,
+        'excludes': ['.git'],
+        'pip': ['dm-control==1.0.9']
+    }
+    ray.init(
+        address=ray_address, 
+        runtime_env=runtime_env
+    )
+    
+    # create resources for train
+    train_resources = dict()
+
+    train_bundle = dict(train_resources)
+    train_bundle['CPU'] = num_cpus
+    train_bundle['GPU'] = num_gpus
+
+    # create resources for monitor
+    monitor_resources = dict()
+    monitor_resources['CPU'] = 1
+    
+    monitor_bundle = dict(monitor_resources)
+
+    # aggregate bundle
+    bundle = collections.defaultdict(lambda:0)
+    n_train_bundles = 1
+    if single_node:
+        n_train_bundles = len(seeds)
+    for _ in range(n_train_bundles):
+        for k, v in train_bundle.items():
+            bundle[k] += v
+    for k, v in monitor_bundle.items():
+        bundle[k] += v
+    bundle = dict(bundle)
+
+    # create placement group
+    print("Creating placement group with resources:")
+    pprint(bundle)
+    pg = placement_group([bundle])
+
+    # run
+    task_name_map = dict()
+    task_refs = list()
+    for i, this_command_args in enumerate(run_command_args):
+        if single_node or i == (len(run_command_args) - 1):
+            print(f'Training worker {i} with placement group.')
+            ray.get(pg.ready())
+            print("Placement Group created!")
+            worker_ray = ray.remote(worker_fn).options(
+                num_cpus=num_cpus,
+                num_gpus=num_gpus,
+                max_retries=max_retries,
+                resources=train_resources,
+                retry_exceptions=True,
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pg)
+            )
+        else:
+            print(f'Training worker {i} without placement group.')
+            worker_ray = ray.remote(worker_fn).options(
+                num_cpus=num_cpus,
+                num_gpus=num_gpus,
+                max_retries=max_retries,
+                resources=train_resources,
+                retry_exceptions=True,
+            )
+        task_ref = worker_ray.remote(
+            this_command_args, data_src, unbuffer_python)
+        task_refs.append(task_ref)
+        task_name_map[task_ref] = f'train_{i}'
+
+    # monitor worker is always packed on the same node
+    # as training worker 0
+    ray.get(pg.ready())
+    monitor_worker_ray = ray.remote(worker_fn).options(
+        num_cpus=1, 
+        num_gpus=0,
+        max_retries=monitor_max_retires,
+        # resources=monitor_resources,
+        retry_exceptions=True,
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg)
+        )
+    monitor_ref = monitor_worker_ray.remote(
+            monitor_command_args, data_src, unbuffer_python)
+    task_name_map[monitor_ref] = 'metrics'
+
+    try:
+        # normal case
+        ready_refs = list()
+        rest_refs = task_refs
+        while len(ready_refs) < len(task_refs):
+            this_ready_refs, rest_refs = ray.wait(rest_refs, 
+                num_returns=1, timeout=None, fetch_local=True)
+            cancel_other_tasks = False
+            for ref in this_ready_refs:
+                task_name = task_name_map[ref]
+                try:
+                    result = ray.get(ref)
+                    print(f"Task {task_name} finished with result: {result}")
+                except KeyboardInterrupt as e:
+                    # skip to outer try catch
+                    raise KeyboardInterrupt
+                except Exception as e:
+                    print(f"Task {task_name} raised exception: {e}")
+                    this_cancel_other_tasks = True
+                    if isinstance(e, ray.exceptions.RayTaskError):
+                        if isinstance(e.cause, ray.exceptions.TaskCancelledError):
+                            this_cancel_other_tasks = False
+                    cancel_other_tasks = cancel_other_tasks or this_cancel_other_tasks
+                ready_refs.append(ref)
+            if cancel_other_tasks:
+                print('Exception! Cancelling all other tasks.')
+                # cancel all other refs
+                for _ref in rest_refs:
+                    ray.cancel(_ref, force=False)
+        print("Training tasks done.")
+        ray.cancel(monitor_ref, force=False)
+    except KeyboardInterrupt:
+        print('KeyboardInterrupt received in the driver.')
+        # a KeyboardInterrupt will be raised in worker
+        _ = [ray.cancel(x, force=False) for x in task_refs + [monitor_ref]]
+        print('KeyboardInterrupt sent to workers.')
+    except Exception as e:
+        # worker will be terminated
+        _ = [ray.cancel(x, force=True) for x in task_refs + [monitor_ref]]
+        raise e
+
+    for ref in task_refs + [monitor_ref]:
+        task_name = task_name_map[ref]
+        try:
+            result = ray.get(ref)
+            print(f"Task {task_name} finished with result: {result}")
+        except KeyboardInterrupt as e:
+            # force kill everything.
+            print("Force killing all workers")
+            _ = [ray.cancel(x, force=True) for x in task_refs]
+            ray.cancel(monitor_ref, force=True)
+        except Exception as e:
+            print(f"Task {task_name} raised exception: {e}")
+
+
+if __name__ == "__main__":
+    main()

third_party/diffusion_policy/setup.py

@@ -0,0 +1,6 @@
+from setuptools import setup, find_packages
+
+setup(
+  name = 'diffusion_policy',
+  packages = find_packages(),
+)

third_party/diffusion_policy/tests/test_block_pushing.py

@@ -0,0 +1,44 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+from diffusion_policy.env.block_pushing.block_pushing_multimodal import BlockPushMultimodal
+from gym.wrappers import FlattenObservation
+from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper
+from diffusion_policy.gym_util.video_wrapper import VideoWrapper
+
+def test():
+    env = MultiStepWrapper(
+            VideoWrapper(
+                FlattenObservation(
+                    BlockPushMultimodal()
+                ),
+                enabled=True,
+                steps_per_render=2
+            ),
+            n_obs_steps=2,
+            n_action_steps=8,
+            max_episode_steps=16
+        )
+    env = BlockPushMultimodal()
+    obs = env.reset()
+    import pdb; pdb.set_trace()
+
+    env = FlattenObservation(BlockPushMultimodal())
+    obs = env.reset()
+    action = env.action_space.sample()
+    next_obs, reward, done, info = env.step(action)
+    print(obs[8:10] + action - next_obs[8:10])
+    import pdb; pdb.set_trace()
+
+    for i in range(3):
+        obs, reward, done, info = env.step(env.action_space.sample())
+    img = env.render()
+    import pdb; pdb.set_trace()
+    print("Done!", done)
+
+if __name__ == '__main__':
+    test()

third_party/diffusion_policy/tests/test_cv2_util.py

@@ -0,0 +1,22 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import numpy as np
+from diffusion_policy.common.cv2_util import get_image_transform
+
+
+def test():
+    tf = get_image_transform((1280,720), (640,480), bgr_to_rgb=False)
+    in_img = np.zeros((720,1280,3), dtype=np.uint8)
+    out_img = tf(in_img)
+    # print(out_img.shape)
+    assert out_img.shape == (480,640,3)
+
+    # import pdb; pdb.set_trace()
+
+if __name__ == '__main__':
+    test()

third_party/diffusion_policy/tests/test_multi_realsense.py

@@ -0,0 +1,82 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import cv2
+import json
+import time
+import numpy as np
+from diffusion_policy.real_world.multi_realsense import MultiRealsense
+from diffusion_policy.real_world.video_recorder import VideoRecorder
+
+def test():
+    config = json.load(open('/home/cchi/dev/diffusion_policy/diffusion_policy/real_world/realsense_config/415_high_accuracy_mode.json', 'r'))
+
+    def transform(data):
+        color = data['color']
+        h,w,_ = color.shape
+        factor = 4
+        color = cv2.resize(color, (w//factor,h//factor), interpolation=cv2.INTER_AREA)
+        # color = color[:,140:500]
+        data['color'] = color
+        return data
+
+    from diffusion_policy.common.cv2_util import get_image_transform
+    color_transform = get_image_transform(
+        input_res=(1280,720),
+        output_res=(640,480), 
+        bgr_to_rgb=False)
+    def transform(data):
+        data['color'] = color_transform(data['color'])
+        return data
+
+    # one thread per camera
+    video_recorder = VideoRecorder.create_h264(
+        fps=30,
+        codec='h264',
+        thread_type='FRAME'
+    )
+
+    with MultiRealsense(
+            resolution=(1280,720),
+            capture_fps=30,
+            record_fps=15,
+            enable_color=True,
+            # advanced_mode_config=config,
+            transform=transform,
+            # recording_transform=transform,
+            # video_recorder=video_recorder,
+            verbose=True
+        ) as realsense:
+        realsense.set_exposure(exposure=150, gain=5)
+        intr = realsense.get_intrinsics()
+        print(intr)
+
+        video_path = 'data_local/test'
+        rec_start_time = time.time() + 1
+        realsense.start_recording(video_path, start_time=rec_start_time)
+        realsense.restart_put(rec_start_time)
+
+        out = None
+        vis_img = None
+        while True:
+            out = realsense.get(out=out)
+
+            # bgr = out['color']
+            # print(bgr.shape)
+            # vis_img = np.concatenate(list(bgr), axis=0, out=vis_img)
+            # cv2.imshow('default', vis_img)
+            # key = cv2.pollKey()
+            # if key == ord('q'):
+            #     break
+
+            time.sleep(1/60)
+            if time.time() > (rec_start_time + 20.0):
+                break
+
+
+if __name__ == "__main__":
+    test()

third_party/diffusion_policy/tests/test_pose_trajectory_interpolator.py

@@ -0,0 +1,126 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+from tqdm import tqdm
+import numpy as np
+import scipy.interpolate as si
+import scipy.spatial.transform as st
+from diffusion_policy.common.pose_trajectory_interpolator import (
+    rotation_distance, 
+    pose_distance, 
+    PoseTrajectoryInterpolator)
+
+
+def test_rotation_distance():
+    def rotation_distance_align(a: st.Rotation, b: st.Rotation) -> float:
+        return st.Rotation.align_vectors(b.as_matrix().T, a.as_matrix().T)[0].magnitude()
+
+    for i in range(10000):
+        a = st.Rotation.from_euler('xyz', np.random.uniform(-7,7,size=3))
+        b = st.Rotation.from_euler('xyz', np.random.uniform(-7,7,size=3))
+        x = rotation_distance(a, b)
+        y = rotation_distance_align(a, b)
+        assert abs(x-y) < 1e-7
+
+def test_pose_trajectory_interpolator():
+    t = np.linspace(-1,5,100)
+    interp = PoseTrajectoryInterpolator(
+        [0,1,3],
+        np.zeros((3,6))
+    )
+    times = interp.times
+    poses = interp.poses
+
+    trimmed_interp = interp.trim(-1,4)
+    assert len(trimmed_interp.times) == 5
+    trimmed_interp(t)
+
+    trimmed_interp = interp.trim(-1,4)
+    assert len(trimmed_interp.times) == 5
+    trimmed_interp(t)
+
+    trimmed_interp = interp.trim(0.5, 3.5)
+    assert len(trimmed_interp.times) == 4
+    trimmed_interp(t)
+
+    trimmed_interp = interp.trim(0.5, 2.5)
+    assert len(trimmed_interp.times) == 3
+    trimmed_interp(t)
+
+    trimmed_interp = interp.trim(0.5, 1.5)
+    assert len(trimmed_interp.times) == 3
+    trimmed_interp(t)
+
+    trimmed_interp = interp.trim(1.2, 1.5)
+    assert len(trimmed_interp.times) == 2
+    trimmed_interp(t)
+
+    trimmed_interp = interp.trim(1.3, 1.3)
+    assert len(trimmed_interp.times) == 1
+    trimmed_interp(t)
+
+    # import pdb; pdb.set_trace()
+
+def test_add_waypoint():
+    # fuzz testing
+    for i in tqdm(range(10000)):
+        rng = np.random.default_rng(i)
+        n_waypoints = rng.integers(1, 5)
+        waypoint_times = np.sort(rng.uniform(0, 1, size=n_waypoints))
+        last_waypoint_time = waypoint_times[-1]
+        insert_time = rng.uniform(-0.1, 1.1)
+        curr_time = rng.uniform(-0.1, 1.1)
+        max_pos_speed = rng.poisson(3) + 1e-3
+        max_rot_speed = rng.poisson(3) + 1e-3
+        waypoint_poses = rng.normal(0, 3, size=(n_waypoints, 6))
+        new_pose = rng.normal(0, 3, size=6)
+
+        if rng.random() < 0.1:
+            last_waypoint_time = None
+            if rng.random() < 0.1:
+                curr_time = None
+
+        interp = PoseTrajectoryInterpolator(
+            times=waypoint_times, 
+            poses=waypoint_poses)
+        new_interp = interp.add_waypoint(
+            pose=new_pose,
+            time=insert_time,
+            max_pos_speed=max_pos_speed,
+            max_rot_speed=max_rot_speed,
+            curr_time=curr_time,
+            last_waypoint_time=last_waypoint_time
+        )
+
+def test_drive_to_waypoint():
+    # fuzz testing
+    for i in tqdm(range(10000)):
+        rng = np.random.default_rng(i)
+        n_waypoints = rng.integers(1, 5)
+        waypoint_times = np.sort(rng.uniform(0, 1, size=n_waypoints))
+        insert_time = rng.uniform(-0.1, 1.1)
+        curr_time = rng.uniform(-0.1, 1.1)
+        max_pos_speed = rng.poisson(3) + 1e-3
+        max_rot_speed = rng.poisson(3) + 1e-3
+        waypoint_poses = rng.normal(0, 3, size=(n_waypoints, 6))
+        new_pose = rng.normal(0, 3, size=6)
+
+        interp = PoseTrajectoryInterpolator(
+            times=waypoint_times, 
+            poses=waypoint_poses)
+        new_interp = interp.drive_to_waypoint(
+            pose=new_pose,
+            time=insert_time,
+            curr_time=curr_time,
+            max_pos_speed=max_pos_speed,
+            max_rot_speed=max_rot_speed
+        )
+
+
+
+if __name__ == '__main__':
+    test_drive_to_waypoint()

third_party/diffusion_policy/tests/test_precise_sleep.py

@@ -0,0 +1,56 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import time
+import numpy as np
+from diffusion_policy.common.precise_sleep import precise_sleep, precise_wait
+
+
+def test_sleep():
+    dt = 0.1
+    tol = 1e-3
+    time_samples = list()
+    for i in range(100):
+        precise_sleep(dt)
+        # time.sleep(dt)
+        time_samples.append(time.monotonic())
+    time_deltas = np.diff(time_samples)
+
+    from matplotlib import pyplot as plt
+    plt.plot(time_deltas)
+    plt.ylim((dt-tol,dt+tol))
+
+
+def test_wait():
+    dt = 0.1
+    tol = 1e-3
+    errors = list()
+    t_start = time.monotonic()
+    for i in range(1,100):
+        t_end_desired = t_start + i * dt
+        time.sleep(t_end_desired - time.monotonic())
+        t_end = time.monotonic()
+        errors.append(t_end - t_end_desired)
+
+    new_errors = list()
+    t_start = time.monotonic()
+    for i in range(1,100):
+        t_end_desired = t_start + i * dt
+        precise_wait(t_end_desired)
+        t_end = time.monotonic()
+        new_errors.append(t_end - t_end_desired)
+    
+    from matplotlib import pyplot as plt
+    plt.plot(errors, label='time.sleep')
+    plt.plot(new_errors, label='sleep/spin hybrid')
+    plt.ylim((-tol,+tol))
+    plt.title('0.1 sec sleep error')
+    plt.legend()
+
+    
+if __name__ == '__main__':
+    test_sleep()

third_party/diffusion_policy/tests/test_replay_buffer.py

@@ -0,0 +1,62 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import zarr
+from diffusion_policy.common.replay_buffer import ReplayBuffer
+
+def test():
+    import numpy as np
+    buff = ReplayBuffer.create_empty_numpy()
+    buff.add_episode({
+        'obs': np.zeros((100,10), dtype=np.float16)
+    })
+    buff.add_episode({
+        'obs': np.ones((50,10)),
+        'action': np.ones((50,2))
+    })
+    # buff.rechunk(256)
+    obs = buff.get_episode(0)
+
+    import numpy as np
+    buff = ReplayBuffer.create_empty_zarr()
+    buff.add_episode({
+        'obs': np.zeros((100,10), dtype=np.float16)
+    })
+    buff.add_episode({
+        'obs': np.ones((50,10)),
+        'action': np.ones((50,2))
+    })
+    obs = buff.get_episode(0)
+    buff.set_chunks({
+        'obs': (100,10),
+        'action': (100,2)
+    })
+
+
+def test_real():
+    import os
+    dist_group = zarr.open(
+        os.path.expanduser('~/dev/diffusion_policy/data/pusht/pusht_cchi_v2.zarr'), 'r')
+
+    buff = ReplayBuffer.create_empty_numpy()
+    key, group = next(iter(dist_group.items()))
+    for key, group in dist_group.items():
+        buff.add_episode(group)
+
+    # out_path = os.path.expanduser('~/dev/diffusion_policy/data/pusht_cchi2_v2_replay.zarr')
+    out_path = os.path.expanduser('~/dev/diffusion_policy/data/test.zarr')
+    out_store = zarr.DirectoryStore(out_path)
+    buff.save_to_store(out_store)
+
+    buff = ReplayBuffer.copy_from_path(out_path, store=zarr.MemoryStore())
+    buff.pop_episode()
+
+
+def test_pop():
+    buff = ReplayBuffer.create_from_path(
+        '/home/chengchi/dev/diffusion_policy/data/pusht_cchi_v3_replay.zarr',
+        mode='rw')

third_party/diffusion_policy/tests/test_ring_buffer.py

@@ -0,0 +1,188 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import time
+import numpy as np
+import multiprocessing as mp
+from multiprocessing.managers import SharedMemoryManager
+from diffusion_policy.shared_memory.shared_memory_ring_buffer import (
+    SharedMemoryRingBuffer,
+    SharedAtomicCounter)
+
+
+def test():
+    shm_manager = SharedMemoryManager()
+    shm_manager.start()
+    ring_buffer = SharedMemoryRingBuffer.create_from_examples(
+        shm_manager,
+        {'timestamp': np.array(0, dtype=np.float64)},
+        buffer_size=128
+    )
+    for i in range(30):
+        ring_buffer.put({
+            'timestamp': np.array(
+                time.perf_counter(), 
+                dtype=np.float64)
+        })
+    print(ring_buffer.get())
+
+
+def _timestamp_worker(ring_buffer, start_event, stop_event):
+    while not stop_event.is_set():
+        start_event.set()
+        ring_buffer.put({
+            'timestamp': np.array(
+                time.time(), 
+                dtype=np.float64)
+        })
+
+
+def test_mp():
+    shm_manager = SharedMemoryManager()
+    shm_manager.start()
+    ring_buffer = SharedMemoryRingBuffer.create_from_examples(
+        shm_manager,
+        {'timestamp': np.array(0, dtype=np.float64)},
+        get_max_k=1,
+        get_time_budget=0.01,
+        put_desired_frequency=1000
+    )
+    start_event = mp.Event()
+    stop_event = mp.Event()
+    worker = mp.Process(target=_timestamp_worker, args=(
+        ring_buffer, start_event, stop_event))
+    worker.start()
+    start_event.wait()
+    for i in range(1000):
+        t = float(ring_buffer.get()['timestamp'])
+        curr_t = time.time()
+        print('latency', curr_t - t)
+    stop_event.set()
+    worker.join()
+
+
+def test_get_last_k():
+    shm_manager = SharedMemoryManager()
+    shm_manager.start()
+    ring_buffer = SharedMemoryRingBuffer.create_from_examples(
+        shm_manager,
+        {'counter': np.array(0, dtype=np.int64)},
+        buffer_size=8
+    )
+
+    from collections import deque
+    k = 4
+    last_k = deque(maxlen=k)
+    for i in range(100):
+        ring_buffer.put({
+            'counter': np.array(i, dtype=np.int64)
+        })
+        last_k.append(i)
+        if i > k:
+            result = ring_buffer.get_last_k(k)['counter']
+            assert np.allclose(result, last_k)
+
+    print(ring_buffer.shared_arrays['counter'].get())
+    result = ring_buffer.get_last_k(4)
+    print(result)
+
+
+def test_timing():
+    shm_manager = SharedMemoryManager()
+    shm_manager.start()
+    ring_buffer = SharedMemoryRingBuffer.create_from_examples(
+        shm_manager,
+        {'counter': np.array(0, dtype=np.int64)},
+        get_max_k=8,
+        get_time_budget=0.1,
+        put_desired_frequency=100
+    )
+    # print(ring_buffer.timestamp_array.get())
+    print('buffer_size', ring_buffer.buffer_size)
+
+    dt = 1 / 150
+    t_init = time.monotonic()
+    for i in range(1000):
+        t_start = time.monotonic()
+        ring_buffer.put({
+            'counter': np.array(i, dtype=np.int64)
+        }, wait=False)
+        if (i % 10 == 0) and (i > 0):
+            result = ring_buffer.get_last_k(8)
+
+        t_end =time.monotonic()
+        desired_t = (i+1) * dt + t_init
+        if desired_t > t_end:
+            time.sleep(desired_t - t_end)
+        hz = 1 / (time.monotonic() - t_start)
+        print(f'{hz}Hz')
+
+
+def _timestamp_image_worker(ring_buffer, img_shape, dt, start_event, stop_event):
+    i = 0
+    t_init = time.monotonic()
+    image = np.ones(img_shape, dtype=np.uint8)
+    while not stop_event.is_set():
+        t_start = time.monotonic()
+        start_event.set()
+        ring_buffer.put({
+            'img': image,
+            'timestamp': time.time(),
+            'counter': i
+        })
+        t_end = time.monotonic()
+        desired_t = (i+1) * dt + t_init
+        # print('alive')
+        if desired_t > t_end:
+            time.sleep(desired_t - t_end)
+        # hz = 1 / (time.monotonic() - t_start)
+        i += 1
+
+
+def test_timing_mp():
+    shm_manager = SharedMemoryManager()
+    shm_manager.start()
+
+    hz = 200
+    img_shape = (1920,1080,3)
+    ring_buffer = SharedMemoryRingBuffer.create_from_examples(
+        shm_manager,
+        examples={
+            'img': np.zeros(img_shape, dtype=np.uint8),
+            'timestamp': time.time(),
+            'counter': 0
+        },
+        get_max_k=60,
+        get_time_budget=0.02,
+        put_desired_frequency=hz
+    )
+    start_event = mp.Event()
+    stop_event = mp.Event()
+    worker = mp.Process(target=_timestamp_image_worker, args=(
+        ring_buffer, img_shape, 1/hz, start_event, stop_event))
+    worker.start()
+    start_event.wait()
+    out = None
+    t_start = time.monotonic()
+    k = 1
+    for i in range(1000):
+        if ring_buffer.count < k:
+            time.sleep(0)
+            continue
+        out = ring_buffer.get_last_k(k=k, out=out)
+        t = float(out['timestamp'][-1])
+        curr_t = time.time()
+        print('latency', curr_t - t)
+    t_end = time.monotonic()
+    print('Get Hz', 1/(t_end-t_start)*1000)
+    stop_event.set()
+    worker.join()
+
+
+if __name__ == '__main__':
+    # test_mp()
+    test_timing_mp()

third_party/diffusion_policy/tests/test_robomimic_image_runner.py

@@ -0,0 +1,38 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+from diffusion_policy.env_runner.robomimic_image_runner import RobomimicImageRunner
+
+def test():
+    import os
+    from omegaconf import OmegaConf
+    cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_image.yaml')
+    cfg = OmegaConf.load(cfg_path)
+    cfg['n_obs_steps'] = 1
+    cfg['n_action_steps'] = 1
+    cfg['past_action_visible'] = False
+    runner_cfg = cfg['env_runner']
+    runner_cfg['n_train'] = 1
+    runner_cfg['n_test'] = 1
+    del runner_cfg['_target_']
+    runner = RobomimicImageRunner(
+        **runner_cfg, 
+        output_dir='/tmp/test')
+
+    # import pdb; pdb.set_trace()
+
+    self = runner
+    env = self.env
+    env.seed(seeds=self.env_seeds)
+    obs = env.reset()
+    for i in range(10):
+        _ = env.step(env.action_space.sample())
+
+    imgs = env.render()
+
+if __name__ == '__main__':
+    test()

third_party/diffusion_policy/tests/test_robomimic_lowdim_runner.py

@@ -0,0 +1,34 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+from diffusion_policy.env_runner.robomimic_lowdim_runner import RobomimicLowdimRunner
+
+def test():
+    import os
+    from omegaconf import OmegaConf
+    cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_lowdim.yaml')
+    cfg = OmegaConf.load(cfg_path)
+    cfg['n_obs_steps'] = 1
+    cfg['n_action_steps'] = 1
+    cfg['past_action_visible'] = False
+    runner_cfg = cfg['env_runner']
+    runner_cfg['n_train'] = 1
+    runner_cfg['n_test'] = 0
+    del runner_cfg['_target_']
+    runner = RobomimicLowdimRunner(
+        **runner_cfg, 
+        output_dir='/tmp/test')
+
+    # import pdb; pdb.set_trace()
+
+    self = runner
+    env = self.env
+    env.seed(seeds=self.env_seeds)
+    obs = env.reset()
+
+if __name__ == '__main__':
+    test()

third_party/diffusion_policy/tests/test_shared_queue.py

@@ -0,0 +1,67 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import numpy as np
+from multiprocessing.managers import SharedMemoryManager
+from diffusion_policy.shared_memory.shared_memory_queue import SharedMemoryQueue, Full, Empty
+
+
+def test():
+    shm_manager = SharedMemoryManager()
+    shm_manager.start()
+    example = {
+        'cmd': 0,
+        'pose': np.zeros((6,))
+    }
+    queue = SharedMemoryQueue.create_from_examples(
+        shm_manager=shm_manager,
+        examples=example,
+        buffer_size=3
+    )
+    raised = False
+    try:
+        queue.get()
+    except Empty:
+        raised = True
+    assert raised
+
+    data = {
+        'cmd': 1,
+        'pose': np.ones((6,))
+    }
+    queue.put(data)
+    result = queue.get()
+    assert result['cmd'] == data['cmd']
+    assert np.allclose(result['pose'], data['pose'])
+
+    queue.put(data)
+    queue.put(data)
+    queue.put(data)
+    assert queue.qsize() == 3
+    raised = False
+    try:
+        queue.put(data)
+    except Full:
+        raised = True
+    assert raised
+
+    result = queue.get_all()
+    assert np.allclose(result['cmd'], [1,1,1])
+    
+    queue.put({'cmd': 0})
+    queue.put({'cmd': 1})
+    queue.put({'cmd': 2})
+    queue.get()
+    queue.put({'cmd': 3})
+
+    result = queue.get_k(3)
+    assert np.allclose(result['cmd'], [1,2,3])
+
+    queue.clear()
+
+if __name__ == "__main__":
+    test()

third_party/diffusion_policy/tests/test_single_realsense.py

@@ -0,0 +1,87 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import cv2
+import json
+import time
+from multiprocessing.managers import SharedMemoryManager
+from diffusion_policy.real_world.single_realsense import SingleRealsense
+
+def test():
+    
+    serials = SingleRealsense.get_connected_devices_serial()
+    # import pdb; pdb.set_trace()
+    serial = serials[0]
+    config = json.load(open('/home/cchi/dev/diffusion_policy/diffusion_policy/real_world/realsense_config/415_high_accuracy_mode.json', 'r'))
+
+    def transform(data):
+        color = data['color']
+        h,w,_ = color.shape
+        factor = 2
+        color = cv2.resize(color, (w//factor,h//factor), interpolation=cv2.INTER_AREA)
+        # color = color[:,140:500]
+        data['color'] = color
+        return data
+
+    # at 960x540 with //3, 60fps and 30fps are indistinguishable
+
+    with SharedMemoryManager() as shm_manager:
+        with SingleRealsense(
+            shm_manager=shm_manager,
+            serial_number=serial,
+            resolution=(1280,720),
+            # resolution=(960,540),
+            # resolution=(640,480),
+            capture_fps=30,
+            enable_color=True,
+            # enable_depth=True,
+            # enable_infrared=True,
+            # advanced_mode_config=config,
+            # transform=transform,
+            # recording_transform=transform
+            # verbose=True
+            ) as realsense:
+            cv2.setNumThreads(1) 
+            realsense.set_exposure(exposure=150, gain=5)
+            intr = realsense.get_intrinsics()
+            print(intr)
+
+
+            video_path = 'data_local/test.mp4'
+            rec_start_time = time.time() + 2
+            realsense.start_recording(video_path, start_time=rec_start_time)
+
+            data = None
+            while True:
+                data = realsense.get(out=data)
+                t = time.time()
+                # print('capture_latency', data['receive_timestamp']-data['capture_timestamp'], 'receive_latency', t - data['receive_timestamp'])
+                # print('receive', t - data['receive_timestamp'])
+
+                # dt = time.time() - data['timestamp']
+                # print(dt)
+                # print(data['capture_timestamp'] - rec_start_time)
+
+                bgr = data['color']
+                # print(bgr.shape)
+                cv2.imshow('default', bgr)
+                key = cv2.pollKey()
+                # if key == ord('q'):
+                #     break
+                # elif key == ord('r'):
+                #     video_path = 'data_local/test.mp4'
+                #     realsense.start_recording(video_path)
+                # elif key == ord('s'):
+                #     realsense.stop_recording()
+                
+                time.sleep(1/60)
+                if time.time() > (rec_start_time + 20.0):
+                    break
+
+
+if __name__ == "__main__":
+    test()

third_party/diffusion_policy/tests/test_timestamp_accumulator.py

@@ -0,0 +1,151 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import numpy as np
+import time
+from diffusion_policy.common.timestamp_accumulator import (
+    get_accumulate_timestamp_idxs,
+    TimestampObsAccumulator,
+    TimestampActionAccumulator
+)
+
+
+def test_index():
+    buffer = np.zeros(16)
+    start_time = 0.0
+    dt = 1/10
+
+    timestamps = np.linspace(0,1,100)
+    gi = list()
+    next_global_idx = 0
+
+    local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, 
+        start_time=start_time, dt=dt, next_global_idx=next_global_idx)
+    assert local_idxs[0] == 0
+    assert global_idxs[0] == 0
+    # print(local_idxs)
+    # print(global_idxs)
+    # print(timestamps[local_idxs])
+    buffer[global_idxs] = timestamps[local_idxs]
+    gi.extend(global_idxs)
+    
+    timestamps = np.linspace(0.5,1.5,100)
+    local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, 
+        start_time=start_time, dt=dt, next_global_idx = next_global_idx)
+    # print(local_idxs)
+    # print(global_idxs)
+    # print(timestamps[local_idxs])
+    # import pdb; pdb.set_trace()
+    buffer[global_idxs] = timestamps[local_idxs]
+    gi.extend(global_idxs)
+    
+    assert np.all(buffer[1:] > buffer[:-1])
+    assert np.all(np.array(gi) == np.array(list(range(len(gi)))))
+    # print(buffer)
+
+    # start over
+    next_global_idx = 0
+    timestamps = np.linspace(0,1,3)
+    local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, 
+        start_time=start_time, dt=dt, next_global_idx = next_global_idx)
+    assert local_idxs[0] == 0
+    assert local_idxs[-1] == 2
+    # print(local_idxs)
+    # print(global_idxs)
+    # print(timestamps[local_idxs])
+
+    # test numerical error issue
+    # this becomes a problem when eps <= 1e-7
+    start_time = time.time()
+    next_global_idx = 0
+    timestamps = np.arange(100000) * dt + start_time
+    local_idxs, global_idxs, next_global_idx = get_accumulate_timestamp_idxs(timestamps, 
+        start_time=start_time, dt=dt, next_global_idx = next_global_idx)
+    assert local_idxs == global_idxs
+    # print(local_idxs)
+    # print(global_idxs)
+    # print(timestamps[local_idxs])
+
+
+def test_obs_accumulator():
+    dt = 1/10
+    ddt = 1/100
+    n = 100
+    d = 6
+    start_time = time.time()
+    toa = TimestampObsAccumulator(start_time, dt)
+    poses = np.arange(n).reshape((n,1))
+    poses = np.repeat(poses, d, axis=1)
+    timestamps = np.arange(n) * ddt + start_time
+
+    toa.put({
+        'pose': poses,
+        'timestamp': timestamps
+    }, timestamps)
+    assert np.all(toa.data['pose'][:,0] == np.arange(10)*10)
+    assert len(toa) == 10
+
+    # add the same thing, result shouldn't change
+    toa.put({
+        'pose': poses,
+        'timestamp': timestamps
+    }, timestamps)
+    assert np.all(toa.data['pose'][:,0] == np.arange(10)*10)
+    assert len(toa) == 10
+    
+    # add lower than desired freuquency to test fill_in
+    dt = 1/10
+    ddt = 1/5
+    n = 10
+    d = 6
+    start_time = time.time()
+    toa = TimestampObsAccumulator(start_time, dt)
+    poses = np.arange(n).reshape((n,1))
+    poses = np.repeat(poses, d, axis=1)
+    timestamps = np.arange(n) * ddt + start_time
+
+    toa.put({
+        'pose': poses,
+        'timestamp': timestamps
+    }, timestamps)
+    assert len(toa) == 1 + (n-1) * 2
+
+    timestamps = (np.arange(n) + 2) * ddt + start_time
+    toa.put({
+        'pose': poses,
+        'timestamp': timestamps
+    }, timestamps)
+    assert len(toa) == 1 + (n-1) * 2 + 4
+
+
+def test_action_accumulator():
+    dt = 1/10
+    n = 10
+    d = 6
+    start_time = time.time()
+    taa = TimestampActionAccumulator(start_time, dt)
+    actions = np.arange(n).reshape((n,1))
+    actions = np.repeat(actions, d, axis=1)
+
+    timestamps = np.arange(n) * dt + start_time
+    taa.put(actions, timestamps)
+    assert np.all(taa.actions == actions)
+    assert np.all(taa.timestamps == timestamps)
+
+    # add another round
+    taa.put(actions-5, timestamps-0.5)
+    assert np.allclose(taa.timestamps, timestamps)
+
+    # add another round
+    taa.put(actions+5, timestamps+0.5)
+    assert len(taa) == 15
+    assert np.all(taa.actions[:,0] == np.arange(15))
+    
+
+
+if __name__ == '__main__':
+    test_action_accumulator()

third_party/diffusion_policy/train.py

@@ -0,0 +1,35 @@
+"""
+Usage:
+Training:
+python train.py --config-name=train_diffusion_lowdim_workspace
+"""
+
+import sys
+# use line-buffering for both stdout and stderr
+sys.stdout = open(sys.stdout.fileno(), mode='w', buffering=1)
+sys.stderr = open(sys.stderr.fileno(), mode='w', buffering=1)
+
+import hydra
+from omegaconf import OmegaConf
+import pathlib
+from diffusion_policy.workspace.base_workspace import BaseWorkspace
+
+# allows arbitrary python code execution in configs using the ${eval:''} resolver
+OmegaConf.register_new_resolver("eval", eval, replace=True)
+
+@hydra.main(
+    version_base=None,
+    config_path=str(pathlib.Path(__file__).parent.joinpath(
+        'diffusion_policy','config'))
+)
+def main(cfg: OmegaConf):
+    # resolve immediately so all the ${now:} resolvers
+    # will use the same time.
+    OmegaConf.resolve(cfg)
+
+    cls = hydra.utils.get_class(cfg._target_)
+    workspace: BaseWorkspace = cls(cfg)
+    workspace.run()
+
+if __name__ == "__main__":
+    main()