replay_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 ReplayBatch {
    Tensor observations;
    Tensor actions;
    Tensor rewards;
    Tensor next_observations;
    Tensor dones;
};
 
class ReplayBuffer {
public:
    ReplayBuffer(int64_t capacity, int64_t observation_dim, int64_t action_dim, DeviceBackend* backend,
                 uint32_t seed = 42);
 
    void add(const Tensor& observation, const Tensor& action, float reward, const Tensor& next_observation,
             bool done);
 
    [[nodiscard]] ReplayBatch sample(int64_t batch_size);
 
    [[nodiscard]] int64_t size() const { return size_; }
 
    [[nodiscard]] int64_t capacity() const { return capacity_; }
 
    [[nodiscard]] int64_t observation_dim() const { return observation_dim_; }
 
    [[nodiscard]] int64_t action_dim() const { return action_dim_; }
 
private:
    [[nodiscard]] int64_t next_index(int64_t bound);
 
    int64_t capacity_;
    int64_t observation_dim_;
    int64_t action_dim_;
    DeviceBackend* backend_;
    uint32_t lcg_state_;
 
    int64_t size_ = 0;
    int64_t write_index_ = 0;
 
    // Host-side (capacity, 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> next_observations_;
    std::vector<float> dones_;
};
 
}  // namespace pulsatrix