AllenAct Documentation

AllenAct is a modular and flexible learning framework specifically designed for Embodied AI research. It provides first-class support for various embodied environments, tasks, and algorithms, with a focus on PyTorch.

Key capabilities include:

• Decoupled Tasks and Environments: Researchers can implement multiple tasks within the same environment.
• Algorithm Support: Includes on-policy algorithms (PPO, DD-PPO, A2C), Imitation Learning, DAgger, and offline training.
• Sequential Training: Easy experimentation with sequences of training routines.
• Simultaneous Losses: Ability to combine multiple losses (e.g., self-supervised loss with PPO loss) during training.
• Multi-agent Support: Built-in support for multi-agent algorithms and tasks.
• Visualizations: Integrated Tensorboard support for first/third-person views and intermediate model tensors.
• Action Spaces: Supports both discrete and continuous action spaces.

The AllenAct repository is organized into several functional directories that separate core framework logic, environment plugins, datasets, and project-specific code:

• allenact/: The core framework. Contains training/inference algorithms (allenact.algorithms), base abstractions (allenact.base_abstractions), and basic embodied AI models (allenact.embodiedai).
• allenact_plugins/: Environment-specific implementations. Each plugin (e.g., {environment}_plugin) contains configs/, data/, and scripts/ for setting up that specific environment.
• datasets/: Task-specific data storage. Includes scripts for automated downloading (e.g., download_navigation_datasets.sh).
• projects/: Project-specific code, including experiment configurations and scripts for result processing or visualization.
• pretrained_model_ckpts/: Storage for pretrained model checkpoints.
• scripts/: Framework-wide utility scripts (e.g., documentation building, code formatting, running tests, or starting an xserver).
• tests/: Unit tests for the allenact core.

Visualization in AllenAct is managed by the VizSuite class (defined in allenact.utils.viz_utils). You can customize visualizations by instantiating different visualization types as plugins to the VizSuite within your ExperimentConfig when in test mode.

Data Sources

VizSuite can pull data from several sources:

• Task output: e.g., 2D trajectories.
• Vector task: e.g., egocentric views.
• Rollout storage: e.g., recurrent memory, taken action logprobs, episode masks.
• ActorCriticOutput: e.g., action probabilities.

Available Visualization Types

• TrajectoryViz: Generic 2D trajectory view.
• AgentViewViz: RGB egocentric view.
• ActorViz: Action probabilities from ActorCriticOutput[CategoricalDistr].
• TensorViz1D: Evolution of a 1D point from RolloutStorage over time.
• TensorViz2D: Evolution of a 2D vector from RolloutStorage over time.
• ThorViz: Specialized 2D trajectory view for RoboThor.

To enable these, override the machine_params method in your ExperimentConfig to call res.set_visualizer(self.get_viz(mode)) when mode == "test".

class PointNavRoboThorRGBPPOVizExperimentConfig(PointNavRoboThorRGBPPOExperimentConfig):
    viz_ep_ids = ["FloorPlan_Train1_1_3", "FloorPlan_Train1_1_4"]
    viz_video_ids = [["FloorPlan_Train1_1_3"], ["FloorPlan_Train1_1_4"]]
    viz: Optional[VizSuite] = None

    def get_viz(self, mode):
        if self.viz is not None:
            return self.viz

        self.viz = VizSuite(
            episode_ids=self.viz_ep_ids,
            mode=mode,
            base_trajectory=TrajectoryViz(path_to_target_location=("task_info", "target",)),
            egeocentric=AgentViewViz(max_video_length=100, episode_ids=self.viz_video_ids),
            action_probs=ActorViz(figsize=(3.25, 10), fontsize=18),
            taken_action_logprobs=TensorViz1D(),
            episode_mask=TensorViz1D(rollout_source=("masks",)),
            rnn_memory=TensorViz2D(rollout_source=("memory", "single_belief")),
            thor_trajectory=ThorViz(figsize=(16, 8), viz_rows_cols=(448, 448), scenes=("FloorPlan_Train{}_{}", 1, 1, 1, 1)),
        )
        return self.viz

    def machine_params(self, mode="train", **kwargs):
        res = super().machine_params(mode, **kwargs)
        if mode == "test":
            res.set_visualizer(self.get_viz(mode))
        return res

To define a complete experiment in allenact, you must implement the ExperimentConfig abstraction. This class encapsulates all configuration for training, validation, and testing. An OnPolicyRunner uses this config to orchestrate OnPolicyTrainer (for training) and OnPolicyInference (for validation/testing).

Key methods to implement in your ExperimentConfig subclass:

• tag(): Returns a str to identify the experiment.
• create_model(**kwargs): Returns an nn.Module (the agent's actor-critic model).
• make_sampler_fn(**kwargs): Returns a TaskSampler instance.
• {train,valid,test}_task_sampler_args(...): Returns a dictionary of initialization parameters for the task samplers used in each phase.
• machine_params(mode="train", **kwargs): Returns a dictionary specifying nprocesses and devices (e.g., GPU IDs).
• training_pipeline(**kwargs): Returns a TrainingPipeline object describing losses, optimizers, and scheduling.

class MiniGridTutorialExperimentConfig(ExperimentConfig):
    @classmethod
    def tag(cls) -> str:
        return "MiniGridTutorial"

    @classmethod
    def create_model(cls, **kwargs) -> nn.Module:
        # Return your model here
        pass

    @classmethod
    def make_sampler_fn(cls, **kwargs) -> TaskSampler:
        # Return your sampler here
        pass

    def train_task_sampler_args(self, ...) -> Dict[str, Any]:
        # Return training sampler args
        pass

    def machine_params(cls, mode="train", **kwargs) -> Dict[str, Any]:
        # Return hardware config
        pass

    @classmethod
    def training_pipeline(cls, **kwargs) -> TrainingPipeline:
        # Return pipeline config
        pass