Small Language Models for Task Planning in Robotics

[ Download Command-steps pairs Dataset ]

• Tabletop Dataset Download
• Kitchen Dataset Download

[ Dataset description ]

We build datasets to finetune small LMs for task planning in a single domain, such as kitchen and tabletop environments, via knowledge distillation from GPT3.5, by CoT prompting on it. Our dataset consists of high-level commands and corresponding low-level actionable steps. We name our dataset as COST, COmmand-STeps dataset.

[ Note that there are two versions of the prompt templates. ]

To allow users to build a dataset for their domain, we propose prompt templates, which are .txt files in the 'code/prompt_templates/fixed_objects or flexible_objects' directory, that can be adaptable to any environment and conditions.

1. Prompt for fixed objects

The prompt template when the objects to be used in action steps are fixed. It is designed to generate action steps using only given objects. We use this prompt template to build the COST dataset for a tabletop domain.

2. Prompt for flexible objects

The prompt template when the objects to be used in action steps are flexible. It is designed to generate not only the action steps but also the objects required by the action steps, for the input command. For general real-world environments, such as kitchen, we use this prompt template.

---

Let's implement a task planner on your domain, utilizing small LMs.

First, Generate your own COmmand-STeps Dataset (COST)

If you already have command-steps dataset for fine-tuning LMs, you can skip this part.

• The case for fixed objects

  1. Save your object list on 'code/datasets/{}/objects.json'.format(your_domain) path.
  2. Once the object list is ready, enter your OPENAI key into Generate_dataset(fixed_objs).ipynb, and run it to build the COST dataset.
  3. Run 'code/Postprocessing_daraset(fixed_objs).ipynb' to post-process the LLMs output and finish building the COST dataset.

• The case for flexible objects

  1. Input the prompt on 'code/prompt_templates/flexible_objects/object_generate_prompt.txt' into LLMs (i.e., GPT3.5), and generate the initial objects list.
  2. Once the object list is ready, enter your OPENAI key into 'code/Generate_dataset(flexible_objs).ipynb', and run it to build the COST dataset.
  3. Run 'code/Postprocessing_daraset(flexible_objs).ipynb' to post-process the LLMs output and finish building the COST dataset.

Second, Fine-tune small LMs with your COST dataset

• Run 'code/Finetuning_gpt2.ipynb' to fine-tune small LMs, such as GPT2-medium and GPT2-base.

Lastly, Test the fine-tuned small LMs

• Run 'code/Test.ipynb' to test your fine-tuned small LMs.