Skill Transformer: A Monolithic Policy for Mobile Manipulation

This is the official PyTorch code implementation for "Skill Transformer: A Monolithic Policy for Mobile Manipulation".

Installation

Refer to Habitat-Lab repo for installation.

• Requires Python >= 3.7: conda create -y -n py37 python=3.7

Important Config Options

Important hyperparameters and options can be found in habitat_baselines/config/rearrange/hab/transformer.yaml

• NUM_ENVIRONMENTS: Num of environments used in evaluation.
• CHECKPOINT_INTERVAL: Frequency to checkpoint the model during training
• TEST_INTERVAL: Frequency to run through validation dataset
• RL.POLICY: Options for the policy class
  • ACtION_DIST: Options to experiment with discrete actions and continuous actions. We found that continuous actions work better in our case.
  • train_planner and train_control: Options to only train skill inference module or action inference module respectively.
• RL.TRAJECTORY_DATASET: Options for the training dataset
  • trajectory_dir: Specify the dataset location.
  • dataset_size: Specify how many episodes to use.
  • files_per_load: Specify how many files (episodes) to load in each epoch. More files means a longer loading time.
  • queue_size: How many epochs of data to load in advance and store in memory
• RL.VALIDATION_DATASET: Options for including a validation dataset to check for overfitting in training.
• TRANSFORMER: Hyperparameters

Collect dataset with DDPPO expert (Habitat 2.0)

• Run python habitat_baselines/run.py --exp-config habitat_baselines/config/rearrange/hab/tp_srl_oracle_plan.yaml --run-type eval
• Add following arguments:
  • TEST_EPISODE_COUNT num_episodes
  • DATASET_SAVE_PATH path_to_your_dataset
  • TASK_CONFIG.SEED seed TASK_CONFIG.SIMULATOR.SEED seed

Train Skill Transformer

• Add the path to dataset in the config file
• Run python habitat_baselines/run.py --exp-config habitat_baselines/config/rearrange/hab/transformer.yaml --run-type train
• Add following arguments:
  • TASK_CONFIG.SEED seed TASK_CONFIG.SIMULATOR.SEED seed

Evaluate Skill Transformer

• Run python habitat_baselines/run.py --exp-config habitat_baselines/config/rearrange/hab/transformer.yaml --run-type eval
• Add following arguments:
  • TASK_CONFIG.SEED seed TASK_CONFIG.SIMULATOR.SEED seed

Code Structure

Structure of the code under habitat_baselines/transformer_policy:

• dataset_utils: Helper functions to load in dataset.
• dataset: The rolling dataset class that loads in dataset and provides data to the trainer in parallel.
• focal_loss: The focal loss implementation.
• action_distribution: Used when some actions are categorical and some are normally distributed.
• transformer_model: Contains the action inference module and the skill inference module.
• transformer_policy: The policy class that specifies the visual encoder, the agent, and the loss functions.
• transformer_trainer: The trainer class that specifies

Citation

@article{huang2023skill,
  title={Skill Transformer: A Monolithic Policy for Mobile Manipulation},
  author={Huang, Xiaoyu and Batra, Dhruv and Rai, Akshara and Szot, Andrew},
  journal={arXiv preprint arXiv:2308.09873},
  year={2023}
}

License

this code is licensed under the CC-BY-NC license, see LICENSE.md for more details