Implementation:Facebookresearch Habitat lab AgentAccessMgr

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

Overview

AgentAccessMgr is an abstract base class that defines the interface for managing the lifecycle of a policy, its data storage, and its updater within the Habitat baselines RL training loop.

Description

AgentAccessMgr serves as the central abstraction through which the trainer interacts with the agent's components. It bundles three concerns: the Policy (how actions are selected from observations), the Storage (how rollout data is collected and batched), and the Updater (how the policy is updated from collected data). The class defines abstract properties for accessing each component (actor_critic, rollouts, updater) and abstract methods for initialization (post_init), state management (get_resume_state, get_save_state, load_state_dict, load_ckpt_state_dict), mode switching (train, eval), and lifecycle hooks (after_update, pre_rollout). It also exposes nbuffers and masks_shape properties for configuring the training loop. Concrete implementations are registered via baseline_registry.register_agent_access_mgr.

Usage

Implement a concrete subclass of AgentAccessMgr to define how a specific agent manages its policy, storage, and updater. Register it with the baseline registry so the trainer can instantiate it by name from configuration.

Code Reference

Source Location

• Repository: Facebookresearch_Habitat_lab
• File: habitat-baselines/habitat_baselines/rl/ppo/agent_access_mgr.py
• Lines: 17-131

Signature

class AgentAccessMgr(ABC):
    @abstractmethod
    def __init__(
        self,
        config: "DictConfig",
        env_spec: EnvironmentSpec,
        is_distrib: bool,
        device,
        resume_state: Optional[Dict[str, Any]],
        num_envs: int,
        percent_done_fn: Callable[[], float],
        lr_schedule_fn: Optional[Callable[[float], float]] = None,
    ):

Import

from habitat_baselines.rl.ppo.agent_access_mgr import AgentAccessMgr

I/O Contract

Inputs

Name	Type	Required	Description
config	DictConfig	Yes	Full training configuration
env_spec	EnvironmentSpec	Yes	Environment action and observation space specifications
is_distrib	bool	Yes	Whether training is distributed across multiple processes
device	torch.device	Yes	Device for tensor allocation
resume_state	Optional[Dict[str, Any]]	Yes	State dictionary for resuming training, or None for fresh start
num_envs	int	Yes	Number of parallel environments
percent_done_fn	Callable[[], float]	Yes	Function returning training progress as a fraction [0, 1]
lr_schedule_fn	Optional[Callable[[float], float]]	No	Learning rate schedule function mapping progress to LR multiplier

Outputs (Abstract Properties)

Name	Type	Description
actor_critic	Policy	The current policy used for action selection
rollouts	Storage	The current rollout data storage
updater	Updater	The current policy updater
nbuffers	int	Number of storage buffers
masks_shape	Tuple	Shape of the masks tensor

Abstract Methods

post_init

def post_init(self, create_rollouts_fn: Optional[Callable] = None) -> None

Called after construction to set up rollout storage.

get_resume_state / get_save_state

def get_resume_state(self) -> Dict[str, Any]
def get_save_state(self) -> Dict[str, Any]

Returns state dictionaries for resuming or saving checkpoints.

load_ckpt_state_dict / load_state_dict

def load_ckpt_state_dict(self, ckpt: Dict) -> None
def load_state_dict(self, state: Dict) -> None

Loads checkpoint state for evaluation or full state for training resumption.

after_update / pre_rollout

def after_update(self) -> None
def pre_rollout(self) -> None

Lifecycle hooks called by the trainer after policy updates and before rollout collection.

Usage Examples

Basic Usage

from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.rl.ppo.agent_access_mgr import AgentAccessMgr

@baseline_registry.register_agent_access_mgr
class MyAgentAccessMgr(AgentAccessMgr):
    def __init__(self, config, env_spec, is_distrib, device,
                 resume_state, num_envs, percent_done_fn, lr_schedule_fn=None):
        # Initialize policy, storage, and updater
        self._policy = build_policy(config, env_spec)
        self._updater = build_updater(config, self._policy)

    def post_init(self, create_rollouts_fn=None):
        self._storage = create_rollouts_fn()

    @property
    def actor_critic(self):
        return self._policy

    @property
    def rollouts(self):
        return self._storage

    @property
    def updater(self):
        return self._updater

    # ... implement remaining abstract methods