Saved in:
Bibliographic Details
Main Authors: Dai, Yang, Ma, Oubo, Zhang, Longfei, Liang, Xingxing, Hu, Shengchao, Wang, Mengzhu, Ji, Shouling, Huang, Jincai, Shen, Li
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2405.12094
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914991477620736
author Dai, Yang
Ma, Oubo
Zhang, Longfei
Liang, Xingxing
Hu, Shengchao
Wang, Mengzhu
Ji, Shouling
Huang, Jincai
Shen, Li
author_facet Dai, Yang
Ma, Oubo
Zhang, Longfei
Liang, Xingxing
Hu, Shengchao
Wang, Mengzhu
Ji, Shouling
Huang, Jincai
Shen, Li
contents Transformer-based trajectory optimization methods have demonstrated exceptional performance in offline Reinforcement Learning (offline RL). Yet, it poses challenges due to substantial parameter size and limited scalability, which is particularly critical in sequential decision-making scenarios where resources are constrained such as in robots and drones with limited computational power. Mamba, a promising new linear-time sequence model, offers performance on par with transformers while delivering substantially fewer parameters on long sequences. As it remains unclear whether Mamba is compatible with trajectory optimization, this work aims to conduct comprehensive experiments to explore the potential of Decision Mamba (dubbed DeMa) in offline RL from the aspect of data structures and essential components with the following insights: (1) Long sequences impose a significant computational burden without contributing to performance improvements since DeMa's focus on sequences diminishes approximately exponentially. Consequently, we introduce a Transformer-like DeMa as opposed to an RNN-like DeMa. (2) For the components of DeMa, we identify the hidden attention mechanism as a critical factor in its success, which can also work well with other residual structures and does not require position embedding. Extensive evaluations demonstrate that our specially designed DeMa is compatible with trajectory optimization and surpasses previous methods, outperforming Decision Transformer (DT) with higher performance while using 30\% fewer parameters in Atari, and exceeding DT with only a quarter of the parameters in MuJoCo.
format Preprint
id arxiv_https___arxiv_org_abs_2405_12094
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?
Dai, Yang
Ma, Oubo
Zhang, Longfei
Liang, Xingxing
Hu, Shengchao
Wang, Mengzhu
Ji, Shouling
Huang, Jincai
Shen, Li
Machine Learning
Transformer-based trajectory optimization methods have demonstrated exceptional performance in offline Reinforcement Learning (offline RL). Yet, it poses challenges due to substantial parameter size and limited scalability, which is particularly critical in sequential decision-making scenarios where resources are constrained such as in robots and drones with limited computational power. Mamba, a promising new linear-time sequence model, offers performance on par with transformers while delivering substantially fewer parameters on long sequences. As it remains unclear whether Mamba is compatible with trajectory optimization, this work aims to conduct comprehensive experiments to explore the potential of Decision Mamba (dubbed DeMa) in offline RL from the aspect of data structures and essential components with the following insights: (1) Long sequences impose a significant computational burden without contributing to performance improvements since DeMa's focus on sequences diminishes approximately exponentially. Consequently, we introduce a Transformer-like DeMa as opposed to an RNN-like DeMa. (2) For the components of DeMa, we identify the hidden attention mechanism as a critical factor in its success, which can also work well with other residual structures and does not require position embedding. Extensive evaluations demonstrate that our specially designed DeMa is compatible with trajectory optimization and surpasses previous methods, outperforming Decision Transformer (DT) with higher performance while using 30\% fewer parameters in Atari, and exceeding DT with only a quarter of the parameters in MuJoCo.
title Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?
topic Machine Learning
url https://arxiv.org/abs/2405.12094