Conghoan branch

.gitattributes

@@ -4,3 +4,4 @@
 outputs/train/aloha_sim_transfer_cube_human/model_10000.safetensors filter=lfs diff=lfs merge=lfs -text
 resnet18-f37072fd.safetensors filter=lfs diff=lfs merge=lfs -text
 outputs/train/act_aloha_sim_transfer_cube_human/model_10000.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ipynb filter=lfs diff=lfs merge=lfs -text

01. Test/test.py

@@ -0,0 +1,21 @@
+# example.py
+import imageio
+import gymnasium as gym
+import numpy as np
+import gym_aloha
+
+env = gym.make("gym_aloha/AlohaInsertion-v0")
+observation, info = env.reset()
+frames = []
+
+for _ in range(1000):
+    action = env.action_space.sample()
+    observation, reward, terminated, truncated, info = env.step(action)
+    image = env.render()
+    frames.append(image)
+
+    if terminated or truncated:
+        observation, info = env.reset()
+
+env.close()
+imageio.mimsave("example.mp4", np.stack(frames), fps=25)

networks.py

@@ -6,7 +6,7 @@
 import numpy as np
 import tinygrad
 from tinygrad import Tensor, nn, dtypes
-from tinygrad.ops import Variable
+# from tinygrad.ops import Variable
 
 
 from utils import *

normalize.py

@@ -52,14 +52,28 @@ def create_stats_buffers(
             # unnormalization). See the logic here
             # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
             if mode == "mean_std":
-                buffer["mean"].assign(stats[key]["mean"].numpy())
+                mean_val = stats[key]["mean"]
+                std_val = stats[key]["std"]
+                # Convert to numpy if needed (handle both numpy arrays and torch tensors)
+                if hasattr(mean_val, 'numpy'):
+                    mean_val = mean_val.numpy()
+                if hasattr(std_val, 'numpy'):
+                    std_val = std_val.numpy()
+                buffer["mean"].assign(mean_val)
                 buffer["mean"].requires_grad = False
-                buffer["std"].assign(stats[key]["std"].numpy())
+                buffer["std"].assign(std_val)
                 buffer["std"].requires_grad = False
             elif mode == "min_max":
-                buffer["min"].assign(stats[key]["min"].numpy())
+                min_val = stats[key]["min"]
+                max_val = stats[key]["max"]
+                # Convert to numpy if needed (handle both numpy arrays and torch tensors)
+                if hasattr(min_val, 'numpy'):
+                    min_val = min_val.numpy()
+                if hasattr(max_val, 'numpy'):
+                    max_val = max_val.numpy()
+                buffer["min"].assign(min_val)
                 buffer["min"].requires_grad = False
-                buffer["max"].assign(stats[key]["max"].numpy())
+                buffer["max"].assign(max_val)
                 buffer["max"].requires_grad = False
 
         stats_buffers[key] = buffer

resnet.py

@@ -90,8 +90,10 @@ def __call__(self, x:Tensor):
 
 class ResNetInstances:
     def resnet18_IMAGENET1K_V1_Generator():
+        import pathlib
         resnet18_IMAGENET1K = ResNet(Block, [2, 2, 2, 2], num_classes=1000)
-        state_dict = nn.state.safe_load("resnet18-f37072fd.safetensors")
+        model_path = pathlib.Path(__file__).parent / "resnet18-f37072fd.safetensors"
+        state_dict = nn.state.safe_load(str(model_path))
         nn.state.load_state_dict(resnet18_IMAGENET1K, state_dict)
         return resnet18_IMAGENET1K
 

test.py

@@ -16,8 +16,11 @@
 import argparse
 
 parser=argparse.ArgumentParser(description="Argument Parser for ACT testing on simulated environments")
-parser.add_argument("env_name", type=str, choices=['AlohaTransferCube-v0', 'AlohaInsertion-v0'], default='AlohaTransferCube-v0')
-parser.add_argument("model_path", type=str)
+# parser.add_argument("--env_name", type=str, choices=['AlohaTransferCube-v0', 'AlohaInsertion-v0'], default='AlohaTransferCube-v0')
+parser.add_argument("--env_name", type=str, choices=['AlohaTransferCube-v0', 'AlohaInsertion-v0'], default='AlohaInsertion-v0')
+# parser.add_argument("--model_path", type=str, default='outputs/train/aloha_sim_transfer_cube_human/model_final.safetensors')
+# parser.add_argument("--model_path", type=str, default='outputs/train/aloha_sim_transfer_cube_human/model_final.safetensors')
+parser.add_argument("--model_path", type=str, default='outputs/train/aloha_sim_insertion_human/model_30000_original.safetensors')
 args=parser.parse_args()
 env_name = args.env_name
 
@@ -115,7 +118,7 @@ def test(state:Tensor, image:Tensor) -> Tensor:
     fps = env.metadata["render_fps"]
 
     # Encode all frames into a mp4 video.
-    video_path = output_directory / "rollout.mp4"
+    video_path = output_directory / "rollout3.mp4"
     imageio.mimsave(str(video_path), numpy.stack(frames), fps=fps)
 
     print(f"Video of the evaluation is available in '{video_path}'.")

train.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+# from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from torch.utils.data import DataLoader
 import torch
 
@@ -17,7 +19,7 @@
 import argparse
 # Start of training code
 parser=argparse.ArgumentParser(description="Argument Parser for ACT training on simulated environments")
-parser.add_argument("env_name", type=str, choices=['aloha_sim_transfer_cube_human', 'aloha_sim_insertion_human'], default='aloha_sim_insertion_human')
+parser.add_argument("--env_name", type=str, choices=['aloha_sim_transfer_cube_human', 'aloha_sim_insertion_human'], default='aloha_sim_insertion_human')
 parser.add_argument("--model_starting_point", type=str)
 parser.add_argument("--model_start_step_count", type=int)
 args=parser.parse_args()
@@ -30,17 +32,18 @@
 # Number of offline training steps (we'll only do offline training for this example.)
 # Adjust as you prefer. 5000 steps are needed to get something worth evaluating.
 training_steps = 100000
-log_freq = 1
+log_freq = 10
 
 # Set up the dataset.
 delta_timestamps = {
     "action": [i / 50.0 for i in range(100)],
 }
 dataset = LeRobotDataset(f'lerobot/{env_name}', delta_timestamps=delta_timestamps)
-print(dataset.stats)
+print(dataset.meta.stats)
+
 
 cfg = ACTConfig()
-policy = ACTPolicy(cfg, dataset_stats=dataset.stats)
+policy = ACTPolicy(cfg, dataset_stats=dataset.meta.stats)
 policy.reset()
 
 step = 0
@@ -65,7 +68,7 @@
         for stats_type, listconfig in stats_dict.items():
             # example of stats_type: min, max, mean, std
             print(f'listconfig: {listconfig}')
-            dataset.stats[key][stats_type] = torch.tensor(listconfig, dtype=torch.float32)
+            dataset.meta.stats[key][stats_type] = torch.tensor(listconfig, dtype=torch.float32)
 
 if cfg.train_backbone_separately == True:
     opt = nn.optim.AdamW(params_not_backbone, lr=1e-5, weight_decay=1e-4)
@@ -89,6 +92,19 @@ def train_step(
     if cfg.train_backbone_separately:
         opt_backbone.zero_grad()
     loss.backward()
+    ########################################################################
+    # Handle unused parameters by assigning zero gradients
+    optimizers_list = [opt]
+    if cfg.train_backbone_separately:
+        optimizers_list.append(opt_backbone)
+    for optimizer in optimizers_list:
+        for param in optimizer.params:
+            if param.grad is None:
+                # Create zero gradient with same shape and device as parameter
+                param.grad = Tensor.zeros(*param.shape, device=param.device, requires_grad=False)
+
+    ########################################################################
+
     if cfg.train_backbone_separately:
         grad_norm_not_backbone = clip_grad_norm_(params_not_backbone, 10.0)
         grad_norm_backbone = clip_grad_norm_(params_backbone, 10.0)
@@ -118,7 +134,16 @@ def train_step(
 with Tensor.train():
     while not done:
         for batch in dataloader:
-            batch = {k: Tensor(v.numpy(), requires_grad=False) for k, v in batch.items()}
+            # batch = {k: Tensor(v.numpy(), requires_grad=False) for k, v in batch.items()}
+            batch_converted = {}
+            for k, v in batch.items():
+                if isinstance(v, torch.Tensor):
+                    batch_converted[k] = Tensor(v.detach().cpu().numpy(), requires_grad=False)
+                else:
+                    batch_converted[k] = v  # Keep strings, lists, etc. as-is
+
+            batch = batch_converted
+
             batch = policy.normalize_batch_inputs_and_targets(batch)
             print(f'batch: {batch}')
             info = train_step(