-
Notifications
You must be signed in to change notification settings - Fork 7
/
mgr.hpp
71 lines (59 loc) · 2.36 KB
/
mgr.hpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
#include <memory>
#include <madrona/py/utils.hpp>
#include <madrona/exec_mode.hpp>
#include <madrona/render/render_mgr.hpp>
namespace madEscape {
// The Manager class encapsulates the linkage between the outside training
// code and the internal simulation state (src/sim.hpp / src/sim.cpp)
//
// Manager is responsible for initializing the simulator, loading physics
// and rendering assets off disk, and mapping ECS components to tensors
// for learning
class Manager {
public:
struct Config {
madrona::ExecMode execMode; // CPU or CUDA
int gpuID; // Which GPU for CUDA backend?
uint32_t numWorlds; // Simulation batch size
uint32_t randSeed; // Seed for random world gen
bool autoReset; // Immediately generate new world on episode end
bool enableBatchRenderer;
uint32_t batchRenderViewWidth = 64;
uint32_t batchRenderViewHeight = 64;
madrona::render::APIBackend *extRenderAPI = nullptr;
madrona::render::GPUDevice *extRenderDev = nullptr;
};
Manager(const Config &cfg);
~Manager();
void step();
// These functions export Tensor objects that link the ECS
// simulation state to the python bindings / PyTorch tensors (src/bindings.cpp)
madrona::py::Tensor resetTensor() const;
madrona::py::Tensor actionTensor() const;
madrona::py::Tensor rewardTensor() const;
madrona::py::Tensor doneTensor() const;
madrona::py::Tensor selfObservationTensor() const;
madrona::py::Tensor partnerObservationsTensor() const;
madrona::py::Tensor roomEntityObservationsTensor() const;
madrona::py::Tensor doorObservationTensor() const;
madrona::py::Tensor lidarTensor() const;
madrona::py::Tensor stepsRemainingTensor() const;
madrona::py::Tensor rgbTensor() const;
madrona::py::Tensor depthTensor() const;
// These functions are used by the viewer to control the simulation
// with keyboard inputs in place of DNN policy actions
void triggerReset(int32_t world_idx);
void setAction(int32_t world_idx,
int32_t agent_idx,
int32_t move_amount,
int32_t move_angle,
int32_t rotate,
int32_t grab);
madrona::render::RenderManager & getRenderManager();
private:
struct Impl;
struct CPUImpl;
struct CUDAImpl;
std::unique_ptr<Impl> impl_;
};
}