categorical_policy_agent.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
 
#include "pulsatrix/agent.hpp"
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/module.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class CategoricalPolicyAgent : public Agent {
public:
    CategoricalPolicyAgent(Module* policy_network, int64_t action_dim, DeviceBackend* backend, uint32_t seed = 42);
 
    [[nodiscard]] Tensor act(const Tensor& observation) override;
 
    [[nodiscard]] Tensor act_greedy(const Tensor& observation);
 
    [[nodiscard]] float log_prob() const;
 
    [[nodiscard]] int64_t action_dim() const { return action_dim_; }
 
private:
    [[nodiscard]] float next_unit();
 
    [[nodiscard]] Tensor policy_logits(const Tensor& observation);
 
    Module* policy_network_;
    int64_t action_dim_;
    DeviceBackend* backend_;
    uint32_t lcg_state_;
    float last_log_prob_ = 0.0f;
    bool has_acted_ = false;
};
 
}  // namespace pulsatrix