C++强化学习通过Python bindings接OpenAI Gym

OpenAI gym是强化学习最为流行的实验环境。某种程度上，其接口已经成为了标准。一方面，很多算法实现都是基于gym开发；另一方面，新的场景也会封装成gym接口。经过这样一层抽象，算法与实验环境充分解耦隔离，可以方便地自由组合。但gym是python的接口，如果想用C++实现强化学习算法，则无法直接与gym相接。一种方案是跨进程：一个进程运行python环境，另一个进程运行强化学习算法，与环境交互数据经过序列化和反序列化通过IPC进行通信。另一种是单进程方案：gym和强化学习算法跑在同一进程，通过python binding来连接。本文尝试通过pybind11来桥接，从而实现在同一进程中gym与强化学习算法通信的目的。

C++机器学习框架采用PyTorch提供的Libtorch。因为在目前主流的几个训练框架中，在C++版本上相比下它是算支持地比较好的，安装也算方便。安装流程见INSTALLING C++ DISTRIBUTIONS OF PYTORCH。官方提供的sample中提供了一个REINFORCE算法（一种虽然比较古老但很经典的RL算法）的例子reinforce.py，我们就先以它为例。

首先来看python部分，参考原sample做一些小改动。将强化学习算法调用接口抽象在RLWrapper。这个类后面会binding到C++。初始化时传入gym环境的状态和动态空间描述，reset()函数通知环境重置并传入初始状态，act()函数根据当前状态根据策略给出动作，update()函数进行策略函数参数学习。

...
def state_space_desc(space):
    if isinstance(space, gym.spaces.Box):
        assert(type(space.shape)==tuple)
        return dict(stype='Box', dtype=str(space.dtype), shape=space.shape)
    else:
        raise NotImplementedError('unknown state space {}'.format(space))

def action_space_desc(space):
    if isinstance(space, gym.spaces.Discrete):
        return dict(stype='Discrete', dtype=str(space.dtype), shape=(space.n,))
    else:
        raise NotImplementedError('unknown action space {}'.format(space))

def main(args):
    env = gym.make(args.env)
    env.seed(args.seed)

    agent = nativerl.RLWrapper(state_space_desc(env.observation_space),
            action_space_desc(env.action_space))

    running_reward = 10
    for i in range(args.epoch):
        obs = env.reset()
        ep_reward = 0
        agent.reset(obs)
        for t in range(1, args.step):
            if args.render:
                env.render()
            action = agent.act(obs)
            obs, reward, done, info = env.step(action)
            agent.update(reward, done)
            ep_reward += reward
            if done:
                break

        running_reward = 0.05 * ep_reward + (1 - 0.05) * running_reward
        agent.episode_finish()

        if i % args.log_itv == 0:
            print("Episode {}\t Last reward: {:.2f}\t step: {}\t Average reward: {:.2f}"
                    .format(i, ep_reward, t, running_reward))
        if env.spec.reward_threshold and running_reward > env.spec.reward_threshold:
            print("Solved. Running reward: {}, Last reward: {}".format(running_reward, t))
            break

    env.close()

然后是RLWrapper的python binding部分。这里主要是将python的对象转为C++的数据结构。

...
namespace py = pybind11;

class RLWrapper {
public:
    RLWrapper(const py::dict& state_space, const py::dict& action_space) {
        spdlog::set_level(spdlog::level::info);
        torch::manual_seed(nrl::kSeed);

        nrl::SpaceDesc ss;
        nrl::SpaceDesc as;
        ss.stype = py::cast<std::string>(state_space["stype"]);
        as.stype = py::cast<std::string>(action_space["stype"]);
        ss.dtype = py::cast<std::string>(state_space["dtype"]);
        as.dtype = py::cast<std::string>(action_space["dtype"]);

        py::tuple shape;
        shape = py::cast<py::tuple>(state_space["shape"]);
        for (const auto& item : shape) {
            ss.shape.push_back(py::cast<int64_t>(item));
        }
        shape = py::cast<py::tuple>(action_space["shape"]);
        for (const auto& item : shape) {
            as.shape.push_back(py::cast<int64_t>(item));
        }

        mStateSpaceDesc = ss;
        mActionSpaceDesc = as;
        mAgent = std::make_shared<nrl::Reinforce>(ss, as);
    }

    void reset(py::array_t<float, py::array::c_style | py::array::forcecast> state) {
        py::buffer_info buf = state.request();
        float* pbuf = static_cast<float*>(buf.ptr);
        
        assert(buf.shape == mStateSpaceDesc.shape);
        mAgent->reset(nrl::Blob{pbuf, mStateSpaceDesc.shape});
    }

    py::object act(py::array_t<float, py::array::c_style | py::array::forcecast> state) {
        py::buffer_info buf = state.request();
        float* pbuf = static_cast<float *>(buf.ptr);
        assert(buf.shape == mStateSpaceDesc.shape);
        torch::Tensor action = mAgent->act(nrl::Blob{pbuf, mStateSpaceDesc.shape}).contiguous().cpu();
        return py::int_(action.item<long>());
    }

    void update(float reward, bool done) {
        mAgent->update(reward, done);
    }

    void episode_finish() {
        spdlog::trace("{}", __func__);
        mAgent->onEpisodeFinished();
    }

    ~RLWrapper() {
    }

private:
    nrl::SpaceDesc mStateSpaceDesc;
    nrl::SpaceDesc mActionSpaceDesc;
    std::shared_ptr<nrl::RLBase> mAgent;
};

PYBIND11_MODULE(nativerl, m) {
    py::class_<RLWrapper>(m, "RLWrapper")
        .def(py::init<const py::dict &, const py::dict &>())
        .def("reset", &RLWrapper::reset)
        .def("episode_finish", &RLWrapper::episode_finish)
        .def("act", &RLWrapper::act)
        .def("update", &RLWrapper::update);
}

可以说这是python和C++的glue层。主要的工作我们放到RLBase类中。它是一个抽象类，定义了几个强化学习的基本接口。我们将REINFORCE算法实现在其继承类Reinforce中：

...
class Reinforce : public RLBase {
public:
    Reinforce(const SpaceDesc& ss, const SpaceDesc& as)
        : mPolicy(std::make_shared<Policy>(ss, as, mDevice)) {
        mPolicy->to(mDevice);

        mRewards = torch::zeros({mCapacity}, torch::TensorOptions(mDevice));
        mReturns = torch::zeros({mCapacity}, torch::TensorOptions(mDevice));

        mOptimizer = std::make_shared<torch::optim::Adam>(mPolicy->parameters(), 
                torch::optim::AdamOptions(mInitLR));
        }

    virtual torch::Tensor act(const Blob& s) override {
        auto state = torch::from_blob(s.pbuf, s.shape).unsqueeze(0).to(mDevice);
        torch::Tensor action;
        torch::Tensor logProb;
        std::tie(action, logProb) = mPolicy->act(state);
        mLogProbs.push_back(logProb);
        return action;
    }

    void update(float r, __attribute__((unused)) bool done) {
        mRewards[mSize++] = r;
        if (mSize >= mCapacity) {
            spdlog::info("buffer has been full, call train()");
            train();
        }
    }

    virtual void onEpisodeFinished() override {
        train();
    }

private:
    void train() {
        spdlog::trace("{}: buffer size = {}", __func__, mSize);

        for (auto i = mSize - 1; i >= 0; --i) {
            if (i == (mSize - 1)) {
                mReturns[i] = mRewards[i];
            } else {
                mReturns[i] = mReturns[i + 1] * mGamma + mRewards[i];
            }
        } 
        auto returns = mReturns.slice(0, 0, mSize);
        returns = (returns - returns.mean()) / (returns.std() + kEps);

        auto logprobs = torch::cat(mLogProbs);

        mOptimizer->zero_grad();
        auto policy_loss = -(logprobs * returns).sum();
        policy_loss.backward();
        mOptimizer->step();

        mLogProbs.clear();
        mSize = 0;
        ++mCount;
        spdlog::debug("{} : episode {}: loss = {}, accumulated reward = {}",
                __func__, mCount, policy_loss.item<float>(), mRewards.sum().item<float>());
    }

    std::shared_ptr<Policy> mPolicy;

    torch::Tensor mRewards;
    std::vector<torch::Tensor> mLogProbs;
    torch::Tensor mReturns;
    int32_t mSize{0};
    int32_t mCapacity{kExpBufferCap};

    std::shared_ptr<torch::optim::Adam> mOptimizer;
    uint32_t mCount{0};

    float mGamma{0.99};
    float mInitLR{1e-2};
};

Sample中的场景为CartPole，场景比较简单，因此其中的策略函数实现为简单的MLP。更为复杂的场景可以替换为复杂的网络结构。

...
class Net : public nn::Module {
public:
    virtual std::tuple<Tensor, Tensor> forward(Tensor x) = 0;

    virtual ~Net() = default;
};

class MLP : public Net {
public:
    MLP(int64_t inputSize, int64_t actionNum) {
        mFC1 = register_module("fc1", nn::Linear(inputSize, mHiddenSize));
        mAction = register_module("action", nn::Linear(mHiddenSize, actionNum));
        mValue = register_module("value", nn::Linear(mHiddenSize, actionNum));
    }

    virtual std::tuple<Tensor, Tensor> forward(Tensor x) override {
        x = mFC1->forward(x);
        x = dropout(x, 0.6, is_training());
        x = relu(x);
        return std::make_tuple(mAction->forward(x), mValue->forward(x));
    }

private:
    nn::Linear mFC1{nullptr};
    nn::Linear mAction{nullptr};
    nn::Linear mValue{nullptr};
    int64_t mHiddenSize{128};
};

class Policy : public torch::nn::Module {
public:
    Policy(const SpaceDesc& ss, const SpaceDesc& as, torch::Device mDevice) 
        : mActionSpaceType(as.stype),
        mActionNum(as.shape[0]),
        mGen(kSeed),
        mUniformDist(0, 1.0) {

        if (ss.shape.size() == 1) {
            mNet = std::make_shared<MLP>(ss.shape[0], as.shape[0]);
        } else {
            mNet = std::make_shared<CNN>(ss.shape, as.shape[0]);
        }
        mNet->to(mDevice);
        register_module("base", mNet);

        torch::Tensor logits = torch::ones({1, as.shape[0]}, torch::TensorOptions(mDevice));
        mUniformCategorical = std::make_shared<Categorical>(nullptr, &logits);

    }

    torch::Tensor forward(torch::Tensor x) {
        x = std::get<0>(mNet->forward(x));
        return torch::softmax(x, 1);
    }

    std::tuple<torch::Tensor, torch::Tensor> act(torch::Tensor state) {
        auto output = forward(state);
        std::shared_ptr<Distribution> dist;
        if (!mActionSpaceType.compare("Discrete")) {
            dist = std::make_shared<Categorical>(&output);
        } else {
            throw std::logic_error("Not implemented : action space");
        }

        float rnd = mUniformDist(mGen);
        float threshold = kEpsEnd + (kEpsStart - kEpsEnd) * exp(-1. * mStep / kEpsDecay);
        ++mStep;
        torch::Tensor action;
        if (rnd > threshold) {
            torch::NoGradGuard no_grad;
            action = dist->sample();
        } else {
            torch::NoGradGuard no_grad;
            action = mUniformCategorical->sample({1}).squeeze(-1);
        }

        auto log_probs = dist->log_prob(action);
        return std::make_tuple(action, log_probs);
    }

private:
    std::string mActionSpaceType;
    int32_t mActionNum;
    int64_t mHiddenSize{128};
    std::shared_ptr<Net> mNet;
    uint64_t mStep{0};

    std::mt19937 mGen;
    std::uniform_real_distribution<float> mUniformDist;
    std::shared_ptr<Categorical> mUniformCategorical;
};

其中的Categorical类为Categorical distribution相关计算，可以根据PyTorch中的python版本重写成C++。

最后，将上面的C++实现编译成so。根据实际情况在CMakeLists.txt中加入：

...
set(CMAKE_CXX_STANDARD 11)

find_package(Torch REQUIRED)

set(NRL_INCLUDE_DIRS
    src
    ${TORCH_INCLUDE_DIRS})

file(GLOB NRL_SOURCES1 "src/*.cpp")
list(APPEND NRL_SOURCES ${NRL_SOURCES1})
message(STATUS "sources: ${NRL_SOURCES}")

add_subdirectory(third_party/pybind11)
add_subdirectory(third_party/spdlog)

pybind11_add_module(nativerl ${NRL_SOURCES})
target_include_directories(nativerl PRIVATE ${NRL_INCLUDE_DIRS})
target_link_libraries(nativerl PRIVATE spdlog::spdlog ${TORCH_LIBRARIES})
...

假设编译出的so位于build目录下，python脚本为example/simple.py。则可以通过命令开始训练：

PYTHONPATH=./build python -m example.simple

正常的话可以看到类似的训练过程log及结果，基本和python版本一致。

[2019-06-22 13:42:22.533] [info] state space type:Box shape size:1
[2019-06-22 13:42:22.534] [info] action space type:Discrete, shape size:1
[2019-06-22 13:42:22.534] [info] Training on GPU (CUDA)
Episode 0	 Last reward: 29.00	 step: 29	 Average reward: 10.95
Episode 10	 Last reward: 17.00	 step: 17	 Average reward: 14.73
Episode 20	 Last reward: 12.00	 step: 12	 Average reward: 17.40
Episode 30	 Last reward: 15.00	 step: 15	 Average reward: 24.47
Episode 40	 Last reward: 18.00	 step: 18	 Average reward: 26.22
Episode 50	 Last reward: 18.00	 step: 18	 Average reward: 23.69
Episode 60	 Last reward: 72.00	 step: 72	 Average reward: 30.21
Episode 70	 Last reward: 19.00	 step: 19	 Average reward: 28.83
Episode 80	 Last reward: 29.00	 step: 29	 Average reward: 32.13
Episode 90	 Last reward: 15.00	 step: 15	 Average reward: 29.64
Episode 100	 Last reward: 30.00	 step: 30	 Average reward: 27.88
Episode 110	 Last reward: 12.00	 step: 12	 Average reward: 26.14
Episode 120	 Last reward: 28.00	 step: 28	 Average reward: 26.32
Episode 130	 Last reward: 11.00	 step: 11	 Average reward: 31.20
Episode 140	 Last reward: 112.00	 step: 112	 Average reward: 35.26
Episode 150	 Last reward: 40.00	 step: 40	 Average reward: 37.14
Episode 160	 Last reward: 40.00	 step: 40	 Average reward: 36.84
Episode 170	 Last reward: 15.00	 step: 15	 Average reward: 41.91
Episode 180	 Last reward: 63.00	 step: 63	 Average reward: 49.78
Episode 190	 Last reward: 21.00	 step: 21	 Average reward: 44.70
Episode 200	 Last reward: 46.00	 step: 46	 Average reward: 41.83
Episode 210	 Last reward: 80.00	 step: 80	 Average reward: 51.55
Episode 220	 Last reward: 151.00	 step: 151	 Average reward: 57.82
Episode 230	 Last reward: 176.00	 step: 176	 Average reward: 62.80
Episode 240	 Last reward: 19.00	 step: 19	 Average reward: 63.17
Episode 250	 Last reward: 134.00	 step: 134	 Average reward: 74.02
Episode 260	 Last reward: 46.00	 step: 46	 Average reward: 71.35
Episode 270	 Last reward: 118.00	 step: 118	 Average reward: 85.88
Episode 280	 Last reward: 487.00	 step: 487	 Average reward: 112.74
Episode 290	 Last reward: 95.00	 step: 95	 Average reward: 139.41
Episode 300	 Last reward: 54.00	 step: 54	 Average reward: 149.20
Episode 310	 Last reward: 417.00	 step: 417	 Average reward: 138.42
Episode 320	 Last reward: 500.00	 step: 500	 Average reward: 179.29
Episode 330	 Last reward: 71.00	 step: 71	 Average reward: 195.88
Episode 340	 Last reward: 309.00	 step: 309	 Average reward: 216.82
Episode 350	 Last reward: 268.00	 step: 268	 Average reward: 214.21
Episode 360	 Last reward: 243.00	 step: 243	 Average reward: 210.89
Episode 370	 Last reward: 266.00	 step: 266	 Average reward: 200.03
Episode 380	 Last reward: 379.00	 step: 379	 Average reward: 220.06
Episode 390	 Last reward: 500.00	 step: 500	 Average reward: 316.20
Episode 400	 Last reward: 500.00	 step: 500	 Average reward: 369.46
Episode 410	 Last reward: 500.00	 step: 500	 Average reward: 421.84
Episode 420	 Last reward: 500.00	 step: 500	 Average reward: 453.20
Episode 430	 Last reward: 500.00	 step: 500	 Average reward: 471.98
Solved. Running reward: 475.9764491024681, Last reward: 500