opentau.policies.layers

Reusable parametric layers shared across VLA policies.

PerGroupLinear is a drop-in replacement for nn.Linear that keeps an independent (weight, bias) per group and selects one row per sample with a (B,) long index — the same per-sample group axis the stacked Normalize/Unnormalize stat buffers use (see opentau.policies.normalize). Its state_dict key names match nn.Linear exactly (<name>.weight / <name>.bias) but carry a leading group axis, so the legacy single-group promotion shim (unsqueeze(0)-style, see PreTrainedPolicy._promote_legacy_norm_buffers_in_state_dict) applies to it.

Classes

PerGroupLinear(in_features, out_features, ...)

nn.Linear with one independent (weight, bias) per group.

class opentau.policies.layers.PerGroupLinear(in_features: int, out_features: int, num_groups: int, bias: bool = True)

Bases: Module

nn.Linear with one independent (weight, bias) per group.

Parameters:

• in_features – Size of each input sample.

• out_features – Size of each output sample.

• num_groups – Number of independent linear maps. 1 makes this numerically identical to a plain nn.Linear — the forward takes an F.linear fast path with the un-cast parameters, so its dtype / autocast behavior matches a plain nn.Linear bit-for-bit.

• bias – If True (default), adds a per-group learnable bias.

Shape:

• weight: (num_groups, out_features, in_features)

• bias: (num_groups, out_features)

These are an nn.Linear parameter with a leading group axis, which is why the state_dict keys stay identical to nn.Linear (just one rank higher). A legacy nn.Linear checkpoint promotes into a num_groups=1 PerGroupLinear by prepending that axis.

Forward:

forward(x, group_index=None) where x is (B, *, in_features) and group_index is a (B,) long tensor in [0, num_groups). group_index=None routes every sample to group 0 (the single-group default; used by direct-call unit tests and any caller that has not threaded an index).

__init__(in_features: int, out_features: int, num_groups: int, bias: bool = True)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

extra_repr() → str

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x: Tensor, group_index: Tensor | None = None) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameters() → None

Initialize every group exactly like nn.Linear.reset_parameters.

Per-row kaiming-uniform weight (a=sqrt(5)) + fan-in-scaled uniform bias, so a freshly built per-group module matches nn.Linear’s init distribution on each row.