Add AbstractEnv, Refactor AlohaEnv, Add rendering_hook in env, Minor modifications, (TODO: Refactor Pusht and Simxarm)

lerobot/common/policies/act/detr_vae.py

@@ -27,7 +27,9 @@ def get_sinusoid_encoding_table(n_position, d_hid):
 class DETRVAE(nn.Module):
     """This is the DETR module that performs object detection"""
 
-    def __init__(self, backbones, transformer, encoder, state_dim, action_dim, num_queries, camera_names):
+    def __init__(
+        self, backbones, transformer, encoder, state_dim, action_dim, num_queries, camera_names, vae
+    ):
         """Initializes the model.
         Parameters:
             backbones: torch module of the backbone to be used. See backbone.py
@@ -42,6 +44,7 @@ class DETRVAE(nn.Module):
         self.camera_names = camera_names
         self.transformer = transformer
         self.encoder = encoder
+        self.vae = vae
         hidden_dim = transformer.d_model
         self.action_head = nn.Linear(hidden_dim, action_dim)
         self.is_pad_head = nn.Linear(hidden_dim, 1)
@@ -86,7 +89,7 @@ class DETRVAE(nn.Module):
         is_training = actions is not None  # train or val
         bs, _ = qpos.shape
         ### Obtain latent z from action sequence
-        if is_training:
+        if self.vae and is_training:
             # project action sequence to embedding dim, and concat with a CLS token
             action_embed = self.encoder_action_proj(actions)  # (bs, seq, hidden_dim)
             qpos_embed = self.encoder_joint_proj(qpos)  # (bs, hidden_dim)
@@ -200,6 +203,7 @@ def build(args):
         action_dim=args.action_dim,
         num_queries=args.num_queries,
         camera_names=args.camera_names,
+        vae=args.vae,
     )
 
     n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)

lerobot/common/policies/act/policy.py

@@ -11,7 +11,6 @@ from lerobot.common.policies.act.detr_vae import build
 
 def build_act_model_and_optimizer(cfg):
     model = build(cfg)
-    model.cuda()
 
     param_dicts = [
         {"params": [p for n, p in model.named_parameters() if "backbone" not in n and p.requires_grad]},
@@ -51,6 +50,8 @@ class ActionChunkingTransformerPolicy(nn.Module):
         self.kl_weight = self.cfg.kl_weight
         logging.info(f"KL Weight {self.kl_weight}")
 
+        self.to(self.device)
+
     def update(self, replay_buffer, step):
         del step
 
@@ -192,20 +193,25 @@ class ActionChunkingTransformerPolicy(nn.Module):
         normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         image = normalize(image)
 
-        is_train_mode = actions is not None
-        if is_train_mode:  # training time
+        is_training = actions is not None
+        if is_training:  # training time
             actions = actions[:, : self.model.num_queries]
             if is_pad is not None:
                 is_pad = is_pad[:, : self.model.num_queries]
 
             a_hat, is_pad_hat, (mu, logvar) = self.model(qpos, image, env_state, actions, is_pad)
-            total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar)
-            loss_dict = {}
+
             all_l1 = F.l1_loss(actions, a_hat, reduction="none")
             l1 = all_l1.mean() if is_pad is None else (all_l1 * ~is_pad.unsqueeze(-1)).mean()
+
+            loss_dict = {}
             loss_dict["l1"] = l1
-            loss_dict["kl"] = total_kld[0]
-            loss_dict["loss"] = loss_dict["l1"] + loss_dict["kl"] * self.kl_weight
+            if self.cfg.vae:
+                total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar)
+                loss_dict["kl"] = total_kld[0]
+                loss_dict["loss"] = loss_dict["l1"] + loss_dict["kl"] * self.kl_weight
+            else:
+                loss_dict["loss"] = loss_dict["l1"]
             return loss_dict
         else:
             action, _, (_, _) = self.model(qpos, image, env_state)  # no action, sample from prior