Differentiable Temporal Logic for Temporal Action Segmentation

This repository contains the software used in our NeurIPS 2022 Paper Don't Pour Cereal into Coffee: Differentiable Temporal Logic for Temporal Action Segmentation (➡️ project page).

Setup

The codebase is built on PyTorch 1.10 and Ubuntu 20.04. In principle it should work on PyTorch>=1.0.

1. Clone this repository.
2. Install dependencies listed in requirements.txt. You will need to install PyTorch if you haven't done so.
3. Download dataset tarballs here. Decompress them into dataset/.

Basic Usage

python main.py -f ${CONFIG_FILE} [-c <KEY1 VAL1> [KEY2 VAL2]...]

To run a predefined experiment, replace ${CONFIG_FILE} with one of the yml files in config/. For example, python main.py -f config/gru_50salads.yml. Please read the corresponding yml file for details of the experiment.

Overriding Defaults

Key-value pairs following -c can be used to override configs. For example, to run the experiment with a different dataset split (say split 2), use python main.py -f config/gru_50salads.yml -c ds_params.commons.split 2.

If you feel like having fun and want to run experiments using other parameters, please check config/defaults.py for a documented list of config keys. You will need to write a yacs config file like those config/*.yml. In particular, there are two important parameters: the weight for logic loss (loss_weights.lg, corresponding to $\lambda$ in the paper) and the smoothing factor (extra_kwargs.rho, corresponding to $\gamma$ in the paper).

Test a Model

To test a trained model, override test_only and load_weight_from options, for example:

python main.py -f config/gru_50salads.yml -c test_only True load_weight_from $PATH_TO_SNAPSHOT/best.state

Use DTL in Your Projects

It is straightforward to incorporate DTL into your own projects. We suggest the following workflow, which would introduce a small modification to your existing codebase.

1. Build constraints and compile them into a formula. Please check out notebooks/rule_extractor.ipynb for how it was done in our work. The process could be different in your specific case.
2. In your training loop, add components for DTL. For example:

from lib.tl import parse

# Existing codebase for model and optimizer initialization...
model = Model()
# omitted ...

# Get ready the logic evaluator
## Create a formula
formula_str = '(F pour_coffee & F boil_water) -> (~pour_coffee W boil_water)'
## Parse the formula to get an evaluator
evaluator = parse(formula_str)
## An ap_map tells the evaluator how to map the name of a proposition to an 
## index, through which it can access the logits of that proposition.
list_of_props = ['pour_coffee', 'boil_water']
ap_map = lambda x: list_of_props.index(x)
## rho is the parameter used by _smin() and _smax() for smooth min() and max()
## Check lib.tl.evaluator._smin and lib.tl.evaluator._smax for details
rho = 1
logic_weight = 0.1

# The training loop
while epoch in range(max_epoch):
    for data in dataloader:
        X, Y = data
        # Y_ should be shaped (batch, num_propositions, time_steps)
        Y_ = model(X)
        
        # Compute the task loss, e.g., nll_loss
        task_loss = task_loss_function(Y, Y_, *args, **kwargs)

        # Compute the logic loss
        tl_score = evaluator(Y_, ap_map=ap_map, rho=rho)
        logic_loss = torch.log(1+torch.exp(-tl_score))
        
        # Combine loss functions
        (task_loss + logic_weight * logic_loss).backward()

        # continue the training loop ...

Credits

The formula parsing code is inspired by mvcisback/py-metric-temporal-logic.

Some data handling and evaluation code are adapted from yiskw713/asrf and ptirupat/MLAD.

We would like to thank the authors for their contribution.

About

Please visit our project page for more details.

Feel free to contact me should you have any questions.

If you find this repository useful in your research, please consider citing our paper:

@inproceedings{xu2022difftl,
      title={Don't Pour Cereal into Coffee: Differentiable Temporal Logic for Temporal Action Segmentation}, 
      author={Ziwei Xu and Yogesh S Rawat and Yongkang Wong and Mohan S Kankanhalli and Mubarak Shah},
      booktitle={{NeurIPS}},
      year={2022}
}

License

MIT License.