From 0f54d10d7ff8153aabfb719c6cca59e41b877fb3 Mon Sep 17 00:00:00 2001 From: Tony Zhao Date: Fri, 14 Apr 2023 11:50:27 -0700 Subject: [PATCH] Change network architecture to have joints as ACT encoder input - does not affect performance for position control - not backward compatible: policies trained before this commit will not load because of additional params. --- detr/models/detr_vae.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/detr/models/detr_vae.py b/detr/models/detr_vae.py index fc25edf..42005c5 100644 --- a/detr/models/detr_vae.py +++ b/detr/models/detr_vae.py @@ -66,9 +66,10 @@ class DETRVAE(nn.Module): # encoder extra parameters self.latent_dim = 32 # final size of latent z # TODO tune self.cls_embed = nn.Embedding(1, hidden_dim) # extra cls token embedding - self.encoder_proj = nn.Linear(14, hidden_dim) # project state to embedding + self.encoder_action_proj = nn.Linear(14, hidden_dim) # project action to embedding + self.encoder_joint_proj = nn.Linear(14, hidden_dim) # project qpos to embedding self.latent_proj = nn.Linear(hidden_dim, self.latent_dim*2) # project hidden state to latent std, var - self.register_buffer('pos_table', get_sinusoid_encoding_table(num_queries+1, hidden_dim)) + self.register_buffer('pos_table', get_sinusoid_encoding_table(1+1+num_queries, hidden_dim)) # [CLS], qpos, a_seq # decoder extra parameters self.latent_out_proj = nn.Linear(self.latent_dim, hidden_dim) # project latent sample to embedding @@ -86,14 +87,16 @@ class DETRVAE(nn.Module): ### Obtain latent z from action sequence if is_training: # project action sequence to embedding dim, and concat with a CLS token - action_embed = self.encoder_proj(actions) # (bs, seq, hidden_dim) + action_embed = self.encoder_action_proj(actions) # (bs, seq, hidden_dim) + qpos_embed = self.encoder_action_proj(qpos) # (bs, hidden_dim) + qpos_embed = torch.unsqueeze(qpos_embed, axis=1) # (bs, 1, hidden_dim) cls_embed = self.cls_embed.weight # (1, hidden_dim) cls_embed = torch.unsqueeze(cls_embed, axis=0).repeat(bs, 1, 1) # (bs, 1, hidden_dim) - encoder_input = torch.cat([cls_embed, action_embed], axis=1) # (bs, seq+1, hidden_dim) + encoder_input = torch.cat([cls_embed, qpos_embed, action_embed], axis=1) # (bs, seq+1, hidden_dim) encoder_input = encoder_input.permute(1, 0, 2) # (seq+1, bs, hidden_dim) # do not mask cls token - cls_is_pad = torch.full((bs, 1), False).to(qpos.device) # False: not a padding - is_pad = torch.cat([cls_is_pad, is_pad], axis=1) # (bs, seq+1) + cls_joint_is_pad = torch.full((bs, 2), False).to(qpos.device) # False: not a padding + is_pad = torch.cat([cls_joint_is_pad, is_pad], axis=1) # (bs, seq+1) # obtain position embedding pos_embed = self.pos_table.clone().detach() pos_embed = pos_embed.permute(1, 0, 2) # (seq+1, 1, hidden_dim)