Skip to content

Defining an experiment#

Let's look at an example experiment configuration for an object navigation example with an actor-critic agent observing RGB images from the environment and target object classes from the task. This is a simplified example where the agent is confined to a single iTHOR scene (FloorPlan1) and needs to find a single object (a tomato). To see how one might running a "full"/"hard" version of navigation within AI2-THOR, see our tutorials PointNav in RoboTHOR and Swapping in a new environment.

The interface to be implemented by the experiment specification is defined in core.base_abstractions.experiment_config. If you'd like to skip ahead and see the finished configuration, see here. We begin by making the following imports:

from math import ceil
from typing import Dict, Any, List, Optional

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR

from core.algorithms.onpolicy_sync.losses import PPO
from core.algorithms.onpolicy_sync.losses.ppo import PPOConfig
from core.base_abstractions.experiment_config import ExperimentConfig
from core.base_abstractions.sensor import SensorSuite
from core.base_abstractions.task import TaskSampler
from plugins.ithor_plugin.ithor_sensors import RGBSensorThor, GoalObjectTypeThorSensor
from plugins.ithor_plugin.ithor_task_samplers import ObjectNavTaskSampler
from plugins.ithor_plugin.ithor_tasks import ObjectNavTask
from projects.objectnav_baselines.models.object_nav_models import (
    ObjectNavBaselineActorCritic,
)
from utils.experiment_utils import Builder, PipelineStage, TrainingPipeline, LinearDecay

Now first method to implement is tag, which provides a string identifying the experiment:

class ObjectNavThorPPOExperimentConfig(ExperimentConfig):
    ...
    @classmethod
    def tag(cls):
        return "ObjectNavThorPPO"
    ...

Model creation#

Next, create_model will be used to instantiate an baseline object navigation actor-critic model:

class ObjectNavThorExperimentConfig(ExperimentConfig):
    ...

    # A simple setting, train/valid/test are all the same single scene
    # and we're looking for a single object
    OBJECT_TYPES = ["Tomato"]
    TRAIN_SCENES = ["FloorPlan1_physics"]
    VALID_SCENES = ["FloorPlan1_physics"]
    TEST_SCENES = ["FloorPlan1_physics"]

    # Setting up sensors and basic environment details
    SCREEN_SIZE = 224
    SENSORS = [
        RGBSensorThor(
            height=SCREEN_SIZE, width=SCREEN_SIZE, use_resnet_normalization=True,
        ),
        GoalObjectTypeThorSensor(object_types=OBJECT_TYPES),
    ]

    ...

    @classmethod
    def create_model(cls, **kwargs) -> nn.Module:
        return ObjectNavBaselineActorCritic(
            action_space=gym.spaces.Discrete(len(ObjectNavTask.class_action_names())),
            observation_space=SensorSuite(cls.SENSORS).observation_spaces,
            goal_sensor_uuid="goal_object_type_ind",
            hidden_size=512,
            object_type_embedding_dim=8,
        )
    ...

Training pipeline#

We now implement a training pipeline which trains with a single stage using PPO.

In the below we use Builder objects, which allow us to defer the instantiation of objects of the class passed as their first argument while allowing passing additional keyword arguments to their initializers. This is necessary when instantiating things like PyTorch optimizers who take as input the list of parameters associated with our agent's model (something we can't know until the create_model function has been called).

class ObjectNavThorPPOExperimentConfig(ExperimentConfig):
    ...
    @classmethod
    def training_pipeline(cls, **kwargs):
        ppo_steps = int(1e6)
        lr = 2.5e-4
        num_mini_batch = 2 if not torch.cuda.is_available() else 6
        update_repeats = 4
        num_steps = 128
        metric_accumulate_interval = cls.MAX_STEPS * 10  # Log every 10 max length tasks
        save_interval = 10000
        gamma = 0.99
        use_gae = True
        gae_lambda = 1.0
        max_grad_norm = 0.5

        return TrainingPipeline(
            save_interval=save_interval,
            metric_accumulate_interval=metric_accumulate_interval,
            optimizer_builder=Builder(optim.Adam, dict(lr=lr)),
            num_mini_batch=num_mini_batch,
            update_repeats=update_repeats,
            max_grad_norm=max_grad_norm,
            num_steps=num_steps,
            named_losses={
                "ppo_loss": PPO(clip_decay=LinearDecay(ppo_steps), **PPOConfig),
            },
            gamma=gamma,
            use_gae=use_gae,
            gae_lambda=gae_lambda,
            advance_scene_rollout_period=cls.ADVANCE_SCENE_ROLLOUT_PERIOD,
            pipeline_stages=[
                PipelineStage(loss_names=["ppo_loss"], max_stage_steps=ppo_steps,),
            ],
            lr_scheduler_builder=Builder(
                LambdaLR, {"lr_lambda": LinearDecay(steps=ppo_steps)}
            ),
        )
    ...

Alternatively, we could use a more sophisticated pipeline that begins training with dataset aggregation (DAgger) before moving to training with PPO. This requires the existence of an expert (implemented in the task definition) that provides optimal actions to agents. We have implemented such a pipeline by extending the above configuration as follows

class ObjectNavThorDaggerThenPPOExperimentConfig(ObjectNavThorPPOExperimentConfig):
    ...
    SENSORS = [
        RGBSensorThor(
            height=SCREEN_SIZE, width=SCREEN_SIZE, use_resnet_normalization=True,
        ),
        GoalObjectTypeThorSensor(object_types=OBJECT_TYPES),
        ExpertActionSensor(nactions=6), # Notice that we have added an expert action sensor.
    ]
    ...
    @classmethod
    def training_pipeline(cls, **kwargs):
        dagger_steps = int(1e4) # Much smaller number of steps as we're using imitation learning
        ppo_steps = int(1e6)
        lr = 2.5e-4
        num_mini_batch = 1 if not torch.cuda.is_available() else 6
        update_repeats = 4
        num_steps = 128
        metric_accumulate_interval = cls.MAX_STEPS * 10  # Log every 10 max length tasks
        save_interval = 10000
        gamma = 0.99
        use_gae = True
        gae_lambda = 1.0
        max_grad_norm = 0.5

        return TrainingPipeline(
            save_interval=save_interval,
            metric_accumulate_interval=metric_accumulate_interval,
            optimizer_builder=Builder(optim.Adam, dict(lr=lr)),
            num_mini_batch=num_mini_batch,
            update_repeats=update_repeats,
            max_grad_norm=max_grad_norm,
            num_steps=num_steps,
            named_losses={
                "ppo_loss": PPO(clip_decay=LinearDecay(ppo_steps), **PPOConfig),
                "imitation_loss": Imitation(), # We add an imitation loss.
            },
            gamma=gamma,
            use_gae=use_gae,
            gae_lambda=gae_lambda,
            advance_scene_rollout_period=cls.ADVANCE_SCENE_ROLLOUT_PERIOD,
            pipeline_stages=[ # The pipeline now has two stages, in the first
                              # we use DAgger (imitation loss + teacher forcing).
                              # In the second stage we no longer use teacher
                              # forcing and add in the ppo loss.
                PipelineStage(
                    loss_names=["imitation_loss"],
                    teacher_forcing=LinearDecay(
                        startp=1.0, endp=0.0, steps=dagger_steps,
                    ),
                    max_stage_steps=dagger_steps,
                ),
                PipelineStage(loss_names=["ppo_loss"], max_stage_steps=ppo_steps,),
            ],
            lr_scheduler_builder=Builder(
                LambdaLR, {"lr_lambda": LinearDecay(steps=ppo_steps)}
            ),
        )

A version of our experiment config file for which we have implemented this two-stage training can be found here. This two-stage configuration ObjectNavThorDaggerThenPPOExperimentConfig is actually implemented as a subclass of ObjectNavThorPPOExperimentConfig. This is a common pattern used in AllenAct and lets one skip a great deal of boilerplate when defining a new experiment as a slight modification of an old one. Of course one must then be careful: changes to the superclass configuration will propagate to all subclassed configurations.

Machine configuration#

In machine_params we define machine configuration parameters that will be used for training, validation and test:

class ObjectNavThorPPOExperimentConfig(core.base_abstractions.experiment_config.ExperimentConfig):
    ...
    @classmethod
    def machine_params(cls, mode="train", **kwargs):
        num_gpus = torch.cuda.device_count()
        has_gpu = num_gpus != 0 

        if mode == "train":
            nprocesses = 20 if has_gpu else 4
            gpu_ids = [0] if has_gpu else []
        elif mode == "valid":
            nprocesses = 1
            gpu_ids = [1 % num_gpus] if has_gpu else []
        elif mode == "test":
            nprocesses = 1
            gpu_ids = [0] if has_gpu else []
        else:
            raise NotImplementedError("mode must be 'train', 'valid', or 'test'.")

        return {"nprocesses": nprocesses, "gpu_ids": gpu_ids}
    ...

In the above we use the availability of cuda (torch.cuda.device_count() != 0) to determine whether we should use parameters appropriate for local machines or for a server. We might optionally add a list of sampler_devices to assign devices (likely those not used for running our agent) to task sampling workers.

Task sampling#

The above has defined the model we'd like to use, the types of losses we wish to use during training, and the machine specific parameters that should be used during training. Critically we have not yet defined which task we wish to train our agent to complete. This is done by implementing the ExperimentConfig.make_sampler_fn function

class ObjectNavThorPPOExperimentConfig(ExperimentConfig):
    ...
    @classmethod
    def make_sampler_fn(cls, **kwargs) -> TaskSampler:
        return ObjectNavTaskSampler(**kwargs)
    ...

Now, before training starts, our trainer will know to generate a collection of task samplers using make_sampler_fn for training (and possibly validation or testing). The kwargs parameters used in the above function call can be different for each training process, we implement such differences using the ExperimentConfig.train_task_sampler_args function

class ObjectNavThorPPOExperimentConfig(ExperimentConfig):
    ...
    def train_task_sampler_args(
        self,
        process_ind: int,
        total_processes: int,
        devices: Optional[List[int]] = None,
        seeds: Optional[List[int]] = None,
        deterministic_cudnn: bool = False,
    ) -> Dict[str, Any]:
        res = self._get_sampler_args_for_scene_split(
            self.TRAIN_SCENES,
            process_ind,
            total_processes,
            seeds=seeds,
            deterministic_cudnn=deterministic_cudnn,
        )
        res["scene_period"] = "manual"
        res["env_args"] = {}
        res["env_args"].update(self.ENV_ARGS)
        res["env_args"]["x_display"] = (
            ("0.%d" % devices[process_ind % len(devices)]) if len(devices) > 0 else None
        )
        return res
    ...

Now training process i out of n total processes will be instantiated with the parameters ObjectNavThorPPOExperimentConfig.train_task_sampler_args(i, n, ...). Similar functions (valid_task_sampler_args and test_task_sampler_args) exist for generating validation and test parameters. Note also that with this function we can assign devices to run our environment for each worker. See the documentation of ExperimentConfig for more information.

Running the experiment#

We are now in the position to run the experiment (with seed 12345) using the command

python main.py object_nav_ithor_ppo_one_object -b projects/tutorials -s 12345