ppo_clipped_loss.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 {
 
class PPOClippedLoss {
public:
    explicit PPOClippedLoss(DeviceBackend* backend);
 
    [[nodiscard]] float forward(const Tensor& new_logits, const Tensor& actions, const Tensor& old_log_probs,
                                const Tensor& advantages, float clip_epsilon);
 
    [[nodiscard]] Tensor backward() const;
 
private:
    DeviceBackend* backend_;
    Tensor last_probs_;
    Tensor last_advantages_;
    Tensor last_ratios_;
    Tensor last_masks_;
    Tensor last_action_indices_;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix