gflownet_forward_policy.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <vector>
 
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/module.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
struct GFlowNetSampledAction {
    Tensor action;
    float log_prob;
};
 
class GFlowNetForwardPolicy {
public:
    GFlowNetForwardPolicy(Module* policy_network, int64_t action_dim, DeviceBackend* backend, uint32_t seed = 42);
 
    [[nodiscard]] GFlowNetSampledAction sample(const Tensor& observation, const std::vector<bool>& valid_actions);
 
    [[nodiscard]] std::vector<float> masked_probs(const Tensor& observation, const std::vector<bool>& valid_actions);
 
    [[nodiscard]] int64_t action_dim() const { return action_dim_; }
 
private:
    [[nodiscard]] float next_unit();
 
    Module* policy_network_;
    int64_t action_dim_;
    DeviceBackend* backend_;
    uint32_t lcg_state_;
};
 
}  // namespace pulsatrix