ゲームボーイ上で強化学習を行う part2

前回の続きです。前回はQ学習をPC上で動作させ、適切なハイパーパラメータを取得しました。

実装

浮動小数点数演算を回避する

先述のプログラムでは浮動小数点数（double型）の演算を用いていましたが、GBDKのCコンパイラではこれをサポートしていません。ゲームボーイのCPUには浮動小数点数を処理する機構がないためです。それどころか、ゲームボーイのCPUには整数の加算・減算命令はありますが乗算・除算命令はありません。GBDKには、整数の乗算・除算を加算・減算命令とビットシフト、条件分岐命令を組み合わせることで処理するルーチンが組み込まれており、ゲーム開発者は命令の有無を気にせずC言語の演算子を利用可能になっています。しかし、乗算は加算の数十倍の時間がかかりますので極力回避する必要があります。なお浮動小数点数演算も理論上はソフトウェア的に実現できますが、整数加算の数千倍の時間がかかるため実用性はないでしょう。浮動小数点数演算がないことに対する対策は、扱う数値を定数倍することで整数として処理できるようにすること、乗算除算の負荷が高いことに対する対策は、ビットシフト命令を用いることが解決策となります。

mountaincar.hの内容です。

#include <stdint.h>
#include "config.h"

typedef int16_t Reward;
typedef int16_t Position;
typedef int16_t Velocity;
typedef uint8_t MountainCarAction; // 0=left, 1=nothing, 2=right

typedef struct {
    Position position;
    Velocity velocity;
    uint8_t steps;
    uint8_t done;
} MountainCarState;

Reward MountainCarStep(MountainCarState *state, MountainCarAction action);
void MountainCarReset(MountainCarState *state);

報酬Reward、座標Positionなどを、16ビット整数(int16_t)で表現するように定義しています。例えば座標は本来[-1.2, 0.6]の範囲でしたが、10000倍して[-12000, 6000]の範囲を取る値として表現します。また、ゲームボーイは8bitCPUであるものの、GBDKでのintは16bit整数のため、8bitで十分な変数はuint8_t, int8_t型で定義することにより極力メモリ使用量を削減します。

mountaincar.cは以下の通りです。

#include <stdlib.h>
#include "config.h"
#include "mountaincar.h"

#define POSITION_SCALE 10000
#define FORCE 10
#define VELOCITY_MAX 700
#define VELOCITY_MIN -700
#define POSITION_MIN -12000
#define POSITION_MAX 6000
#define POSITION_DONE 5000

/*
raw_values = np.arange(512//2,18000+512,512)
(np.cos(3*(raw_values / scale - 1.2))*-0.0025*10000).astype(np.int8)
*/
static const int8_t gravity_table[] = {
    23,  24,  24,  24,  24,  23,  21,  19,  16,  13,  10,   6,   2,
        -1,  -4,  -8, -12, -15, -18, -20, -22, -23, -24, -24, -24, -23,
       -22, -20, -17, -14, -11,  -8,  -4,   0,   3, 6
};

static Velocity gravity(Position position) {
    // cos(3 * position) * (-0.0025);
    return gravity_table[(position - POSITION_MIN) >> 9];
}

Reward MountainCarStep(MountainCarState *state, MountainCarAction action) {
    Position position = state->position;
    Velocity velocity = state->velocity;
    switch (action) {
        case 0:
            velocity -= FORCE;
            break;
        case 2:
            velocity += FORCE;
            break;
    }
    velocity += gravity(position);
    if (velocity > VELOCITY_MAX) {
        velocity = VELOCITY_MAX;
    } else if (velocity < VELOCITY_MIN) {
        velocity = VELOCITY_MIN;
    }
    position += velocity;
    if (position > POSITION_MAX) {
        position = POSITION_MAX;
        velocity = 0;
    } else if (position < POSITION_MIN) {
        position = POSITION_MIN;
        velocity = 0;
    }
    state->position = position;
    state->velocity = velocity;
    state->steps += 1;


    Reward reward = REWARD_PER_STEP;

    if (position >= POSITION_DONE) {
        reward += REWARD_AT_GOAL;
    }

    if (position >= POSITION_DONE || state->steps >= MAX_EPISODE_LEN) {
        state->done = 1;
    }

    return reward;
}

void MountainCarReset(MountainCarState *state) {
    // uniform random between -0.6 and -0.4
    // (0.2 * rand() / RAND_MAX - 0.6) * 10000;
    
#ifdef __SDCC
    // randw(): 0 to 0xffff
    // randw()>>5: 0 to 2047
    state->position = (int16_t)(randw() >> 5) - 6000;
#else
    // rand(): 0 to 0x7fff
    // rand()>>4: 0 to 2047
    state->position = (rand() >> 4) - 6000;
#endif
    state->velocity = 0;
    state->steps = 0;
    state->done = 0;
}

まず、エージェントが与える入力によって速度が変化する項はPC版では(action - 1) * 0.001という乗算を用いていましたが、switch文により定数を加算・減算するように変更します。

    switch (action) {
        case 0:
            velocity -= FORCE;
            break;
        case 2:
            velocity += FORCE;
            break;
    }

重力によって速度が変化する項は、PC版ではcos(3 * position) * (-0.0025)という三角関数を用いた式になっていましたが、GBDKでは三角関数は搭載されていませんので独自に実装する必要があります。しかし、三角関数を直接実装することはこのタスクでは重要ではなく、そのあと-0.0025を乗算する処理が入るため、乗算も回避すべきです。そのため、cos(3 * position) * (-0.0025)を一括で計算するgravity関数に切り出し、極力低コストな実装を考えます。その実装は以下の通りです。

static const int8_t gravity_table[] = {
    23,  24,  24,  24,  24,  23,  21,  19,  16,  13,  10,   6,   2,
        -1,  -4,  -8, -12, -15, -18, -20, -22, -23, -24, -24, -24, -23,
       -22, -20, -17, -14, -11,  -8,  -4,   0,   3, 6
};

static Velocity gravity(Position position) {
    return gravity_table[(position + 12000) >> 9];
}

// 使用方法: velocity += gravity(position);

Position, VelocityはPC倍の値の10,000倍にスケーリングされていることに注意してください。まず、position + 12000により、負の数を回避し値を[0, 18000]の範囲に収めます。9ビット右シフトにより、数値を512 (2の9乗)で割ったのと同じ効果が得られるため、値が[0, 35]の範囲に変換されます。各値に対応するgravityの値を、gravity_tableに定数として埋め込んでおき、配列参照によって取り出します。この定数は、以下のPythonコードにより事前に計算しました。

raw_values = np.arange(512//2,18000+512,512)
(np.cos(3*(raw_values / scale - 1.2))*-0.0025*10000).astype(np.int8)

本来のPositionは18000通りの値を取りえますが、テーブルの要素数は36であり精度が失われています。そのためdouble版と厳密に同じ計算結果にはならないことに留意してください。

車の初期位置を乱数で設定する箇所では、プリプロセッサマクロによる分岐を行っています。GBDKでのコンパイル時には__SDCCというマクロが定義されているため、これによりPC用ビルドと区別します。GBDKでの乱数生成はrandw()関数を用います。これはuint16_t型で[0, 65535]の値を返します。車の座標は[-6000,-4000]の範囲にしてほしいので、5ビット右シフトして[0, 2047]の範囲に変換し、6000を引くことで[-6000,-3953]の範囲の値を得ます。ビットシフトで実装するために、若干本来の範囲とずれることを許容しています。

#ifdef __SDCC
    // randw(): 0 to 0xffff
    // randw()>>5: 0 to 2047
    state->position = (int16_t)(randw() >> 5) - 6000;
#else
    // rand(): 0 to 0x7fff
    // rand()>>4: 0 to 2047
    state->position = (rand() >> 4) - 6000;
#endif

次にQ学習アルゴリズムの実装です。qlearning.hは以下の通りです。

#include "mountaincar.h"

typedef struct {
    Reward q_table[N_STATE_0 * N_STATE_1][N_ACTIONS];
} QLearningState;

QLearningState *QLearningStateCreate();
MountainCarAction GetBestAction(const QLearningState *q_state, const MountainCarState* state);
Reward TestEpisode(const QLearningState *q_state, uint8_t *steps);
Reward TrainEpisode(QLearningState *q_state);

QLearningStateにはQテーブルが含まれていますが、Reward (=uint16_t)型の要素を256(状態数)×3(行動数)個保持しており、1536バイトを占めます。Qテーブルの値は学習中に少しずつ変化させていく必要があり、8ビットでは粗すぎるため16ビット使用します。Qテーブルは今回のプログラムで最も容量の大きなデータ構造となります。ゲームボーイのRAM容量は8kBですので余裕をもって収めることができました。今回の環境は状態の軸が座標と速度の2つでしたが、仮に状態の軸がもう一つ増えると容量が厳しくなりそうです。

qlearning.cは以下の通りです。

#include <stdlib.h>
#include "config.h"
#include "mountaincar.h"
#include "qlearning.h"

QLearningState *QLearningStateCreate() {
    QLearningState *state = malloc(sizeof(QLearningState));
    for (int i = 0; i < N_STATE_0 * N_STATE_1; i++) {
        for (int j = 0; j < N_ACTIONS; j++) {
            state->q_table[i][j] = 0;
        }
    }
    return state;
}

// (np.linspace(-1.2,0.6,16+1)*10000)[1:-1].astype(np.int32)
static const Position position_thresholds[] = {-10875,  -9750,  -8625,  -7500,  -6375,  -5250,  -4125,  -3000,
        -1875,   -749,    374,   1499,   2624,   3749,   4874};
static const Velocity velocity_thresholds[] = {-612, -525, -437, -350, -262, -175,  -87,    0,   87,  175,  262,
        350,  437,  525,  612};

static int get_state_index(const MountainCarState* state)
{
    // int position = (int)((state->position / 10000.0 + 1.2) / 1.8 * N_STATE_0);
    // int velocity = (int)((state->velocity / 10000.0 + 0.07) / 0.14 * N_STATE_1);
    uint8_t position;
    for (position = 0; position < N_STATE_0 - 1; position++) {
        if (state->position < position_thresholds[position]) {
            break;
        }
    }
    uint8_t velocity;
    for (velocity = 0; velocity < N_STATE_1 - 1; velocity++) {
        if (state->velocity < velocity_thresholds[velocity]) {
            break;
        }
    }
    return position * N_STATE_1 + velocity;
}

MountainCarAction GetBestAction(const QLearningState *q_state, const MountainCarState* state) {
    MountainCarAction action = 0;
    int state_index = get_state_index(state);
    Reward max_q = q_state->q_table[state_index][0];
    for (uint8_t i = 1; i < N_ACTIONS; i++) {
        if (q_state->q_table[state_index][i] > max_q) {
            max_q = q_state->q_table[state_index][i];
            action = i;
        }
    }
    return action;
}

Reward TestEpisode(const QLearningState *q_state, uint8_t *steps) {
    MountainCarState state;
    MountainCarReset(&state);
    Reward total_reward = 0;
    *steps = 0;
    while (!state.done) {
        // choose action
        MountainCarAction action = GetBestAction(q_state, &state);
        Reward reward = MountainCarStep(&state, action);
        total_reward += reward;
        (*steps)++;
    }
    return total_reward;
}

Reward TrainEpisode(QLearningState *q_state) {
    MountainCarState state;
    MountainCarReset(&state);
    Reward total_reward = 0;
    while (!state.done) {
        // choose action
        MountainCarAction action = 0;
        int state_index = get_state_index(&state);
        if (rand() < EPSILON_RAND_MAX) {
            // random action
            action = rand() % N_ACTIONS;
        } else {
            Reward max_q = q_state->q_table[state_index][0];
            for (uint8_t i = 1; i < N_ACTIONS; i++) {
                if (q_state->q_table[state_index][i] > max_q) {
                    max_q = q_state->q_table[state_index][i];
                    action = i;
                }
            }
        }
        Reward reward = MountainCarStep(&state, action);
        total_reward += reward;

        // update model
        Reward max_q_next = 0;
        if (!state.done) {
            int next_state_index = get_state_index(&state);
            max_q_next = q_state->q_table[next_state_index][0];
            for (uint8_t i = 1; i < N_ACTIONS; i++) {
                if (q_state->q_table[next_state_index][i] > max_q_next) {
                    max_q_next = q_state->q_table[next_state_index][i];
                }
            }
        }

        Reward q_value = q_state->q_table[state_index][action];
        q_state->q_table[state_index][action] = ((reward + max_q_next - (max_q_next >> GAMMA_BITSHIFT) - q_value) >> ALPHA_BITSHIFT) + q_value;
    }

    return total_reward;
}

get_state_index関数では、座標をN_STATE_0分割した区間のうちのどこに該当するかを算出する処理があります。浮動小数点数演算では、 (int)((state->position / 10000.0 + 1.2) / 1.8 * N_STATE_0) という処理でした。その代わりに、区間の境界となる値を定数テーブル position_thresholds に格納し、該当するインデックスを探索する処理で置き換えました。N_STATE_0=16なので線形探索で十分ですが、要素が多くなるなら二分探索のほうが効率的です。関数の最後でposition * N_STATE_1という掛け算がありますが、掛け算の一方が定数かつ2の累乗(N_STATE_1=16)の場合は、コンパイラが自動的にビットシフトに置換してくれます。

// (np.linspace(-1.2,0.6,16+1)*10000)[1:-1].astype(np.int32)
static const Position position_thresholds[] = {-10875,  -9750,  -8625,  -7500,  -6375,  -5250,  -4125,  -3000,
        -1875,   -749,    374,   1499,   2624,   3749,   4874};

static int get_state_index(const MountainCarState* state)
{
    uint8_t position;
    for (position = 0; position < N_STATE_0 - 1; position++) {
        if (state->position < position_thresholds[position]) {
            break;
        }
    }
    // velocityについて同様の処理が続く
}

次にεグリーディ法によるランダムな行動の実装です。ランダムに行動するか否かの判定は、[0,255]の乱数を与えるrand()関数と、EPSILON_RAND_MAX=(int)(0.168 * RAND_MAX)の比較で行います。浮動小数点数演算を記載していますが、定数ですのでPC上でコンパイル時に計算が行われますので問題ありません。ランダムな行動の生成には剰余演算で手抜きをしています。突き詰めるなら[0,255]の区間を3等分して境界値と比較するようなコードを書くことになるでしょう。

if (rand() < EPSILON_RAND_MAX) {
    // random action
    action = rand() % N_ACTIONS;
}

最後にQテーブルの更新式です。ビットシフトにより掛け算を回避しています。ビットシフトで処理できるようにするため、定数を丸めています。また、 max_q_next * 0.96875は、2を掛ける、2で割るしかできないビットシフトでは一見実装が難しいですが、max_q_next - max_q_next / 32という式に変形するとビットシフトと減算で処理可能です。

#define ALPHA_BITSHIFT 6 // 0.015625 ~= 0.0126; x * alphaの代わりに x >> ALPHA_BITSHIFTを使う
#define GAMMA_BITSHIFT 5 // 1-0.03125=0.96875 ~= 0.963; x * gammaの代わりに x - (x >> GAMMA_BITSHIFT)を使う
Reward q_value = q_state->q_table[state_index][action];
q_state->q_table[state_index][action] = ((reward + max_q_next - (max_q_next >> GAMMA_BITSHIFT) - q_value) >> ALPHA_BITSHIFT) + q_value;
// doubleを用いた場合の等価な式
// q_state->q_table[state_index][action] = (reward + max_q_next * 0.96875 - q_value) * 0.015625 + q_value;

学習アルゴリズムが実装できたので、まずはPC上で動作させてデバッグを行います。コードは以下のようになります。また、学習済みのQLearningStateをファイルに保存します。このファイルは、ゲームボーイのUIをデバッグする際に使います。

// PC用main関数
#include <stdio.h>
#include <stdlib.h>
#include "config.h"
#include "mountaincar.h"
#include "qlearning.h"

int main() {
    QLearningState *q_state = QLearningStateCreate();
    for (int epoch = 0; epoch < 1000; epoch++) {
        srand((unsigned int)epoch);
        for (int i = 0; i < 100; i++) {
            TrainEpisode(q_state);
        }
        // fix test case
        srand((unsigned int)1);
        int32_t total_reward = 0;
        for (int i = 0; i < 10; i++) {
            uint8_t steps;
            total_reward += TestEpisode(q_state, &steps);
        }
        printf("%d,Avg:%d\n", epoch, total_reward / 10);
    }

    FILE* result_file = fopen("trained.bin", "wb");
    fwrite(q_state, sizeof(QLearningState), 1, result_file);
    fclose(result_file);
    return 0;
}

ここまでの実装で、PC上で浮動小数点数演算を使わずに強化学習を実装することができました。次はゲームボーイソフトとしてのビルドを行います。