From 2aa06c8917a1bf16308cd7dec7c406742476f61e Mon Sep 17 00:00:00 2001 From: Logic Date: Sun, 15 Mar 2026 18:54:50 +0800 Subject: [PATCH] fix(pusht): stabilize DiT pusht training on current stack --- diffusion_policy/env_runner/pusht_keypoints_runner.py | 5 ++--- requirements-pusht-5090.txt | 1 + 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/diffusion_policy/env_runner/pusht_keypoints_runner.py b/diffusion_policy/env_runner/pusht_keypoints_runner.py index a16bd58..14bbad6 100644 --- a/diffusion_policy/env_runner/pusht_keypoints_runner.py +++ b/diffusion_policy/env_runner/pusht_keypoints_runner.py @@ -8,8 +8,7 @@ import dill import math import wandb.sdk.data_types.video as wv from diffusion_policy.env.pusht.pusht_keypoints_env import PushTKeypointsEnv -from diffusion_policy.gym_util.async_vector_env import AsyncVectorEnv -# from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv +from diffusion_policy.gym_util.sync_vector_env import SyncVectorEnv from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper from diffusion_policy.gym_util.video_recording_wrapper import VideoRecordingWrapper, VideoRecorder @@ -133,7 +132,7 @@ class PushTKeypointsRunner(BaseLowdimRunner): env_prefixs.append('test/') env_init_fn_dills.append(dill.dumps(init_fn)) - env = AsyncVectorEnv(env_fns) + env = SyncVectorEnv(env_fns) # test env # env.reset(seed=env_seeds) diff --git a/requirements-pusht-5090.txt b/requirements-pusht-5090.txt index 9bf8f9d..caffc6c 100644 --- a/requirements-pusht-5090.txt +++ b/requirements-pusht-5090.txt @@ -22,6 +22,7 @@ pymunk==6.2.1 wandb==0.13.3 threadpoolctl==3.1.0 shapely==1.8.5.post1 +matplotlib==3.6.1 imageio==2.22.0 imageio-ffmpeg==0.4.7 termcolor==2.0.1