trajectory_balance_loss.hpp

Go to the documentation of this file.

 
#pragma once
 
namespace pulsatrix {
 
class TrajectoryBalanceLoss {
public:
    TrajectoryBalanceLoss() = default;
 
    [[nodiscard]] float forward(float sum_log_pf, float sum_log_pb, float log_reward, float log_z);
 
    [[nodiscard]] float delta() const;
 
    [[nodiscard]] float grad_log_z() const;
 
    [[nodiscard]] float grad_weight_for_log_pf() const;
 
private:
    float last_delta_ = 0.0f;
    bool has_forwarded_ = false;
 
    void require_forwarded() const;
};
 
}  // namespace pulsatrix