suolyerandyangping
d18fa65fa1
* add mot * update libero example * translate zh to en * fix load model from hf * lint * lint --------- Co-authored-by: yangping <yangping@x2robot.com>
809 lines
30 KiB
Python
809 lines
30 KiB
Python
import torch
|
|
|
|
import torch.nn as nn
|
|
|
|
from typing import Union
|
|
import math
|
|
|
|
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
|
|
from torch.distributions import Beta
|
|
|
|
|
|
def print_rank_last(message):
|
|
"""If distributed is initialized, print only on last rank."""
|
|
if torch.distributed.is_initialized():
|
|
if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1):
|
|
print(message, flush=True)
|
|
else:
|
|
print(message, flush=True)
|
|
|
|
|
|
class Normalizer(nn.Module):
|
|
@classmethod
|
|
def from_ckpt(cls, ckpt_path):
|
|
instance = cls.__new__(cls)
|
|
nn.Module.__init__(instance)
|
|
|
|
instance.min = nn.ParameterDict()
|
|
instance.delta = nn.ParameterDict()
|
|
instance.min_key = "min"
|
|
instance.delta_key = "delta"
|
|
|
|
ckpt = torch.load(ckpt_path, map_location="cpu")
|
|
|
|
for key, value in ckpt.items():
|
|
# Parse key: "min.robot_name" -> prefix="min", name="robot_name"
|
|
try:
|
|
prefix, name = key.split(".", 1)
|
|
if hasattr(instance, prefix):
|
|
getattr(instance, prefix)[name] = nn.Parameter(
|
|
value, requires_grad=False
|
|
)
|
|
print("prefix", prefix)
|
|
print("name", name)
|
|
except ValueError:
|
|
continue
|
|
|
|
return instance
|
|
|
|
def __init__(
|
|
self, action_statistic_dof, dof_config, min_key="min", delta_key="delta"
|
|
):
|
|
super(Normalizer, self).__init__()
|
|
|
|
self.min_key = min_key
|
|
self.delta_key = delta_key
|
|
|
|
action_statistic = {}
|
|
for robot_name in action_statistic_dof.keys():
|
|
action_statistic[robot_name] = {}
|
|
all_dof_min = []
|
|
all_dof_delta = []
|
|
for k in dof_config:
|
|
if k in action_statistic_dof[robot_name]:
|
|
if (
|
|
min_key in action_statistic_dof[robot_name][k]
|
|
and delta_key in action_statistic_dof[robot_name][k]
|
|
):
|
|
all_dof_min.extend(action_statistic_dof[robot_name][k][min_key])
|
|
all_dof_delta.extend(
|
|
action_statistic_dof[robot_name][k][delta_key]
|
|
)
|
|
else:
|
|
if robot_name == "x2_normal" or "libero" in robot_name:
|
|
print_rank_last(
|
|
f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} "
|
|
)
|
|
print_rank_last(
|
|
f"not in action_statistic_dof[{robot_name}][{k}], use default min 0.0 and delta 1.0"
|
|
)
|
|
all_dof_min.extend([0.0] * dof_config[k])
|
|
all_dof_delta.extend([1.0] * dof_config[k])
|
|
else:
|
|
if robot_name == "x2_normal" or "libero" in robot_name:
|
|
print_rank_last(
|
|
f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0"
|
|
)
|
|
all_dof_min.extend([0.0] * dof_config[k])
|
|
all_dof_delta.extend([1.0] * dof_config[k])
|
|
all_dof_min = torch.tensor(all_dof_min)
|
|
all_dof_delta = torch.tensor(all_dof_delta)
|
|
action_statistic[robot_name][min_key] = all_dof_min
|
|
action_statistic[robot_name][delta_key] = all_dof_delta
|
|
|
|
self.min = nn.ParameterDict(
|
|
{
|
|
k: nn.Parameter(action_statistic[k][min_key], requires_grad=False)
|
|
for k in action_statistic.keys()
|
|
}
|
|
)
|
|
self.delta = nn.ParameterDict(
|
|
{
|
|
k: nn.Parameter(action_statistic[k][delta_key], requires_grad=False)
|
|
for k in action_statistic.keys()
|
|
}
|
|
)
|
|
|
|
for k, v in action_statistic.items():
|
|
print_rank_last(
|
|
f"Normalizer: {k} min {action_statistic[k][min_key]} delta {action_statistic[k][delta_key]}"
|
|
)
|
|
|
|
def normalize_data(self, xs, dataset_names):
|
|
new_xs = []
|
|
dataset_names = [name for name in dataset_names if name != "x2_multimodal"]
|
|
for x, dataset_name in zip(xs, dataset_names):
|
|
x = (x - self.min[dataset_name]) / (self.delta[dataset_name])
|
|
x = x * 2 - 1
|
|
x = torch.clamp(x, -1, 1)
|
|
new_xs.append(x)
|
|
new_xs = torch.stack(new_xs)
|
|
return new_xs
|
|
|
|
def unnormalize_data(self, xs, dataset_names, dof_mask=None):
|
|
new_xs = []
|
|
dataset_names = [name for name in dataset_names if name != "x2_multimodal"]
|
|
dof_mask = dof_mask if dof_mask is not None else [None] * len(xs)
|
|
for x, dataset_name, mask in zip(xs, dataset_names, dof_mask):
|
|
x = (x + 1) / 2
|
|
if mask is not None:
|
|
mask = mask[0].bool()
|
|
action_space_delta = self.delta[dataset_name][mask]
|
|
action_space_min = self.min[dataset_name][mask]
|
|
else:
|
|
action_space_delta = self.delta[dataset_name]
|
|
action_space_min = self.min[dataset_name]
|
|
x = x * action_space_delta + action_space_min
|
|
new_xs.append(x)
|
|
new_xs = torch.stack(new_xs)
|
|
return new_xs
|
|
|
|
|
|
class SinusoidalPosEmb(nn.Module):
|
|
def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0):
|
|
super().__init__()
|
|
self.dim = dim
|
|
if dim % 2 != 0:
|
|
raise ValueError(f"embedding_dim ({dim}) must be divisible by 2")
|
|
self.min_period = min_period
|
|
self.max_period = max_period
|
|
|
|
def forward(self, x):
|
|
device = x.device
|
|
half_dim = self.dim // 2
|
|
emb = math.log(10000) / (half_dim - 1)
|
|
emb = torch.exp(
|
|
torch.arange(half_dim, device=device, dtype=torch.float32) * -emb
|
|
)
|
|
emb = x[:, None] * emb[None, :]
|
|
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
|
return emb
|
|
|
|
|
|
class Downsample1d(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.conv = nn.Conv1d(dim, dim, 3, 2, 1)
|
|
|
|
def forward(self, x):
|
|
return self.conv(x)
|
|
|
|
|
|
class Upsample1d(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1)
|
|
|
|
def forward(self, x):
|
|
return self.conv(x)
|
|
|
|
|
|
class Conv1dBlock(nn.Module):
|
|
"""
|
|
Conv1d --> GroupNorm --> Mish
|
|
"""
|
|
|
|
def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):
|
|
super().__init__()
|
|
|
|
self.block = nn.Sequential(
|
|
nn.Conv1d(
|
|
inp_channels, out_channels, kernel_size, padding=kernel_size // 2
|
|
),
|
|
nn.GroupNorm(n_groups, out_channels),
|
|
nn.Mish(),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.block(x)
|
|
|
|
|
|
class ConditionalResidualBlock1D(nn.Module):
|
|
def __init__(self, in_channels, out_channels, cond_dim, kernel_size=3, n_groups=8):
|
|
super().__init__()
|
|
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups),
|
|
Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups),
|
|
]
|
|
)
|
|
|
|
# FiLM modulation https://arxiv.org/abs/1709.07871
|
|
# predicts per-channel scale and bias
|
|
cond_channels = out_channels * 2
|
|
self.out_channels = out_channels
|
|
self.cond_encoder = nn.Sequential(
|
|
nn.Mish(),
|
|
nn.Linear(cond_dim, cond_channels),
|
|
nn.Dropout(0.1),
|
|
nn.Unflatten(-1, (-1, 1)),
|
|
)
|
|
|
|
# make sure dimensions compatible
|
|
self.residual_conv = (
|
|
nn.Conv1d(in_channels, out_channels, 1)
|
|
if in_channels != out_channels
|
|
else nn.Identity()
|
|
)
|
|
|
|
def forward(self, x, cond):
|
|
"""
|
|
x : [ batch_size x in_channels x horizon ]
|
|
cond : [ batch_size x cond_dim]
|
|
|
|
returns:
|
|
out : [ batch_size x out_channels x horizon ]
|
|
"""
|
|
out = self.blocks[0](x)
|
|
embed = self.cond_encoder(cond)
|
|
|
|
embed = embed.reshape(embed.shape[0], 2, self.out_channels, 1)
|
|
scale = embed[:, 0, ...]
|
|
bias = embed[:, 1, ...]
|
|
out = scale * out + bias
|
|
|
|
out = self.blocks[1](out)
|
|
out = out + self.residual_conv(x)
|
|
return out
|
|
|
|
|
|
class ConditionalUnet1D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_dim,
|
|
global_cond_dim,
|
|
diffusion_step_embed_dim=256,
|
|
down_dims=[256, 512, 1024],
|
|
# down_dims=[512, 1024, 2048],
|
|
kernel_size=5,
|
|
n_groups=8,
|
|
):
|
|
"""
|
|
input_dim: Dim of actions.
|
|
global_cond_dim: Dim of global conditioning applied with FiLM
|
|
in addition to diffusion step embedding. This is usually obs_horizon * obs_dim
|
|
diffusion_step_embed_dim: Size of positional encoding for diffusion iteration k
|
|
down_dims: Channel size for each UNet level.
|
|
The length of this array determines numebr of levels.
|
|
kernel_size: Conv kernel size
|
|
n_groups: Number of groups for GroupNorm
|
|
"""
|
|
|
|
super().__init__()
|
|
all_dims = [input_dim] + list(down_dims)
|
|
start_dim = down_dims[0]
|
|
|
|
dsed = diffusion_step_embed_dim
|
|
diffusion_step_encoder = nn.Sequential(
|
|
SinusoidalPosEmb(dsed),
|
|
nn.Linear(dsed, dsed * 4),
|
|
nn.Mish(),
|
|
nn.Linear(dsed * 4, dsed),
|
|
)
|
|
cond_dim = dsed + global_cond_dim
|
|
|
|
in_out = list(zip(all_dims[:-1], all_dims[1:]))
|
|
mid_dim = all_dims[-1]
|
|
self.mid_modules = nn.ModuleList(
|
|
[
|
|
ConditionalResidualBlock1D(
|
|
mid_dim,
|
|
mid_dim,
|
|
cond_dim=cond_dim,
|
|
kernel_size=kernel_size,
|
|
n_groups=n_groups,
|
|
),
|
|
ConditionalResidualBlock1D(
|
|
mid_dim,
|
|
mid_dim,
|
|
cond_dim=cond_dim,
|
|
kernel_size=kernel_size,
|
|
n_groups=n_groups,
|
|
),
|
|
]
|
|
)
|
|
|
|
down_modules = nn.ModuleList([])
|
|
for ind, (dim_in, dim_out) in enumerate(in_out):
|
|
is_last = ind >= (len(in_out) - 1)
|
|
down_modules.append(
|
|
nn.ModuleList(
|
|
[
|
|
ConditionalResidualBlock1D(
|
|
dim_in,
|
|
dim_out,
|
|
cond_dim=cond_dim,
|
|
kernel_size=kernel_size,
|
|
n_groups=n_groups,
|
|
),
|
|
ConditionalResidualBlock1D(
|
|
dim_out,
|
|
dim_out,
|
|
cond_dim=cond_dim,
|
|
kernel_size=kernel_size,
|
|
n_groups=n_groups,
|
|
),
|
|
Downsample1d(dim_out) if not is_last else nn.Identity(),
|
|
]
|
|
)
|
|
)
|
|
|
|
up_modules = nn.ModuleList([])
|
|
for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
|
|
is_last = ind >= (len(in_out) - 1)
|
|
up_modules.append(
|
|
nn.ModuleList(
|
|
[
|
|
ConditionalResidualBlock1D(
|
|
dim_out * 2,
|
|
dim_in,
|
|
cond_dim=cond_dim,
|
|
kernel_size=kernel_size,
|
|
n_groups=n_groups,
|
|
),
|
|
ConditionalResidualBlock1D(
|
|
dim_in,
|
|
dim_in,
|
|
cond_dim=cond_dim,
|
|
kernel_size=kernel_size,
|
|
n_groups=n_groups,
|
|
),
|
|
Upsample1d(dim_in) if not is_last else nn.Identity(),
|
|
]
|
|
)
|
|
)
|
|
|
|
final_conv = nn.Sequential(
|
|
Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size),
|
|
nn.Conv1d(start_dim, input_dim, 1),
|
|
)
|
|
|
|
self.diffusion_step_encoder = diffusion_step_encoder
|
|
self.up_modules = up_modules
|
|
self.down_modules = down_modules
|
|
self.final_conv = final_conv
|
|
|
|
# print("number of parameters: {:e}".format(
|
|
# sum(p.numel() for p in self.parameters()))
|
|
# )
|
|
|
|
def forward(
|
|
self,
|
|
sample: torch.Tensor,
|
|
timestep: Union[torch.Tensor, float, int],
|
|
global_cond=None,
|
|
):
|
|
"""
|
|
x: (B,T,input_dim)
|
|
timestep: (B,) or int, diffusion step
|
|
global_cond: (B,global_cond_dim)
|
|
output: (B,T,input_dim)
|
|
"""
|
|
# (B,T,C)
|
|
sample = sample.moveaxis(-1, -2)
|
|
# (B,C,T)
|
|
|
|
# 1. time
|
|
timesteps = timestep
|
|
if not torch.is_tensor(timesteps):
|
|
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
|
timesteps = torch.tensor(
|
|
[timesteps], dtype=torch.long, device=sample.device
|
|
)
|
|
elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
|
|
timesteps = timesteps[None].to(sample.device)
|
|
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
|
timesteps = timesteps.expand(sample.shape[0])
|
|
|
|
global_feature = self.diffusion_step_encoder(timesteps)
|
|
|
|
if global_cond is not None:
|
|
global_feature = torch.cat([global_feature, global_cond], axis=-1)
|
|
x = sample
|
|
h = []
|
|
for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules):
|
|
x = resnet(x, global_feature)
|
|
x = resnet2(x, global_feature)
|
|
h.append(x)
|
|
x = downsample(x) # bs, 2048, 5
|
|
|
|
for mid_module in self.mid_modules:
|
|
x = mid_module(x, global_feature) # bs, 2048, 5
|
|
|
|
for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules):
|
|
x = torch.cat((x, h.pop()), dim=1)
|
|
x = resnet(x, global_feature)
|
|
x = resnet2(x, global_feature)
|
|
x = upsample(x) # bs, 512, 20
|
|
|
|
x = self.final_conv(x) # bs, 14, 20
|
|
|
|
# (B,C,T)
|
|
x = x.moveaxis(-1, -2)
|
|
# (B,T,C)
|
|
return x
|
|
|
|
|
|
class DP_Action_head(nn.Module):
|
|
def __init__(
|
|
self,
|
|
action_dim=14,
|
|
transformer_dim=896,
|
|
global_cond_dim=1806,
|
|
load_pretrained=True,
|
|
):
|
|
super().__init__()
|
|
self.action_dim = action_dim
|
|
self.transformer_dim = transformer_dim
|
|
self.load_pretrained = load_pretrained
|
|
self.noise_scheduler = DDPMScheduler(
|
|
num_train_timesteps=132,
|
|
beta_schedule="squaredcos_cap_v2",
|
|
clip_sample=True,
|
|
prediction_type="epsilon",
|
|
)
|
|
if self.load_pretrained:
|
|
self.global_cond_dim = global_cond_dim
|
|
self.condition_proj = nn.Sequential(
|
|
nn.Linear(self.transformer_dim, 2 * self.transformer_dim),
|
|
nn.ReLU(),
|
|
nn.Linear(2 * self.transformer_dim, 2 * self.global_cond_dim),
|
|
nn.ReLU(),
|
|
nn.Linear(2 * self.global_cond_dim, self.global_cond_dim),
|
|
)
|
|
|
|
self.noise_pred_net = ConditionalUnet1D(
|
|
input_dim=self.action_dim,
|
|
# down_dims=[256,512,1024],
|
|
down_dims=[512, 1024, 2048],
|
|
global_cond_dim=self.global_cond_dim,
|
|
)
|
|
|
|
# load pretrained model
|
|
action_pretrained_path = "/x2robot/liangyuxin/workspace/DiffusionPolicy/big_mix_0718_mn/30_noise_pred_net.pth"
|
|
print("load noise_pred_net from:", action_pretrained_path, flush=True)
|
|
self.noise_pred_net.load_state_dict(torch.load(action_pretrained_path))
|
|
else:
|
|
self.noise_pred_net = ConditionalUnet1D(
|
|
input_dim=self.action_dim,
|
|
down_dims=[256, 512, 1024],
|
|
global_cond_dim=self.transformer_dim,
|
|
)
|
|
|
|
def forward(self, naction, condition, sample_times):
|
|
bs = naction.shape[0]
|
|
noise_shape = (
|
|
naction.shape[0] * sample_times,
|
|
naction.shape[1],
|
|
naction.shape[2],
|
|
)
|
|
noise = torch.randn(noise_shape, device=naction.device)
|
|
naction = (
|
|
naction.unsqueeze(1)
|
|
.repeat(1, sample_times, 1, 1)
|
|
.reshape(bs * sample_times, naction.shape[1], naction.shape[2])
|
|
)
|
|
|
|
timesteps = torch.randint(
|
|
0,
|
|
self.noise_scheduler.config.num_train_timesteps,
|
|
(bs * sample_times,),
|
|
device=naction.device,
|
|
).long()
|
|
condition = condition.to(self.condition_proj[0].weight.data.dtype)
|
|
if self.load_pretrained:
|
|
condition = self.condition_proj(condition)
|
|
|
|
noisy_actions = self.noise_scheduler.add_noise(naction, noise, timesteps)
|
|
noise_pred = self.noise_pred_net(
|
|
noisy_actions, timesteps, global_cond=condition
|
|
)
|
|
return noise, noise_pred
|
|
|
|
@torch.no_grad()
|
|
def predict(self, condition, naction=None):
|
|
bs = condition.shape[0]
|
|
condition = condition.to(self.condition_proj[0].weight.data.dtype)
|
|
if self.load_pretrained:
|
|
condition = self.condition_proj(condition)
|
|
|
|
if naction is not None:
|
|
noise_shape = (naction.shape[0], naction.shape[1], naction.shape[2])
|
|
else:
|
|
noise_shape = (bs, 16, self.action_dim) # tobe parameterized
|
|
noise = torch.randn(noise_shape, device=condition.device)
|
|
naction_pred = noise
|
|
# init scheduler
|
|
self.noise_scheduler.set_timesteps(
|
|
self.noise_scheduler.config.num_train_timesteps
|
|
)
|
|
|
|
for k in self.noise_scheduler.timesteps:
|
|
# predict noise
|
|
noise_pred = self.noise_pred_net(
|
|
sample=naction_pred, timestep=k, global_cond=condition
|
|
)
|
|
|
|
# inverse diffusion step (remove noise)
|
|
naction_pred = self.noise_scheduler.step(
|
|
model_output=noise_pred, timestep=k, sample=naction_pred
|
|
).prev_sample
|
|
|
|
return naction, naction_pred
|
|
|
|
|
|
class ActionProcessor(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.dof_config = config.dof_config
|
|
self.agent_pos_config = config.agent_pos_config
|
|
self.action_dim = sum([v for k, v in self.dof_config.items()])
|
|
self.propri_dim = sum([v for k, v in self.agent_pos_config.items()])
|
|
|
|
print_rank_last(
|
|
f"self.dof_config: {self.dof_config}; action_dim: {self.action_dim}; self.agent_pos_config: {self.agent_pos_config}; propri_dim: {self.propri_dim}"
|
|
)
|
|
|
|
self.action_hidden_size = config.action_hidden_size
|
|
self.state_hidden_size = config.state_hidden_size
|
|
self.hidden_size = config.hidden_size
|
|
|
|
if not self.config.use_state_string_representation:
|
|
if self.config.proj_with_mask:
|
|
self.propri_proj = nn.Linear(
|
|
self.propri_dim * 2, self.state_hidden_size, bias=False
|
|
)
|
|
else:
|
|
self.propri_proj = nn.Linear(
|
|
self.propri_dim, self.state_hidden_size, bias=False
|
|
)
|
|
|
|
# noise scheduler configing
|
|
if getattr(self.config, "use_flow_action_expert", True):
|
|
noise_scheduler_config = config.noise_scheduler
|
|
self.beta_alpha = noise_scheduler_config.get("beta_alpha", 1.5)
|
|
self.beta_beta = noise_scheduler_config.get("beta_beta", 1.0)
|
|
self.s = noise_scheduler_config.get("s", 0.999)
|
|
alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda")
|
|
beta_tensor = torch.tensor(self.beta_beta, dtype=torch.float32).to("cuda")
|
|
self.beta_dist = Beta(alpha_tensor, beta_tensor)
|
|
self.time_embed = SinusoidalPosEmb(self.action_hidden_size)
|
|
|
|
# project to hidden space
|
|
if self.config.proj_with_mask:
|
|
self.w1 = nn.Linear(
|
|
self.action_dim * 2, self.action_hidden_size, bias=False
|
|
)
|
|
else:
|
|
self.w1 = nn.Linear(
|
|
self.action_dim, self.action_hidden_size, bias=False
|
|
)
|
|
if not self.config.use_adarms:
|
|
self.w2 = nn.Linear(
|
|
self.action_hidden_size * 2, self.action_hidden_size, bias=False
|
|
)
|
|
self.w3 = nn.Linear(
|
|
self.action_hidden_size, self.action_hidden_size, bias=False
|
|
)
|
|
self.act_fn = nn.SiLU()
|
|
else:
|
|
self.time_mlp_in = nn.Linear(
|
|
self.action_hidden_size, self.action_hidden_size
|
|
)
|
|
self.time_mlp_out = nn.Linear(
|
|
self.action_hidden_size, self.action_hidden_size
|
|
)
|
|
self.act_fn = nn.SiLU()
|
|
|
|
# project back to action space
|
|
self.action_proj_back = nn.Linear(
|
|
self.action_hidden_size, self.action_dim, bias=False
|
|
)
|
|
self.mse_loss = nn.MSELoss(reduction="none")
|
|
|
|
def set_normalizer(self, normalizer_action, normalizer_propri):
|
|
self.normalizer_action = normalizer_action
|
|
self.normalizer_propri = normalizer_propri
|
|
|
|
# dataset_name = self.config["data"]["lerobot_config"]["repo_id"]
|
|
# print("normalizer_propri min", self.normalizer_propri.min.__getattr__(dataset_name), flush=True)
|
|
# print("normalizer_propri delta", self.normalizer_propri.delta.__getattr__(dataset_name), flush=True)
|
|
# print("normalizer_action min", self.normalizer_action.min.__getattr__(dataset_name), flush=True)
|
|
# print("normalizer_action delta", self.normalizer_action.delta.__getattr__(dataset_name), flush=True)
|
|
|
|
def sample_time(self, batch_size, device, dtype):
|
|
"""
|
|
Sampling Time Step
|
|
Generates random numbers in the range [0, 1] using a Beta distribution, and then scales them.
|
|
|
|
Parameters:
|
|
batch_size (int): Batch size
|
|
device: Device type
|
|
dtype: Data type
|
|
|
|
Returns:
|
|
torch.Tensor: Sampled time steps, with shape [batch_size]
|
|
"""
|
|
sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype)
|
|
time = (1 - sample) * self.s
|
|
return time
|
|
|
|
def proprioception_proj(
|
|
self, proprioception, dataset_names=None, dof_mask=None, use_history=False
|
|
):
|
|
"""
|
|
proprioception: [batch_size, 1, action_dim]
|
|
dataset_names: [batch_size]
|
|
dof_mask: [batch_size, action_dim]
|
|
"""
|
|
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
|
|
dtype=self.propri_proj.weight.dtype
|
|
)
|
|
if dof_mask is not None:
|
|
if self.config.proj_with_mask:
|
|
proprioception = torch.cat(
|
|
[proprioception, dof_mask], dim=-1
|
|
) # .unsqueeze(1)
|
|
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
|
|
dtype=self.propri_proj.weight.dtype
|
|
)
|
|
proprio_embed = self.propri_proj(
|
|
proprioception
|
|
) # [batch_size, 1, state_hidden_size]
|
|
|
|
if self.state_hidden_size < self.hidden_size:
|
|
# padding to hidden size
|
|
padding_size = self.hidden_size - self.state_hidden_size
|
|
padding = torch.zeros(
|
|
(proprio_embed.shape[0], 1, padding_size),
|
|
device=proprio_embed.device,
|
|
dtype=proprio_embed.dtype,
|
|
)
|
|
proprio_embed = torch.cat([proprio_embed, padding], dim=-1)
|
|
|
|
return proprio_embed # [batch_size, 1, hidden_size]
|
|
|
|
def forward(self, action_chunk, dataset_names, dof_mask=None):
|
|
"""
|
|
Parameters:
|
|
action_chunk (torch.Tensor): Action sequence, shape [batch_size, action_chunk_len, action_dim]
|
|
dataset_names: [batch_size]
|
|
dof_mask: [batch_size, action_dim]
|
|
|
|
Returns:
|
|
torch.Tensor: Processed action representation, shape [batch_size, seq_len, hidden_size]
|
|
"""
|
|
with torch.autocast("cuda", dtype=torch.float32):
|
|
action_chunk = action_chunk.to(dtype=torch.float32)
|
|
batch_size = action_chunk.shape[0]
|
|
device = action_chunk.device
|
|
dtype = action_chunk.dtype
|
|
|
|
# 1. add noise to action_chunk
|
|
noise = torch.randn_like(action_chunk)
|
|
time = self.sample_time(batch_size, device, dtype)
|
|
time_expanded = time.unsqueeze(-1).unsqueeze(-1)
|
|
noisy_action = (1 - time_expanded) * noise + time_expanded * action_chunk
|
|
flow = action_chunk - noise
|
|
|
|
# 2. sinusoidal positional encoding for timesteps
|
|
time_embed = self.time_embed(time).to(torch.float32)
|
|
|
|
self.noise = noise
|
|
self.noisy_action = noisy_action # for new x-pred
|
|
|
|
# 3.action_chunk_nosiy + t_pos_emb -> MLP_act_chunk -> action_chunk_nosiy_emb_with_t (dim=trans * chunk)
|
|
if dof_mask is not None:
|
|
noisy_action = torch.cat([noisy_action, dof_mask], dim=-1)
|
|
|
|
noisy_action = noisy_action.to(dtype=self.w1.weight.dtype)
|
|
action_embed = self.w1(noisy_action)
|
|
|
|
self.time_expanded = time_expanded # for new x-pred
|
|
|
|
if not self.config.use_adarms:
|
|
time_embed = (
|
|
time_embed.unsqueeze(1)
|
|
.repeat(1, action_embed.shape[1], 1)
|
|
.to(dtype=self.w2.weight.dtype)
|
|
)
|
|
concat_embed = torch.cat([action_embed, time_embed], dim=-1)
|
|
concat_embed = self.w2(concat_embed)
|
|
action_time_embed = self.w3(self.act_fn(concat_embed))
|
|
adarms_cond = None
|
|
else:
|
|
time_embed = self.time_mlp_in(time_embed)
|
|
time_embed = self.act_fn(time_embed)
|
|
time_embed = self.time_mlp_out(time_embed)
|
|
time_embed = self.act_fn(time_embed)
|
|
action_time_embed = action_embed
|
|
adarms_cond = time_embed
|
|
|
|
if self.action_hidden_size < self.hidden_size:
|
|
# padding to hidden size
|
|
padding_size = self.hidden_size - self.action_hidden_size
|
|
padding = torch.zeros(
|
|
(
|
|
action_time_embed.shape[0],
|
|
action_time_embed.shape[1],
|
|
padding_size,
|
|
),
|
|
device=action_time_embed.device,
|
|
dtype=action_time_embed.dtype,
|
|
)
|
|
action_time_embed = torch.cat([action_time_embed, padding], dim=-1)
|
|
|
|
return action_time_embed, flow, adarms_cond
|
|
|
|
def step(self, timestep, noisy_action, dof_mask=None):
|
|
# noisy_action: bs, pred_horizon, action_dim
|
|
# timestep: bs
|
|
with torch.autocast("cuda", dtype=torch.float32):
|
|
if dof_mask is not None and self.config.proj_with_mask:
|
|
if dof_mask.shape[1] == 1:
|
|
dof_mask = dof_mask.unsqueeze(1).repeat(1, noisy_action.shape[1], 1)
|
|
noisy_action = torch.cat([noisy_action, dof_mask], dim=-1)
|
|
|
|
noisy_action = noisy_action.to(dtype=self.w1.weight.dtype)
|
|
time_embed = self.time_embed(timestep).to(torch.float32) # bs,hidden_size
|
|
action_embed = self.w1(noisy_action)
|
|
|
|
if not self.config.use_adarms:
|
|
time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1)
|
|
time_embed = time_embed.to(device=noisy_action.device).to(
|
|
dtype=noisy_action.dtype
|
|
)
|
|
concat_embed = torch.cat([action_embed, time_embed], dim=-1)
|
|
concat_embed = self.w2(concat_embed)
|
|
embed = self.w3(self.act_fn(concat_embed)) # is this right?
|
|
adarms_cond = None
|
|
else:
|
|
time_embed = time_embed.to(dtype=self.time_mlp_in.weight.dtype)
|
|
time_embed = self.time_mlp_in(time_embed)
|
|
time_embed = self.act_fn(time_embed)
|
|
time_embed = self.time_mlp_out(time_embed)
|
|
time_embed = self.act_fn(time_embed)
|
|
embed = action_embed
|
|
adarms_cond = time_embed
|
|
|
|
if self.action_hidden_size < self.hidden_size:
|
|
# padding to hidden size
|
|
padding_size = self.hidden_size - self.action_hidden_size
|
|
padding = torch.zeros(
|
|
(embed.shape[0], embed.shape[1], padding_size),
|
|
device=embed.device,
|
|
dtype=embed.dtype,
|
|
)
|
|
embed = torch.cat([embed, padding], dim=-1)
|
|
|
|
return embed, adarms_cond
|
|
|
|
def flow_loss(
|
|
self,
|
|
action_hidden_states,
|
|
flow,
|
|
action_chunk,
|
|
dof_mask=None,
|
|
flow_loss_mask=None,
|
|
):
|
|
with torch.autocast("cuda", dtype=torch.float32):
|
|
action_pred = self.action_proj_back(
|
|
action_hidden_states[:, : self.action_hidden_size]
|
|
)
|
|
v_pred = action_pred
|
|
loss = self.mse_loss(v_pred, flow)
|
|
if dof_mask is not None:
|
|
dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1])
|
|
loss = loss * dof_mask
|
|
|
|
if flow_loss_mask is not None:
|
|
flow_loss_mask = (
|
|
flow_loss_mask.unsqueeze(-1)
|
|
.reshape(-1, 1)
|
|
.expand(-1, loss.shape[-1])
|
|
)
|
|
loss = loss * flow_loss_mask
|
|
return loss
|