tanh_gaussian_policy.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
struct TanhGaussianSample {
    Tensor action;
    Tensor log_prob;
};
 
struct TanhGaussianGrad {
    Tensor grad_mean;
    Tensor grad_log_std;
};
 
class TanhGaussianPolicy {
public:
    static constexpr float kLogProbStabilizer = 1e-6f;
 
    explicit TanhGaussianPolicy(DeviceBackend* backend);
 
    [[nodiscard]] TanhGaussianSample forward(const Tensor& mean, const Tensor& log_std, const Tensor& epsilon);
 
    [[nodiscard]] TanhGaussianGrad backward(const Tensor& grad_action, const Tensor& grad_log_prob) const;
 
private:
    DeviceBackend* backend_;
    // The cached action (not u) -- every backward term is expressible in a, std and epsilon,
    // and caching a avoids recomputing tanh in backward().
    Tensor last_action_;
    Tensor last_std_;
    Tensor last_epsilon_;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix