Upload 2 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+pusht_policy.mp4 filter=lfs diff=lfs merge=lfs -text

pusht_policy.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de267c80a2133d2af9783649b663f0f2c29fa75e33b5fd0025eeb4473432e6d1
+size 195654

run_pusht.py

@@ -0,0 +1,102 @@
+import os
+import gymnasium as gym
+import gym_pusht
+import torch
+import imageio
+from huggingface_hub import hf_hub_download
+import safetensors.torch
+from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.policies.factory import make_pre_post_processors
+from lerobot.envs.utils import preprocess_observation
+
+def main():
+    # 1. Download checkpoint and load config
+    print("Downloading config from lerobot/diffusion_pusht...")
+    cfg = PreTrainedConfig.from_pretrained('lerobot/diffusion_pusht')
+    
+    # We override the observation.image feature shape to (3, 384, 384) to match the environment defaults,
+    # which instantiates the model's pos_grid as [144, 2] instead of [9, 2] (checkpoint size).
+    cfg.input_features['observation.image'].shape = (3, 384, 384)
+
+    # Build the DiffusionPolicy
+    print("Building DiffusionPolicy...")
+    policy = DiffusionPolicy(cfg)
+    print("Initial pos_grid shape in model:", policy.diffusion.rgb_encoder.pool.pos_grid.shape)
+
+    # Load weights with strict=False
+    print("Downloading and loading safetensors model weights...")
+    model_file = hf_hub_download(repo_id='lerobot/diffusion_pusht', filename='model.safetensors')
+    state_dict = safetensors.torch.load_file(model_file)
+    policy.load_state_dict(state_dict, strict=False)
+
+    # 2. Patch the pos_grid shape mismatch so inference works
+    print("Patching the pos_grid shape mismatch...")
+    checkpoint_pos_grid = state_dict['diffusion.rgb_encoder.pool.pos_grid']
+    policy.diffusion.rgb_encoder.pool.register_buffer('pos_grid', checkpoint_pos_grid)
+    print("Patched pos_grid shape in model:", policy.diffusion.rgb_encoder.pool.pos_grid.shape)
+
+    # Move policy to correct device and set to eval mode
+    policy.to(cfg.device)
+    policy.eval()
+
+    # 3. Create preprocessor / postprocessor with the extracted dataset stats
+    print("Creating preprocessor and postprocessor...")
+    dataset_stats = {
+        'observation.image': {
+            'mean': state_dict['normalize_inputs.buffer_observation_image.mean'],
+            'std': state_dict['normalize_inputs.buffer_observation_image.std'],
+        },
+        'observation.state': {
+            'max': state_dict['normalize_inputs.buffer_observation_state.max'],
+            'min': state_dict['normalize_inputs.buffer_observation_state.min'],
+        },
+        'action': {
+            'max': state_dict['normalize_targets.buffer_action.max'],
+            'min': state_dict['normalize_targets.buffer_action.min'],
+        }
+    }
+    preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_stats)
+
+    # 4. Instantiate the gym environment
+    print("Creating PushT environment...")
+    env = gym.make('gym_pusht/PushT-v0', render_mode='rgb_array', obs_type='pixels_agent_pos')
+    
+    # Reset env and cache initial frame
+    policy.reset()
+    obs, info = env.reset()
+    frames = [env.render()]
+
+    # Run rollout for 300 steps
+    print("Running 300 steps rollout...")
+    for step in range(300):
+        # Format observations to LeRobot format
+        obs_t = preprocess_observation(obs)
+        obs_t = preprocessor(obs_t)
+
+        # Select action
+        with torch.no_grad():
+            action = policy.select_action(obs_t)
+        action = postprocessor(action)
+
+        # Extract numpy action and apply to env (drop batch dimension)
+        action_numpy = action.to("cpu").numpy()[0]
+        obs, reward, terminated, truncated, info = env.step(action_numpy)
+        
+        # Render frame
+        frame = env.render()
+        frames.append(frame)
+
+        if terminated or truncated:
+            obs, info = env.reset()
+
+    # Close env
+    env.close()
+
+    # 5. Save the frames as pusht_policy.mp4
+    print("Saving video to pusht_policy.mp4...")
+    imageio.mimsave("pusht_policy.mp4", frames, fps=10)
+    print("Done! Video saved successfully.")
+
+if __name__ == "__main__":
+    main()