fix keypoint bug on different shapes

config/image_pusht_dp_tf_hgr_20pc.yaml

@@ -157,7 +157,7 @@ task:
         type: rgb
 task_name: pusht_image
 training:
-  checkpoint_every: 50
+  checkpoint_every: 20
   debug: false
   device: cuda:0
   gradient_accumulate_every: 1

diffusion_policy/env/pusht/pusht_env.py

@@ -244,6 +244,9 @@ def seed(self, seed=None):
     def _handle_collision(self, arbiter, space, data):
         self.n_contact_points += len(arbiter.contact_point_set.points)
 
+    def set_block_shape(self, shape):
+        self.block_shape = shape
+
     def _set_state(self, state):
         if isinstance(state, np.ndarray):
             state = state.tolist()

diffusion_policy/env/pusht/pusht_keypoints_env.py

@@ -4,6 +4,12 @@
 from diffusion_policy.env.pusht.pymunk_keypoint_manager import PymunkKeypointManager
 import numpy as np
 
+KP_MANAGER_DICT = {
+    "tee": PymunkKeypointManager.create_from_pusht_env(PushTEnv(), block_shape="tee"),
+    "gamma": PymunkKeypointManager.create_from_pusht_env(PushTEnv(), block_shape="gamma"),
+    "al": PymunkKeypointManager.create_from_pusht_env(PushTEnv(), block_shape="al"),
+    "vee": PymunkKeypointManager.create_from_pusht_env(PushTEnv(), block_shape="vee"),
+}
 class PushTKeypointsEnv(PushTEnv):
     def __init__(self,
             legacy=False,
@@ -63,9 +69,8 @@ def __init__(self,
         self.keypoint_visible_rate = keypoint_visible_rate
         self.agent_keypoints = agent_keypoints
         self.draw_keypoints = draw_keypoints
-        self.kp_manager = PymunkKeypointManager(
-            local_keypoint_map=local_keypoint_map,
-            color_map=color_map)
+        self.kp_manager = KP_MANAGER_DICT["tee"]
+
         self.draw_kp_map = None
 
     @classmethod
@@ -74,6 +79,10 @@ def genenerate_keypoint_manager_params(cls):
         kp_manager = PymunkKeypointManager.create_from_pusht_env(env)
         kp_kwargs = kp_manager.kwargs
         return kp_kwargs
+
+    def set_block_shape(self, shape):
+        self.block_shape = shape
+        self.kp_manager = KP_MANAGER_DICT[shape]
 
     def _get_obs(self):
         # get keypoints

diffusion_policy/env/pusht/pymunk_keypoint_manager.py

@@ -50,7 +50,7 @@ def kwargs(self):
         }
 
     @classmethod
-    def create_from_pusht_env(cls, env, n_block_kps=9, n_agent_kps=3, seed=0, **kwargs):
+    def create_from_pusht_env(cls, env, n_block_kps=9, n_agent_kps=3, seed=0, block_shape="tee", **kwargs):
         rng = np.random.default_rng(seed=seed)
         local_keypoint_map = dict()
         for name in ['block','agent']:
@@ -60,7 +60,14 @@ def create_from_pusht_env(cls, env, n_block_kps=9, n_agent_kps=3, seed=0, **kwar
                 self.agent = obj = self.add_circle((256, 400), 15)
                 n_kps = n_agent_kps
             else:
-                self.block = obj = self.add_tee((256, 300), 0)
+                if block_shape == "tee":
+                    self.block = obj = self.add_tee((256, 300), 0)
+                elif block_shape == "gamma":
+                    self.block = obj = self.add_gamma((256, 300), 0)
+                elif block_shape == "al":
+                    self.block = obj = self.add_al((256, 300), 0)
+                elif block_shape == "vee":
+                    self.block = obj = self.add_vee((256, 300), 0)
                 n_kps = n_block_kps
 
             self.screen = pygame.Surface((512,512))

diffusion_policy/env_runner/pusht_image_runner.py

@@ -122,7 +122,7 @@ def init_fn(env, seed=seed, enable_render=enable_render, shape=shape_name):
 
                     # set shape
                     assert isinstance(env.env.env, PushTImageEnv)
-                    env.env.env.block_shape = shape
+                    env.env.env.set_block_shape(shape)
 
                 env_seeds.append(seed)
                 env_prefixs.append(f'test/{shape_name}_')

diffusion_policy/env_runner/pusht_keypoints_runner.py

@@ -134,7 +134,7 @@ def init_fn(env, seed=seed, enable_render=enable_render, shape=shape_name):
 
                     # set shape
                     assert isinstance(env.env.env, PushTKeypointsEnv)
-                    env.env.env.block_shape = shape
+                    env.env.env.set_block_shape(shape)
 
                 env_seeds.append(seed)
                 env_prefixs.append(f'test/{shape_name}_')

eval.sh

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

hgr.sh

@@ -1 +1 @@
-python hgr_pusht.py -o data/hgr_pusht_1000epoch_base_20.zarr -c /local2/xingcheng/diffusion_policy/data/outputs/2024.11.25/00.19.26_train_diffusion_transformer_lowdim_pusht_lowdim/checkpoints/epoch=1000-test_mean_score=0.845.ckpt -d cuda:0
+python hgr_pusht.py -o data/hgr_pusht_1000epoch_base_20_kpright.zarr -c /local2/xingcheng/diffusion_policy/data/outputs/2024.11.25/00.19.26_train_diffusion_transformer_lowdim_pusht_lowdim/checkpoints/epoch=1000-test_mean_score=0.845.ckpt -d cuda:7

hgr_pusht.py

@@ -94,9 +94,9 @@ def main(output, checkpoint, device, num_trajs_per_shape, reward_thres):
             episode = list()
             # set seed for env
             env.seed(seed)
-            env.block_shape = shape
+            env.set_block_shape(shape)
             menv.seed(seed)
-            menv.env.block_shape = shape
+            menv.env.set_block_shape(shape)
             # reset env and get observations (including info and render for recording)
             obs = env.reset()
             info = env._get_info()
@@ -176,7 +176,7 @@ def main(output, checkpoint, device, num_trajs_per_shape, reward_thres):
 
                 # First transform the initial position of the block
                 env.seed(seed)
-                env.block_shape = shape
+                env.set_block_shape(shape)
                 obs = env.reset()
                 init_info = env._get_info()
                 block_pose = init_info['block_pose']