rollout_buffer.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <vector>
 
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
struct RolloutBatch {
    Tensor observations;
    Tensor actions;
    Tensor returns;
    Tensor log_probs;
};
 
class RolloutBuffer {
public:
    RolloutBuffer(int64_t max_length, int64_t observation_dim, int64_t action_dim, DeviceBackend* backend);
 
    void add(const Tensor& observation, const Tensor& action, float reward, float log_prob, bool done);
 
    [[nodiscard]] RolloutBatch compute_returns(float gamma) const;
 
    [[nodiscard]] Tensor rewards() const;
 
    [[nodiscard]] Tensor dones() const;
 
    void clear();
 
    [[nodiscard]] int64_t size() const { return size_; }
 
    [[nodiscard]] int64_t max_length() const { return max_length_; }
 
    [[nodiscard]] int64_t observation_dim() const { return observation_dim_; }
 
    [[nodiscard]] int64_t action_dim() const { return action_dim_; }
 
private:
    int64_t max_length_;
    int64_t observation_dim_;
    int64_t action_dim_;
    DeviceBackend* backend_;
 
    int64_t size_ = 0;
 
    // Host-side (max_length, width) row-major storage -- see the class's host-boundary note.
    std::vector<float> observations_;
    std::vector<float> actions_;
    std::vector<float> rewards_;
    std::vector<float> log_probs_;
    std::vector<float> dones_;
};
 
}  // namespace pulsatrix