hypergrid_env.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <vector>
 
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/environment.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class HyperGridEnv : public Environment {
public:
    explicit HyperGridEnv(DeviceBackend* backend, int64_t ndim = 2, int64_t side_length = 8, float r0 = 0.1f,
                           float r1 = 0.5f, float r2 = 2.0f, int64_t max_steps = 64);
 
    [[nodiscard]] Tensor reset() override;
 
    [[nodiscard]] Tensor reset(const Tensor& initial_state) override;
 
    [[nodiscard]] StepResult step(const Tensor& action) override;
 
    [[nodiscard]] float reward(const Tensor& state) const;
 
    [[nodiscard]] float backward_log_prob(const Tensor& state, int64_t action) const;
 
    [[nodiscard]] std::vector<bool> valid_actions_mask(const Tensor& state) const;
 
    [[nodiscard]] int64_t observation_dim() const override { return ndim_; }
 
    [[nodiscard]] int64_t action_dim() const override { return ndim_ + 1; }
 
    [[nodiscard]] bool is_discrete() const override { return true; }
 
    [[nodiscard]] int64_t ndim() const { return ndim_; }
 
    [[nodiscard]] int64_t side_length() const { return side_length_; }
 
    [[nodiscard]] int64_t max_steps() const { return max_steps_; }
 
    [[nodiscard]] int64_t step_count() const { return step_count_; }
 
private:
    [[nodiscard]] Tensor observation() const;
 
    void decode_state(const Tensor& state, std::vector<int64_t>& out) const;
 
    DeviceBackend* backend_;
    int64_t ndim_;
    int64_t side_length_;
    float r0_;
    float r1_;
    float r2_;
    int64_t max_steps_;
 
    std::vector<int64_t> coord_;
    int64_t step_count_ = 0;
    bool has_reset_ = false;
    bool done_ = false;
};
 
}  // namespace pulsatrix