Implementation:Facebookresearch Habitat lab ActionEmbedding

Knowledge Sources	Facebookresearch_Habitat_lab
Domains	Embodied_AI, Action_Representation
Last Updated	2026-02-15 00:00 GMT

Overview

The ActionEmbedding module provides NeRF-style sinusoidal embeddings for continuous (Box) actions, learned embeddings for discrete actions, and a composite ActionEmbedding class that handles dictionary action spaces containing multiple sub-spaces.

Description

This module defines three nn.Module classes:

BoxActionEmbedding implements NeRF-style positional encoding for continuous actions. It normalizes actions to [-1, 1] using the action space bounds, then applies sinusoidal frequency encoding with logarithmically-spaced frequencies: [sin(x * 2^t * pi), cos(x * 2^t * pi)] for t in 0..dim_per_action/2. The output dimension per action is dim_per_action.

DiscreteActionEmbedding uses a standard nn.Embedding table with an extra entry (index 0) serving as a start/padding token. Discrete actions are offset by +1 before lookup, and masked positions are set to the zero-th entry.

ActionEmbedding is the top-level module that iterates over all sub-spaces in a Habitat ActionSpace dictionary. For each sub-space, it creates the appropriate embedding module (Box or Discrete) and tracks the corresponding action slices. On forward, it applies each embedding module to its slice and concatenates the results. If all sub-spaces are EmptySpace, it falls back to a single discrete embedding.

Usage

Use ActionEmbedding in policy architectures that need dense representations of previous actions for recurrent processing. It handles mixed continuous/discrete action spaces automatically.

Code Reference

Source Location

• Repository: Facebookresearch_Habitat_lab
• File: habitat-baselines/habitat_baselines/rl/models/action_embedding.py
• Lines: 18-144

Signature

class BoxActionEmbedding(nn.Module):
    def __init__(self, action_space: gym.spaces.Box, dim_per_action: int = 32):
    def forward(self, action, masks=None):

class DiscreteActionEmbedding(nn.Module):
    def __init__(self, action_space: gym.spaces.Discrete, dim_per_action: int):
    def forward(self, action, masks=None):

class ActionEmbedding(nn.Module):
    def __init__(self, action_space: ActionSpace, dim_per_action: int = 32):
    def forward(self, action, masks=None):

Import

from habitat_baselines.rl.models.action_embedding import ActionEmbedding, BoxActionEmbedding, DiscreteActionEmbedding

I/O Contract

Inputs (ActionEmbedding)

Name	Type	Required	Description
action_space	ActionSpace	Yes	Habitat dictionary action space containing Box, Discrete, or EmptySpace sub-spaces
dim_per_action	int	No	Embedding dimension per individual action (default: 32)
action	Tensor	Yes	Action tensor passed to forward()
masks	Tensor	No	Mask tensor; unmasked positions are zeroed/set to start token

Outputs

Name	Type	Description
embedding	Tensor	Concatenated action embeddings with total dimension accessible via output_size property

Key Properties

output_size

@property
def output_size(self) -> int

Returns the total output dimension of all concatenated sub-embeddings.

Usage Examples

Basic Usage

import torch
from habitat.core.spaces import ActionSpace
import gym.spaces as spaces
from habitat_baselines.rl.models.action_embedding import ActionEmbedding

# Define a mixed action space
action_space = ActionSpace({
    "arm_action": spaces.Box(low=-1.0, high=1.0, shape=(7,)),
    "grip_action": spaces.Discrete(2),
})

# Create the embedding module
action_emb = ActionEmbedding(action_space, dim_per_action=32)
print(f"Output size: {action_emb.output_size}")

# Embed a batch of actions
batch_size = 16
actions = torch.randn(batch_size, 8)  # 7 continuous + 1 discrete
masks = torch.ones(batch_size, 8)
embedded = action_emb(actions, masks)
# embedded shape: (16, output_size)