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: