backup wip

lerobot/common/datasets/pusht.py

@@ -8,8 +8,10 @@ import torch
 import tqdm
 from gym_pusht.envs.pusht import pymunk_to_shapely
 
+from lerobot.common.datasets._diffusion_policy_replay_buffer import (
+    ReplayBuffer as DiffusionPolicyReplayBuffer,
+)
 from lerobot.common.datasets.utils import download_and_extract_zip, load_data_with_delta_timestamps
-from lerobot.common.policies.diffusion.replay_buffer import ReplayBuffer as DiffusionPolicyReplayBuffer
 
 # as define in env
 SUCCESS_THRESHOLD = 0.95  # 95% coverage,