Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,10 @@
|
||||
"""Action model components: normalizer, processor, head, MoE."""
|
||||
|
||||
from wall_x.model.core.action.head import SinusoidalPosEmb
|
||||
from wall_x.model.core.action.normalizer import (
|
||||
Normalizer,
|
||||
create_normalizers,
|
||||
normalize_data_with_virtual_tail,
|
||||
unnormalize_data_with_virtual_tail,
|
||||
)
|
||||
from wall_x.model.core.action.processor import ActionProcessor
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Action head helpers used by the VLA action processor."""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class SinusoidalPosEmb(nn.Module):
|
||||
"""Sinusoidal timestep embedding for action flow timesteps."""
|
||||
|
||||
def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0):
|
||||
super().__init__()
|
||||
if dim % 2 != 0:
|
||||
raise ValueError(f"embedding_dim ({dim}) must be divisible by 2")
|
||||
self.dim = dim
|
||||
self.min_period = min_period
|
||||
self.max_period = max_period
|
||||
|
||||
def forward(self, x):
|
||||
half_dim = self.dim // 2
|
||||
exponent = math.log(10000) / (half_dim - 1)
|
||||
frequencies = torch.exp(
|
||||
torch.arange(half_dim, device=x.device, dtype=torch.float32) * -exponent
|
||||
)
|
||||
emb = x[:, None] * frequencies[None, :]
|
||||
return torch.cat((emb.sin(), emb.cos()), dim=-1)
|
||||
@@ -0,0 +1,126 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint as cp
|
||||
from transformers.activations import ACT2FN
|
||||
|
||||
from wall_x.model.core.ops import permute, unpermute
|
||||
|
||||
|
||||
class TokenTypeRouter(nn.Module):
|
||||
def __init__(self, num_experts: int):
|
||||
super().__init__()
|
||||
self.num_experts = num_experts
|
||||
|
||||
def forward(self, token_types: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Route tokens to experts based on token_type.
|
||||
|
||||
Args:
|
||||
token_types (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each token type.
|
||||
|
||||
Returns:
|
||||
experts_indices (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each assigned expert index.
|
||||
"""
|
||||
# Simple rule: assign by token_type modulo the expert count
|
||||
experts_indices = token_types % self.num_experts
|
||||
return experts_indices
|
||||
|
||||
|
||||
class BlockSparseMLP(nn.Module):
|
||||
def __init__(self, config, use_selective_recompute: bool = False):
|
||||
super().__init__()
|
||||
self.hidden_size = config["hidden_size"]
|
||||
self.intermediate_size = config["intermediate_size"]
|
||||
self.hidden_act = config["hidden_act"]
|
||||
|
||||
self.use_selective_recompute = use_selective_recompute
|
||||
|
||||
self.gate_up_proj = nn.Linear(
|
||||
self.hidden_size, 2 * self.intermediate_size, bias=False
|
||||
)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
|
||||
self.act_fn = ACT2FN[self.hidden_act]
|
||||
|
||||
# FIXME: The full MLP is recomputed for now; recomputing only activations can be optimized later.
|
||||
def _full_mlp(self, hidden_state):
|
||||
gate_up_out = self.gate_up_proj(hidden_state)
|
||||
gate_out, up_out = gate_up_out.split(
|
||||
[self.intermediate_size, self.intermediate_size], dim=-1
|
||||
)
|
||||
|
||||
act_out = self.act_fn(gate_out) * up_out
|
||||
return self.down_proj(act_out)
|
||||
|
||||
def forward(self, hidden_state):
|
||||
if self.use_selective_recompute:
|
||||
# Checkpoint-recompute the whole expert MLP
|
||||
return cp.checkpoint(
|
||||
self._full_mlp,
|
||||
hidden_state,
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
return self._full_mlp(hidden_state)
|
||||
|
||||
|
||||
class SparseMoeBlock(nn.Module):
|
||||
def __init__(self, config, num_experts: int, use_selective_recompute: bool = False):
|
||||
super().__init__()
|
||||
self.num_experts = num_experts
|
||||
self.use_selective_recompute = use_selective_recompute
|
||||
|
||||
# Pass use_selective_recompute to each expert
|
||||
self.experts = nn.ModuleList(
|
||||
[
|
||||
BlockSparseMLP(
|
||||
config.experts[i], use_selective_recompute=use_selective_recompute
|
||||
)
|
||||
for i in range(num_experts)
|
||||
]
|
||||
)
|
||||
|
||||
if not hasattr(config, "dim_inputs") or not config.dim_inputs:
|
||||
raise ValueError("config.dim_inputs must be set")
|
||||
|
||||
self.dim_inputs = config.dim_inputs
|
||||
self.permuted = config.mot_opt
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
experts_indices: torch.Tensor,
|
||||
start_indices: torch.Tensor,
|
||||
end_indices: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
|
||||
if self.permuted:
|
||||
permuted_inputs = hidden_states
|
||||
else:
|
||||
batch_size, seq_length, hidden_dim = hidden_states.shape
|
||||
|
||||
flat_hidden = hidden_states.reshape(-1, hidden_dim)
|
||||
experts_indices = experts_indices.reshape(-1)
|
||||
probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1)
|
||||
permuted_inputs, row_id_map = permute(flat_hidden, experts_indices)
|
||||
|
||||
# buffer
|
||||
final_output = torch.zeros_like(permuted_inputs)
|
||||
|
||||
# Expert forward, including selective recompute
|
||||
for expert_idx, expert in enumerate(self.experts):
|
||||
start, end = start_indices[expert_idx], end_indices[expert_idx]
|
||||
if start == end:
|
||||
continue
|
||||
|
||||
dim_input = self.dim_inputs[expert_idx]
|
||||
expert_input = permuted_inputs[start:end, :dim_input]
|
||||
|
||||
partial_output = expert(expert_input)
|
||||
final_output[start:end, :dim_input] = partial_output[:, :dim_input]
|
||||
|
||||
if self.permuted:
|
||||
return final_output
|
||||
else:
|
||||
final_output = unpermute(final_output, row_id_map, probs)
|
||||
return final_output.reshape(batch_size, seq_length, hidden_dim)
|
||||
@@ -0,0 +1,447 @@
|
||||
import json
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from wall_x.utils.constant import _ACTION_KEY_FULL_MAPPING as _MODEL_KEY_TO_RAW_KEY
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def resolve_normalizer_dataset_names(
|
||||
dataset_names,
|
||||
normalizer,
|
||||
batch_size,
|
||||
*,
|
||||
source="dataset_names",
|
||||
allow_skip_names=None,
|
||||
):
|
||||
"""Validate dataset names before indexing a normalizer ParameterDict.
|
||||
|
||||
The normalizer key is data-dependent. Callers should pass the
|
||||
dataset names carried by the batch/model inputs, matching the qact/VLA
|
||||
path. This helper intentionally only validates presence, batch-size
|
||||
alignment, and key membership; it does not remap unknown names or fall
|
||||
back to a default key, since silent fallback can apply the wrong action
|
||||
scale without failing loudly. Pass ``allow_skip_names`` explicitly only
|
||||
for names that should bypass normalizer lookup, such as multimodal-only
|
||||
rows that do not carry action values.
|
||||
"""
|
||||
if dataset_names is None:
|
||||
raise KeyError(f"Missing {source}; cannot choose normalizer key.")
|
||||
|
||||
if isinstance(dataset_names, str):
|
||||
names = [dataset_names]
|
||||
elif isinstance(dataset_names, (list, tuple)):
|
||||
names = [str(name) for name in dataset_names]
|
||||
else:
|
||||
raise TypeError(
|
||||
f"{source} must be a string or sequence of strings, got "
|
||||
f"{type(dataset_names).__name__}"
|
||||
)
|
||||
|
||||
if len(names) != batch_size:
|
||||
raise ValueError(
|
||||
f"{source} count ({len(names)}) does not match batch size "
|
||||
f"({batch_size}). names={names}"
|
||||
)
|
||||
|
||||
available = list(getattr(normalizer, "delta", {}).keys())
|
||||
if not available:
|
||||
raise KeyError("Normalizer has no registered dataset keys.")
|
||||
|
||||
allowed_skip = set(allow_skip_names or ())
|
||||
missing = sorted(
|
||||
{name for name in names if name not in available and name not in allowed_skip}
|
||||
)
|
||||
if missing:
|
||||
raise KeyError(
|
||||
f"{source} contains keys not present in normalizer: {missing}. "
|
||||
f"available={available}"
|
||||
)
|
||||
|
||||
return names
|
||||
|
||||
|
||||
def _normalizer_dataset_name_list(dataset_names):
|
||||
if dataset_names is None:
|
||||
return []
|
||||
if isinstance(dataset_names, str):
|
||||
return [dataset_names]
|
||||
return list(dataset_names)
|
||||
|
||||
|
||||
def _normalizer_width(normalizer, dataset_names):
|
||||
names = _normalizer_dataset_name_list(dataset_names)
|
||||
if not names:
|
||||
return None
|
||||
if any(name not in normalizer.delta for name in names):
|
||||
return None
|
||||
widths = [int(normalizer.delta[name].shape[0]) for name in names]
|
||||
if not widths or len(set(widths)) != 1:
|
||||
return None
|
||||
return widths[0]
|
||||
|
||||
|
||||
def _has_uniform_virtual_tail_mask(dof_mask, width):
|
||||
if dof_mask is None:
|
||||
return False
|
||||
mask = dof_mask.reshape(-1, dof_mask.shape[-1]).bool()
|
||||
if mask.shape[-1] <= width:
|
||||
return False
|
||||
prefix_active = mask[:, :width].all(dim=1)
|
||||
tail_inactive = (~mask[:, width:]).all(dim=1)
|
||||
return bool((prefix_active & tail_inactive).all().item())
|
||||
|
||||
|
||||
def normalize_data_with_virtual_tail(normalizer, values, dataset_names, dof_mask):
|
||||
"""Normalize real prefix dims when model tensors include a virtual tail."""
|
||||
width = _normalizer_width(normalizer, dataset_names)
|
||||
if width is None or values.shape[-1] == width:
|
||||
return normalizer.normalize_data(values, dataset_names)
|
||||
if values.shape[-1] < width or not _has_uniform_virtual_tail_mask(dof_mask, width):
|
||||
return normalizer.normalize_data(values, dataset_names)
|
||||
|
||||
out = values.clone()
|
||||
out[..., :width] = normalizer.normalize_data(values[..., :width], dataset_names)
|
||||
out[..., width:] = 0
|
||||
return out
|
||||
|
||||
|
||||
def unnormalize_data_with_virtual_tail(normalizer, values, dataset_names, dof_mask):
|
||||
"""Unnormalize real prefix dims when model tensors include a virtual tail."""
|
||||
width = _normalizer_width(normalizer, dataset_names)
|
||||
if width is None or values.shape[-1] == width:
|
||||
return normalizer.unnormalize_data(values, dataset_names, None)
|
||||
if values.shape[-1] < width or not _has_uniform_virtual_tail_mask(dof_mask, width):
|
||||
return normalizer.unnormalize_data(values, dataset_names, dof_mask)
|
||||
|
||||
out = values.clone()
|
||||
out[..., :width] = normalizer.unnormalize_data(
|
||||
values[..., :width], dataset_names, None
|
||||
)
|
||||
out[..., width:] = 0
|
||||
return out
|
||||
|
||||
|
||||
def print_rank_last(message):
|
||||
"""If distributed is initialized, log only on last rank."""
|
||||
if torch.distributed.is_initialized():
|
||||
if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1):
|
||||
logger.info(message)
|
||||
else:
|
||||
logger.info(message)
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return instance
|
||||
|
||||
@classmethod
|
||||
def from_lerobot_norm_stats(
|
||||
cls,
|
||||
action_stats,
|
||||
dataset_name,
|
||||
):
|
||||
|
||||
# Create instance without calling __init__
|
||||
instance = cls.__new__(cls)
|
||||
nn.Module.__init__(instance)
|
||||
|
||||
# Set containers
|
||||
instance.min = nn.ParameterDict()
|
||||
instance.delta = nn.ParameterDict()
|
||||
|
||||
# Fill dataset entry
|
||||
instance.min[dataset_name] = nn.Parameter(action_stats.min, requires_grad=False)
|
||||
instance.delta[dataset_name] = nn.Parameter(
|
||||
action_stats.delta, requires_grad=False
|
||||
)
|
||||
|
||||
# Record keys
|
||||
instance.min_key = "min"
|
||||
instance.delta_key = "delta"
|
||||
|
||||
return instance
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
action_statistic_dof,
|
||||
dof_config,
|
||||
min_key="min",
|
||||
delta_key="delta",
|
||||
name="normalizer",
|
||||
):
|
||||
super(Normalizer, self).__init__()
|
||||
|
||||
self.min_key = min_key
|
||||
self.delta_key = delta_key
|
||||
|
||||
action_statistic = {}
|
||||
normalizer_missing_information = []
|
||||
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 not in action_statistic_dof[robot_name]
|
||||
and k.replace("master_", "follow_", 1)
|
||||
in action_statistic_dof[robot_name]
|
||||
):
|
||||
k = k.replace("master_", "follow_", 1)
|
||||
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:
|
||||
normalizer_missing_information.append(
|
||||
f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} 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:
|
||||
# k is a model key; action_statistic_dof may store raw keys - try fallback lookup
|
||||
raw_key = _MODEL_KEY_TO_RAW_KEY.get(k)
|
||||
if (
|
||||
raw_key is not None
|
||||
and raw_key in action_statistic_dof[robot_name]
|
||||
):
|
||||
stat = action_statistic_dof[robot_name][raw_key]
|
||||
if min_key in stat and delta_key in stat:
|
||||
all_dof_min.extend(stat[min_key])
|
||||
all_dof_delta.extend(stat[delta_key])
|
||||
else:
|
||||
normalizer_missing_information.append(
|
||||
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])
|
||||
else:
|
||||
normalizer_missing_information.append(
|
||||
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
|
||||
|
||||
if not torch.distributed.is_initialized() or torch.distributed.get_rank() == (
|
||||
torch.distributed.get_world_size() - 1
|
||||
):
|
||||
# export normalizer_missing_information to file
|
||||
with open(f"normalizer_missing_information_{name}.txt", "w") as f:
|
||||
for info in normalizer_missing_information:
|
||||
f.write(info + "\n")
|
||||
logger.info(
|
||||
"Normalizer missing information saved to normalizer_missing_information_%s.txt",
|
||||
name,
|
||||
)
|
||||
|
||||
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()
|
||||
}
|
||||
)
|
||||
|
||||
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):
|
||||
# if dataset_name == "ex_normal":
|
||||
# dataset_name = "x2_normal"
|
||||
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
|
||||
|
||||
|
||||
def _pad_lerobot_stats(stats, target_dim, label):
|
||||
"""Tail-pad flat LeRobot stats to the model's configured action/state dim."""
|
||||
target_dim = int(target_dim)
|
||||
current_dim = int(stats.min.numel())
|
||||
if current_dim == target_dim:
|
||||
return stats
|
||||
if current_dim > target_dim:
|
||||
raise ValueError(
|
||||
f"LeRobot {label} norm stats dim ({current_dim}) exceeds configured "
|
||||
f"{label} dim ({target_dim})"
|
||||
)
|
||||
|
||||
pad_dim = target_dim - current_dim
|
||||
min_stat = torch.cat(
|
||||
[stats.min.flatten(), torch.zeros(pad_dim, dtype=stats.min.dtype)]
|
||||
)
|
||||
delta = torch.cat(
|
||||
[stats.delta.flatten(), torch.ones(pad_dim, dtype=stats.delta.dtype)]
|
||||
)
|
||||
max_stat = torch.cat(
|
||||
[stats.max.flatten(), torch.ones(pad_dim, dtype=stats.max.dtype)]
|
||||
)
|
||||
logger.info(
|
||||
"Padded LeRobot %s normalizer stats from %d to %d dims",
|
||||
label,
|
||||
current_dim,
|
||||
target_dim,
|
||||
)
|
||||
return type(stats)(min=min_stat, max=max_stat, delta=delta)
|
||||
|
||||
|
||||
def pad_normalizer_to_dim(normalizer, target_dim, label):
|
||||
"""Tail-pad every dataset entry to ``target_dim`` (virtual ``action_padding`` tail)."""
|
||||
target_dim = int(target_dim)
|
||||
for name in list(normalizer.min.keys()):
|
||||
current_dim = int(normalizer.min[name].numel())
|
||||
if current_dim == target_dim:
|
||||
continue
|
||||
if current_dim > target_dim:
|
||||
raise ValueError(
|
||||
f"Normalizer {label}[{name!r}] dim ({current_dim}) exceeds configured "
|
||||
f"{label} dim ({target_dim})"
|
||||
)
|
||||
|
||||
pad_dim = target_dim - current_dim
|
||||
dtype = normalizer.min[name].dtype
|
||||
normalizer.min[name] = nn.Parameter(
|
||||
torch.cat(
|
||||
[
|
||||
normalizer.min[name].flatten(),
|
||||
torch.zeros(pad_dim, dtype=dtype),
|
||||
]
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
normalizer.delta[name] = nn.Parameter(
|
||||
torch.cat(
|
||||
[
|
||||
normalizer.delta[name].flatten(),
|
||||
torch.ones(pad_dim, dtype=dtype),
|
||||
]
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
logger.info(
|
||||
"Padded %s normalizer[%r] from %d to %d dims",
|
||||
label,
|
||||
name,
|
||||
current_dim,
|
||||
target_dim,
|
||||
)
|
||||
|
||||
|
||||
def create_normalizers_from_lerobot_norm_stats(
|
||||
norm_stats, dataset_name, action_dim, propri_dim
|
||||
):
|
||||
"""Create model normalizers from LeRobot flat action/state norm stats."""
|
||||
action_stats = _pad_lerobot_stats(norm_stats["action"], action_dim, "action")
|
||||
propri_stats = _pad_lerobot_stats(norm_stats["state"], propri_dim, "state")
|
||||
return (
|
||||
Normalizer.from_lerobot_norm_stats(action_stats, dataset_name),
|
||||
Normalizer.from_lerobot_norm_stats(propri_stats, dataset_name),
|
||||
)
|
||||
|
||||
|
||||
def create_normalizers(config, action_statistic_dof=None):
|
||||
"""Create action and proprioception normalizers from explicit stats.
|
||||
|
||||
Normalization statistics must come from the config, checkpoint, or dataset.
|
||||
Public Wall-X builds intentionally do not bundle private default stats.
|
||||
"""
|
||||
|
||||
if action_statistic_dof is None:
|
||||
custom_path = config.get("customized_action_statistic_dof")
|
||||
if custom_path:
|
||||
with open(custom_path, "r") as f:
|
||||
action_statistic_dof = json.load(f)
|
||||
else:
|
||||
raise ValueError(
|
||||
"create_normalizers requires action statistics. Provide "
|
||||
"`customized_action_statistic_dof` in the config or pass "
|
||||
"`action_statistic_dof` from the checkpoint/dataset."
|
||||
)
|
||||
|
||||
min_key = config.get("min_key", "min")
|
||||
delta_key = config.get("delta_key", "delta")
|
||||
|
||||
# Required keys for constructing action/proprio normalizers.
|
||||
missing_required = []
|
||||
if config.get("dof_config") is None:
|
||||
missing_required.append("dof_config")
|
||||
if config.get("agent_pos_config") is None:
|
||||
missing_required.append("agent_pos_config")
|
||||
if missing_required:
|
||||
raise KeyError(
|
||||
"create_normalizers requires non-None config keys: "
|
||||
+ ", ".join(missing_required)
|
||||
)
|
||||
|
||||
normalizer_action = Normalizer(
|
||||
action_statistic_dof,
|
||||
config["dof_config"],
|
||||
min_key=min_key,
|
||||
delta_key=delta_key,
|
||||
)
|
||||
normalizer_propri = Normalizer(
|
||||
action_statistic_dof,
|
||||
config["agent_pos_config"],
|
||||
min_key=min_key,
|
||||
delta_key=delta_key,
|
||||
)
|
||||
return normalizer_action, normalizer_propri, action_statistic_dof
|
||||
@@ -0,0 +1,317 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from wall_x.model.core.action.head import SinusoidalPosEmb
|
||||
from wall_x.model.core.action.normalizer import print_rank_last
|
||||
|
||||
|
||||
class ActionProcessor(nn.Module):
|
||||
"""
|
||||
Action processor
|
||||
|
||||
Main responsibilities:
|
||||
1. Add noise to action sequences
|
||||
2. Generate time encodings
|
||||
3. Project actions into the model hidden space
|
||||
|
||||
Uses a Beta distribution to control noise scheduling and provide flexible noise injection.
|
||||
"""
|
||||
|
||||
def __init__(self, config):
|
||||
"""
|
||||
Args:
|
||||
config: Configuration object containing:
|
||||
- action_dim: action-space dimension
|
||||
- hidden_size: model hidden size
|
||||
- noise_scheduler: noise scheduler configuration
|
||||
"""
|
||||
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 = getattr(config, "hidden_size", config.dim_inputs[0])
|
||||
|
||||
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.s = noise_scheduler_config.get("s", 0.999)
|
||||
self.time_shift = noise_scheduler_config.get(
|
||||
"time_shift", 1.0
|
||||
) # time shift factor
|
||||
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
|
||||
|
||||
def get_inference_times(self, num_steps, device, dtype):
|
||||
"""
|
||||
Get inference timesteps
|
||||
|
||||
Apply time shift and scaling
|
||||
|
||||
Args:
|
||||
num_steps (int): number of inference steps
|
||||
device: Device type
|
||||
dtype: dtype
|
||||
|
||||
Returns:
|
||||
torch.Tensor: inference timestep sequence
|
||||
"""
|
||||
times = torch.linspace(0.0, 1.0, num_steps + 1, device=device, dtype=dtype)
|
||||
if self.time_shift != 1.0:
|
||||
times = (self.time_shift * times) / (1 + (self.time_shift - 1) * times)
|
||||
times = times * self.s
|
||||
return times
|
||||
|
||||
def proprioception_proj(
|
||||
self, proprioception, dataset_names=None, dof_mask=None, use_history=False
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
proprioception: [batch_size, 1, action_dim]
|
||||
dataset_names: [batch_size]
|
||||
dof_mask: [batch_size, action_dim]
|
||||
"""
|
||||
with torch.autocast("cuda", dtype=torch.float32):
|
||||
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, sample_time, dof_mask=None):
|
||||
"""
|
||||
Args:
|
||||
action_chunk (torch.Tensor): action sequence with shape [batch_size, action_chunk_len, action_dim]
|
||||
dataset_names: [batch_size]
|
||||
dof_mask: [batch_size, action_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: processed action representation with shape [batch_size, seq_len, hidden_size]
|
||||
"""
|
||||
with torch.autocast("cuda", dtype=torch.float32):
|
||||
action_chunk = action_chunk.to(dtype=torch.float32)
|
||||
|
||||
# 1. add noise to action_chunk
|
||||
noise = torch.randn_like(action_chunk)
|
||||
time_expanded = sample_time.unsqueeze(-1).unsqueeze(-1)
|
||||
noisy_action = (
|
||||
1 - time_expanded
|
||||
) * noise + time_expanded * action_chunk # denoise from 0 to 1; integration does not need a negative sign
|
||||
flow = action_chunk - noise # used to compute loss
|
||||
|
||||
# 2. sinusoidal positional encoding for timesteps
|
||||
time_embed = self.time_embed(sample_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:
|
||||
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]
|
||||
)
|
||||
|
||||
if getattr(self.config, "use_x_pred", False):
|
||||
noisy_action_flat = self.noisy_action.reshape(
|
||||
-1, self.noisy_action.shape[-1]
|
||||
)
|
||||
time_expanded_flat = self.time_expanded.expand(
|
||||
-1, self.noisy_action.shape[1], -1
|
||||
).reshape(-1, 1)
|
||||
v_pred = (action_pred - noisy_action_flat) / torch.clamp(
|
||||
1 - time_expanded_flat, min=0.05
|
||||
)
|
||||
x_pred = action_pred
|
||||
else:
|
||||
v_pred = action_pred
|
||||
time_expanded_flat = self.time_expanded.expand(
|
||||
-1, self.noisy_action.shape[1], -1
|
||||
).reshape(-1, 1)
|
||||
x_pred = (1 - time_expanded_flat) * v_pred + self.noisy_action.reshape(
|
||||
-1, self.noisy_action.shape[-1]
|
||||
)
|
||||
|
||||
if getattr(self.config, "use_x_loss", False):
|
||||
loss = self.mse_loss(
|
||||
x_pred,
|
||||
action_chunk.reshape(-1, action_chunk.shape[-1]).to(
|
||||
dtype=x_pred.dtype
|
||||
),
|
||||
)
|
||||
else:
|
||||
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
|
||||
Reference in New Issue
Block a user