JC6123/aloha
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Change network architecture to have joints as ACT encoder input

Browse Source 
- does not affect performance for position control
- not backward compatible: policies trained before this commit will not load because of additional params.
This commit is contained in:
Tony Zhao
2023-04-14 11:50:27 -07:00
parent 25a02bfdd1
commit 0f54d10d7f
1 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
detr/models/detr_vae.py
View File

@@ -66,9 +66,10 @@ class DETRVAE(nn.Module):
# encoder extra parameters # encoder extra parameters
self.latent_dim = 32 # final size of latent z # TODO tune self.latent_dim = 32 # final size of latent z # TODO tune
self.cls_embed = nn.Embedding(1, hidden_dim) # extra cls token embedding 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.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 # decoder extra parameters
self.latent_out_proj = nn.Linear(self.latent_dim, hidden_dim) # project latent sample to embedding 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 ### Obtain latent z from action sequence
if is_training: if is_training:
# project action sequence to embedding dim, and concat with a CLS token # 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 = self.cls_embed.weight # (1, hidden_dim)
cls_embed = torch.unsqueeze(cls_embed, axis=0).repeat(bs, 1, 1) # (bs, 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) encoder_input = encoder_input.permute(1, 0, 2) # (seq+1, bs, hidden_dim)
# do not mask cls token # do not mask cls token
cls_is_pad = torch.full((bs, 1), False).to(qpos.device) # False: not a padding cls_joint_is_pad = torch.full((bs, 2), False).to(qpos.device) # False: not a padding
is_pad = torch.cat([cls_is_pad, is_pad], axis=1) # (bs, seq+1) is_pad = torch.cat([cls_joint_is_pad, is_pad], axis=1) # (bs, seq+1)
# obtain position embedding # obtain position embedding
pos_embed = self.pos_table.clone().detach() pos_embed = self.pos_table.clone().detach()
pos_embed = pos_embed.permute(1, 0, 2) # (seq+1, 1, hidden_dim) pos_embed = pos_embed.permute(1, 0, 2) # (seq+1, 1, hidden_dim)