rfx.utils
Utility functions and helpers for multi-embodiment training, including DOF padding and observation/action transforms.
Padding Functions
DOF padding utilities for training models across robots with different numbers of joints.
PaddingConfig
Configuration for state/action padding.
from rfx.utils import PaddingConfig
config = PaddingConfig(
state_dim=24,
action_dim=12,
max_state_dim=64,
max_action_dim=64
)
Constructor
PaddingConfig(
state_dim: int,
action_dim: int,
max_state_dim: int = 64,
max_action_dim: int = 64
)
Actual state dimension for this robot.
Actual action dimension for this robot.
Maximum state dimension across all robots in the dataset.
Maximum action dimension across all robots in the dataset.
pad_state
from rfx.utils import pad_state
import torch
state = torch.randn(32, 24) # (batch, state_dim)
padded = pad_state(state, state_dim=24, max_state_dim=64)
print(padded.shape) # (32, 64)
pad_state(
state: torch.Tensor,
state_dim: int,
max_state_dim: int
) -> torch.Tensor
State tensor of shape (batch, state_dim) or (batch, seq_len, state_dim).
Padded state tensor of shape (batch, max_state_dim) or (batch, seq_len, max_state_dim).
pad_action
from rfx.utils import pad_action
import torch
action = torch.randn(32, 12) # (batch, action_dim)
padded = pad_action(action, action_dim=12, max_action_dim=64)
print(padded.shape) # (32, 64)
pad_action(
action: torch.Tensor,
action_dim: int,
max_action_dim: int
) -> torch.Tensor
Action tensor of shape (batch, action_dim) or (batch, seq_len, action_dim).
Padded action tensor of shape (batch, max_action_dim) or (batch, seq_len, max_action_dim).
unpad_action
from rfx.utils import unpad_action
import torch
padded_action = torch.randn(32, 64) # (batch, max_action_dim)
action = unpad_action(padded_action, action_dim=12)
print(action.shape) # (32, 12)
unpad_action(
action: torch.Tensor,
action_dim: int
) -> torch.Tensor
Actual action dimension to extract.
Unpadded action tensor of shape (..., action_dim).
Observation and action transforms for training.
ObservationNormalizer
Running mean/std normalizer for observations using Welford’s algorithm.
from rfx.utils import ObservationNormalizer
import torch
normalizer = ObservationNormalizer(state_dim=24, clip=10.0)
# Update statistics
states = torch.randn(1000, 24)
normalizer.update(states)
# Normalize observations
obs = {"state": torch.randn(32, 24)}
normalized = normalizer.normalize(obs)
Constructor
ObservationNormalizer(
state_dim: int,
clip: float = 10.0,
eps: float = 1e-8
)
Dimension of the state to normalize.
Maximum absolute value after normalization.
Small constant for numerical stability.
Methods
update
normalizer.update(state: torch.Tensor) -> None
Batch of states, shape (batch_size, state_dim).
normalize
normalizer.normalize(obs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]
obs
dict[str, torch.Tensor]
required
Observation dictionary containing "state", optionally "images" and "language".
Normalized observation dictionary.
to_dict
normalizer.to_dict() -> dict[str, Any]
Dictionary containing all normalizer state.
from_dict
@classmethod
ObservationNormalizer.from_dict(d: dict[str, Any]) -> ObservationNormalizer
Serialized normalizer state.
Reconstructed normalizer.
ActionChunker
Action chunking for temporal abstraction. Ensembles multiple action predictions over a temporal horizon.
from rfx.utils import ActionChunker
import torch
chunker = ActionChunker(
horizon=4,
action_dim=12,
ensemble_mode="exponential",
temperature=0.5
)
# Add action chunks (e.g., from a model predicting 4 steps ahead)
for _ in range(4):
chunk = torch.randn(1, 4, 12) # (batch, horizon, action_dim)
chunker.add_chunk(chunk)
# Get ensembled action for current timestep
action = chunker.get_action() # (batch, action_dim)
Constructor
ActionChunker(
horizon: int,
action_dim: int,
ensemble_mode: str = "exponential",
temperature: float = 0.5
)
Number of timesteps to predict ahead.
Dimension of each action.
How to ensemble predictions: "first", "average", or "exponential".
Temperature for exponential weighting (only used if ensemble_mode=“exponential”).
Methods
add_chunk
chunker.add_chunk(chunk: torch.Tensor) -> None
Action chunk of shape (batch, horizon, action_dim).
get_action
chunker.get_action() -> torch.Tensor
Action tensor of shape (batch, action_dim).
reset
Clear all stored chunks.
Ensemble Modes
"first": Use only the first prediction from the most recent chunk"average": Average all predictions for the current timestep"exponential": Weighted average with exponentially decaying weights based on prediction age
