空中懸停、翻滾轉(zhuǎn)身、成功著陸，我用強化學(xué)習(xí)「回收」了SpaceX的火箭

520jefferson 2021-11-14

展開全文

機器之心報道

編輯：杜偉、陳萍

我自己造了個「火箭」，還把它回收了。

SpaceX 作為一家太空探索技術(shù)公司是美國一家民營航天制造商和太空運輸公司，由伊隆 · 馬斯克于 2002 年創(chuàng)辦，目標是降低太空運輸?shù)某杀荆⑦M行火星探索。SpaceX 成立近 20 年以來，吸引了無數(shù)的火箭愛好者。

對于個人來說，怎樣實現(xiàn)這個火箭夢呢？據(jù)了解，SpaceX 制造一枚獵鷹 9 號的費用實際在 3040 萬美元左右。對于個人來說，這個費用簡直是天方夜譚，更別說涉及到的技術(shù)等問題了。

有困難就要想辦法解決，作為 SpaceX 的超級粉絲，來自密歇根大學(xué)安娜堡分校的博士后研究員 Zhengxia Zou 也是個火箭迷，一直夢想擁有自己的火箭。最近，他研究了一個有趣的問題，即我們是否可以「建造」一個虛擬火箭，并通過強化學(xué)習(xí)解決火箭回收這個具有挑戰(zhàn)性的問題。在實驗中，Zou 嘗試了關(guān)于火箭懸停和降落的兩個任務(wù)。

由于這是 Zou 的第一個強化學(xué)習(xí)項目，包括環(huán)境、火箭動力學(xué)、RL 智能體等，Zou 表示盡量從頭開始實現(xiàn)所有內(nèi)容，并希望通過這些底層的編碼，能夠?qū)娀瘜W(xué)習(xí)有更深入的了解，包括基礎(chǔ)算法，智能體與環(huán)境的交互，以及獎勵的設(shè)計。

項目主頁：https://jiupinjia./rocket-recycling/
GitHub 地址：https://github.com/jiupinjia/rocket-recycling

不過對于這個項目，也有網(wǎng)友提出質(zhì)疑：「如果我們能夠用經(jīng)典的控制方法來完成這項任務(wù)，為什么在 SpaceX 之前沒有人做過呢？」對于這一質(zhì)疑，有網(wǎng)友表示：「SpaceX 沒有使用強化學(xué)習(xí)，他們使用論文《 Lossless Convexification of Nonconvex Control Bound and Pointing Constraints of the Soft Landing Optimal Control Problem 》中的方法來解決火箭著陸問題，性能優(yōu)于 RL。」

實現(xiàn)懸停和降落的智能體以及環(huán)境

Zou 嘗試了懸停和降落這兩個任務(wù)。如下圖所示，火箭被簡化成二維平面上的剛體，并且考慮了基本圓柱體動力學(xué)模型，并假設(shè)空氣阻力與速度成正比。火箭底部安裝了一個推力矢量發(fā)動機，該發(fā)動機為火箭提供不同方向的推力值 (0.2g, 1.0g, 2.0g)。噴嘴添加角速度約束，角速度最大旋轉(zhuǎn)速度為 30 度 / 秒。

通過上述基本設(shè)置，動作空間（action space）被定義為發(fā)動機離散控制信號的集合，包括噴管的推力加速度和角速度。狀態(tài)空間（state-space）由火箭位置、速度、角度、角速度、噴嘴角度和模擬時間組成。

對于著陸任務(wù)，Zou 遵循了 Starship SN10 的 Belly Flop 式降落（腹部朝下的翻轉(zhuǎn)動作以實現(xiàn)空中轉(zhuǎn)身）的基本參數(shù)。初始速度設(shè)置為 - 50m/s，火箭方向設(shè)置為 90 度（水平），著陸 burn height 設(shè)置為離地 500 米。

圖源：https://twitter.com/thejackbeyer/status/1367364251233497095

獎勵函數(shù)非常簡單。

對于懸停任務(wù)：基于兩個規(guī)則給出 step-reward：1）火箭與預(yù)定義目標點之間的距離——它們越近，分配的獎勵越大；2）火箭機身的角度（盡量保持直立）。

對于著陸任務(wù)：觀察看觸地瞬間的速度和角度——當(dāng)觸地速度小于安全閾值并且角度接近 0 度（直立）時，則認為它是成功著陸并獲得豐厚獎勵。其余規(guī)則與懸停任務(wù)相同。

除了進行上述任務(wù)外，有網(wǎng)友表示，「希望將燃料的研究也納入進來，燃料的有效使用是火箭技術(shù)的主要關(guān)注點，因此看到有效著陸?yīng)剟顚⑹且患檬隆ｉ_始時火箭有給定的燃料量，剩余的燃料量成為學(xué)習(xí)者的另一個信號——既用于控制也作為獎勵。」

實現(xiàn)效果

Zou 實現(xiàn)了上述環(huán)境，并且訓(xùn)練了一個基于策略的智能體（actor-critic）來解決這個問題。在超過 20,000 個訓(xùn)練 episodes 之后，episode 獎勵最終收斂地非常好。

下圖左為懸停任務(wù)上不同訓(xùn)練 episode 數(shù)量時的獎勵；圖右為著陸任務(wù)上不同 episode 數(shù)量時的獎勵。

以下幾個動圖分別展示了經(jīng)過不同訓(xùn)練 episode 后學(xué)習(xí)到的 RL 行為：

左：訓(xùn)練 episode 為 0（隨機智能體），右：訓(xùn)練 episode 為 100。

左：訓(xùn)練 episode 為 2,000，右：訓(xùn)練 episode 為 10,000。

20,000 個訓(xùn)練 episode 之后，左：執(zhí)行懸停任務(wù)的完全訓(xùn)練智能體，右：執(zhí)行著陸任務(wù)的完全訓(xùn)練智能體。

與 SN10 著陸的動效對比

盡管環(huán)境和獎勵的設(shè)置很簡單，但經(jīng)過訓(xùn)練的智能體已經(jīng)很好地學(xué)會了 Belly Flop 式降落。

如下動圖展示了真實的 Starship SN10 和從強化學(xué)習(xí)中學(xué)到的智能體在著陸時的比較：

智能體訓(xùn)練與測試

訓(xùn)練智能體，需要./example_train.py。

測試智能體的流程如下：

import torchfrom rocket import Rocketfrom policy import ActorCriticimport osimport glob

# Decide which device we want to run ondevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

if __name__ == '__main__':

task = 'hover' # 'hover' or 'landing' max_steps = 800 ckpt_dir = glob.glob(os.path.join(task+'_ckpt', '*.pt'))[-1] # last ckpt env = Rocket(task=task, max_steps=max_steps) net = ActorCritic(input_dim=env.state_dims, output_dim=env.action_dims).to(device) if os.path.exists(ckpt_dir): checkpoint = torch.load(ckpt_dir) net.load_state_dict(checkpoint['model_G_state_dict'])

state = env.reset() for step_id in range(max_steps): action, log_prob, value = net.get_action(state) state, reward, done, _ = env.step(action) env.render(window_name='test') if env.already_crash: break

作者簡介

項目作者 Zhengxia Zou 現(xiàn)為密歇根大學(xué)安娜堡分校的博士后研究員，此前先后于 2013 年和 2018 年取得北京航空航天大學(xué)的學(xué)士和博士學(xué)位。他的主要研究興趣包括計算機視覺及其在遙感、自動駕駛汽車和電子游戲等領(lǐng)域的相關(guān)應(yīng)用。

谷歌學(xué)術(shù)主頁：https://scholar.google.com/citations?user=DzwoyZsAAAAJ&hl=en

Zhengxia Zou 參與撰寫的論文被 AAAI、CVPR、ICCV、IJCAI、ACM MM 等多個學(xué)術(shù)頂會接收。他還曾擔(dān)任 NeurIPS、AAAI、ACCV 和 WACV 等多個學(xué)術(shù)會議的程序委員，以及 ICLR 會議及 IEEE Transactions on Image Processing 等多份期刊的審稿人。

機器之心此前報道過多篇他參與的研究，包括如下：