Focus-Then-Decide: Segmentation-Assisted Reinforcement Learning

This respiratory contains the code for AAAI-2024 accepted paper: Focus-Then-Decide: Segmentation-Assisted Reinforcement Learning

Project page    Paper

Installation

You can directly install the dependencies using the provided setup.sh file.

./setup.sh

Run

Below are running commands of running FTD and baseline algorithms. Please replace algo, env, and task with corresponding arguments (e.g. ./scripts/run.sh ftd franka reach).

./scripts/run.sh algo env task

If everything goes well, this should yield an output of the form

Working directory: logs/franka_reach/ftd/20231228-170531
Observation space: (81, 84, 84)
Action space: (8,)
=====Start training=====
Evaluating: logs/franka_reach/ftd/20231228-170531
| eval | S: 0 | ERTEST: 1.8151e+01
| train | E: 1 | S: 250 | D: 93.8 s | R: 0.0000e+00 | ALOSS: 0.0000e+00 | CLOSS: 0.0000e+00 | RPredLOSS: 0.0000e+00 | APredLOSS: 0.0000e+00
| train | E: 2 | S: 500 | D: 98.4 s | R: 8.2735e+00 | ALOSS: 0.0000e+00 | CLOSS: 0.0000e+00 | RPredLOSS: 0.0000e+00 | APredLOSS: 0.0000e+00

To run the ablation study, please refer to ./scripts/run_ablation.sh

For more parameter settings, please refer to arguments.py.

If you encounter problems related to rendering, please refer to respiratory of DMC

Visualization

pendulum-swingup

cartpole-swingup

finger-spin

hopper-stand

hopper-hop

cheetah-run

walker-walk

walker-run

franka-reach

Citation

If you find our work useful in your research, please consider citing our work as follows:

@article{Chen_Focus_2024,
    author={Chen, Chao and Xu, Jiacheng and Liao, Weijian and Ding, Hao and Zhang, Zongzhang and Yu, Yang and Zhao, Rui},
    title={Focus-Then-Decide: Segmentation-Assisted Reinforcement Learning}, 
    journal={Proceedings of the AAAI Conference on Artificial Intelligence},
    volume={38}, 
    number={10}, 
    year={2024},
    pages={11240-11248},
    DOI={10.1609/aaai.v38i10.29002},
    url={https://ojs.aaai.org/index.php/AAAI/article/view/29002}
}