Export render and depth tensors (#78)

* Update mgr.cpp * Fix viewer * Fix padding agent values on reset * Render padding agents too * Export valid state * Fix mean calculations (#84) * Local rendering (#79) * Add scale to road observations (#83) * Fix scale road issue * save progress * Looking good * Add json for viewer * Remove M_PI sub * Reintroduce means * Remove floor plane --------- Co-authored-by: SamanKazemkhani <[email protected]>
Emerge-Lab · May 4, 2024 · dd226e0 · dd226e0
1 parent c456e50
commit dd226e0
Show file tree

Hide file tree

Showing 8 changed files with 156 additions and 34 deletions.
diff --git a/scripts/test.py b/scripts/test.py
@@ -14,6 +14,7 @@
 params.collisionBehaviour = gpudrive.CollisionBehaviour.Ignore  # Set appropriate value
 params.datasetInitOptions = gpudrive.DatasetInitOptions.PadN  # Set appropriate value
 params.rewardParams = reward_params  # Set the rewardParams attribute to the instance created above
+params.maxNumControlledVehicles = 0
 
 # Now use the 'params' instance when creating SimManager
 sim = gpudrive.SimManager(
@@ -22,8 +23,19 @@
     num_worlds=1,
     auto_reset=True,
     json_path="nocturne_data",
-    params=params
+    params=params,
+    enable_batch_renderer=True, # Optional parameter
+    batch_render_view_width=1024,
+    batch_render_view_height=1024
 )
 
-sim.step()
-print(sim.map_observation_tensor().to_torch().shape)
+frames = []
+for steps in range(90):
+    sim.step()
+    rgb_tensor = sim.rgb_tensor().to_torch()
+    frames.append(rgb_tensor.cpu().numpy()[0,3,:,:,:3]) 
+    # Use the obs tensor for further processing
+
+import imageio
+
+imageio.mimsave('movie.gif', frames, duration=0.1)  # Save the frames as a gif
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -70,7 +70,7 @@ target_link_libraries(gpudrive PRIVATE
 if (TARGET madrona_viz)
     add_executable(viewer viewer.cpp)
     target_link_libraries(viewer PRIVATE
-        madrona_mw_core gpudrive_mgr madrona_viz)
+        madrona_mw_core gpudrive_mgr madrona_viz nlohmann_json::nlohmann_json)
 
     target_compile_definitions(viewer PRIVATE
         -DDATA_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../data/"

diff --git a/src/bindings.cpp b/src/bindings.cpp
@@ -71,20 +71,26 @@ namespace gpudrive
         // Bindings for Manager class
         nb::class_<Manager>(m, "SimManager")
             .def(
-                "__init__", [](Manager *self, madrona::py::PyExecMode exec_mode, int64_t gpu_id, int64_t num_worlds, bool auto_reset, std::string jsonPath, Parameters params)
+                "__init__", [](Manager *self, madrona::py::PyExecMode exec_mode, int64_t gpu_id, int64_t num_worlds, bool auto_reset, std::string jsonPath, Parameters params, bool enable_batch_renderer, uint32_t batch_render_view_width, uint32_t batch_render_view_height)
                 { new (self) Manager(Manager::Config{
                       .execMode = exec_mode,
                       .gpuID = (int)gpu_id,
                       .numWorlds = (uint32_t)num_worlds,
                       .autoReset = auto_reset,
                       .jsonPath = jsonPath,
-                      .params = params}); },
+                      .params = params,
+                      .enableBatchRenderer = enable_batch_renderer,
+                      .batchRenderViewWidth = batch_render_view_width,
+                      .batchRenderViewHeight = batch_render_view_height});},
                 nb::arg("exec_mode"),
                 nb::arg("gpu_id"),
                 nb::arg("num_worlds"),
                 nb::arg("auto_reset"),
                 nb::arg("json_path"),
-                nb::arg("params"))
+                nb::arg("params"),
+                nb::arg("enable_batch_renderer") = false,
+                nb::arg("batch_render_view_width") = 64,
+                nb::arg("batch_render_view_height") = 64)
             .def("step", &Manager::step)
             .def("reset", &Manager::triggerReset)
             .def("reset_tensor", &Manager::resetTensor)
@@ -103,7 +109,9 @@ namespace gpudrive
             .def("absolute_self_observation_tensor",
                  &Manager::absoluteSelfObservationTensor)
             .def("valid_state_tensor", &Manager::validStateTensor)
-            .def("info_tensor", &Manager::infoTensor);
+            .def("info_tensor", &Manager::infoTensor)
+            .def("rgb_tensor", &Manager::rgbTensor)
+            .def("depth_tensor", &Manager::depthTensor);
     }
 
 }
diff --git a/src/headless.cpp b/src/headless.cpp
@@ -69,8 +69,10 @@ int main(int argc, char *argv[])
                 .rewardType = RewardType::DistanceBased,
                 .distanceToGoalThreshold = 0.5,
                 .distanceToExpertThreshold = 0.5
-            }
-        }
+            },
+            .maxNumControlledVehicles = 0
+        },
+        .enableBatchRenderer = true
     });
 
     std::random_device rd;
@@ -92,33 +94,33 @@ int main(int argc, char *argv[])
     auto info_printer = mgr.infoTensor().makePrinter();
 
     auto printObs = [&]() {
-        printf("Self\n");
-        self_printer.print();
+        // printf("Self\n");
+        // self_printer.print();
 
-        printf("Actions\n");
-        action_printer.print();
+        // printf("Actions\n");
+        // action_printer.print();
 
         printf("Model \n");
         model_printer.print();
 
-        printf("Partner Obs\n");
-        partner_obs_printer.print();
+        // printf("Partner Obs\n");
+        // partner_obs_printer.print();
 
-        printf("Map Obs\n");
-        map_obs_printer.print();
-        printf("\n");
+        // printf("Map Obs\n");
+        // map_obs_printer.print();
+        // printf("\n");
 
-        printf("Shape\n");
-        shapePrinter.print();
+        // printf("Shape\n");
+        // shapePrinter.print();
 
-        printf("Reward\n");
-        rewardPrinter.print();
+        // printf("Reward\n");
+        // rewardPrinter.print();
 
-        printf("Done\n");
-        donePrinter.print();
+        // printf("Done\n");
+        // donePrinter.print();
 
-        printf("Controlled State\n");
-        controlledStatePrinter.print();
+        // printf("Controlled State\n");
+        // controlledStatePrinter.print();
 
         printf("Agent Map Obs\n");
         agent_map_obs_printer.print();

diff --git a/src/level_gen.cpp b/src/level_gen.cpp
@@ -272,6 +272,13 @@ static inline Entity createAgentPadding(Engine &ctx) {
     ctx.get<StepsRemaining>(agent).t = consts::episodeLen;
     ctx.get<ControlledState>(agent) = ControlledState{.controlledState = ControlMode::EXPERT};
 
+    if (ctx.data().enableRender) {
+        render::RenderingSystem::attachEntityToView(ctx,
+                agent,
+                90.f, 0.001f,
+                1.5f * math::up);
+    }
+
     return agent;
 }
 
@@ -306,7 +313,7 @@ void createPaddingEntities(Engine &ctx) {
 }
 
 void createPersistentEntities(Engine &ctx, Map *map) {
-
+    // createFloorPlane(ctx);
     ctx.data().mean = {0, 0};
     ctx.data().mean.x = map->mean.x;
     ctx.data().mean.y = map->mean.y;

diff --git a/src/mgr.cpp b/src/mgr.cpp
@@ -81,8 +81,8 @@ static inline Optional<render::RenderManager> initRenderManager(
         .agentViewWidth = mgr_cfg.batchRenderViewWidth,
         .agentViewHeight = mgr_cfg.batchRenderViewHeight,
         .numWorlds = mgr_cfg.numWorlds,
-        .maxViewsPerWorld = 2, // FIXME?
-        .maxInstancesPerWorld = 450,
+        .maxViewsPerWorld = consts::kMaxAgentCount, // FIXME?
+        .maxInstancesPerWorld = 3000,
         .execMode = mgr_cfg.execMode,
         .voxelCfg = {},
     });
@@ -260,7 +260,7 @@ static void loadRenderObjects(render::RenderManager &render_mgr)
     });
 
     render_mgr.configureLighting({
-        { true, math::Vector3{1.0f, 1.0f, -2.0f}, math::Vector3{1.0f, 1.0f, 1.0f} }
+        { true, math::Vector3{1.0f, 1.0f, -2.0f}, math::Vector3{50.0f, 50.0f, 1.0f} }
     });
 }
 
@@ -628,6 +628,14 @@ Manager::~Manager() {}
 void Manager::step()
 {
     impl_->run();
+
+    if (impl_->renderMgr.has_value()) {
+        impl_->renderMgr->readECS();
+    }
+
+    if (impl_->cfg.enableBatchRenderer) {
+        impl_->renderMgr->batchRender();
+    }
 }
 
 Tensor Manager::resetTensor() const
@@ -800,6 +808,33 @@ void Manager::triggerReset(int32_t world_idx)
     }
 }
 
+Tensor Manager::rgbTensor() const
+{
+    const uint8_t *rgb_ptr = impl_->renderMgr->batchRendererRGBOut();
+
+    assert(rgb_ptr != nullptr);
+
+    return Tensor((void*)rgb_ptr, TensorElementType::UInt8, {
+        impl_->cfg.numWorlds,
+        consts::kMaxAgentCount,
+        impl_->cfg.batchRenderViewHeight,
+        impl_->cfg.batchRenderViewWidth,
+        4,
+    }, impl_->cfg.gpuID);
+}
+
+Tensor Manager::depthTensor() const
+{
+    const float *depth_ptr = impl_->renderMgr->batchRendererDepthOut();
+
+    return Tensor((void *)depth_ptr, TensorElementType::Float32, {
+        impl_->cfg.numWorlds,
+        consts::kMaxAgentCount,
+        impl_->cfg.batchRenderViewHeight,
+        impl_->cfg.batchRenderViewWidth,
+        1,
+    }, impl_->cfg.gpuID);
+}
 
 void Manager::setAction(int32_t world_idx, int32_t agent_idx,
                         float acceleration, float steering, float headAngle) {

diff --git a/src/mgr.hpp b/src/mgr.hpp
@@ -65,6 +65,8 @@ class Manager {
     MGR_EXPORT madrona::py::Tensor absoluteSelfObservationTensor() const;
     MGR_EXPORT madrona::py::Tensor validStateTensor() const;
     MGR_EXPORT madrona::py::Tensor infoTensor() 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
     MGR_EXPORT void triggerReset(int32_t world_idx);

diff --git a/src/viewer.cpp b/src/viewer.cpp
@@ -8,9 +8,48 @@
 #include <fstream>
 #include <optional>
 
+#include <iostream>
+// #include "json_serialization.hpp"
+#include <nlohmann/json.hpp>    
+
 using namespace madrona;
 using namespace madrona::viz;
 
+std::pair<float, float> calc_mean(const nlohmann::json &j)
+{
+    std::pair<float, float> mean = {0, 0};
+    int64_t numEntities = 0;
+    for (const auto &obj : j["objects"])
+    {
+        int i = 0;
+        for (const auto &pos : obj["position"])
+        {
+            if (obj["valid"][i++] == false)
+                continue;
+            numEntities++;
+            float newX = pos["x"];
+            float newY = pos["y"];
+            // Update mean incrementally
+            mean.first += (newX - mean.first) / numEntities;
+            mean.second += (newY - mean.second) / numEntities;
+        }
+    }
+    for (const auto &obj : j["roads"])
+    {
+        for (const auto &point : obj["geometry"])
+        {
+            numEntities++;
+            float newX = point["x"];
+            float newY = point["y"];
+
+            // Update mean incrementally
+            mean.first += (newX - mean.first) / numEntities;
+            mean.second += (newY - mean.second) / numEntities;
+        }
+    }
+    return mean;
+}
+
 static HeapArray<float> readReplayLog(const char *path) {
     std::ifstream replay_log(path, std::ios::binary);
     replay_log.seekg(0, std::ios::end);
@@ -65,18 +104,19 @@ int main(int argc, char *argv[])
 #endif
 
     WindowManager wm {};
-    WindowHandle window = wm.makeWindow("Escape Room", 2730, 1536);
+    WindowHandle window = wm.makeWindow("Escape Room", 640, 480);
     render::GPUHandle render_gpu = wm.initGPU(0, { window.get() });
 
     Manager mgr({
         .execMode = exec_mode,
         .gpuID = 0,
         .numWorlds = num_worlds,
         .autoReset = replay_log.has_value(),
-        .jsonPath = "../maps",
+        .jsonPath = "/home/aarav/gpudrive/nocturne_data",
         .params = {
             .polylineReductionThreshold = 1.0,
             .observationRadius = 100.0,
+            .maxNumControlledVehicles = 0
         },
         .enableBatchRenderer = enable_batch_renderer,
         .extRenderAPI = wm.gpuAPIManager().backend(),
@@ -88,11 +128,27 @@ int main(int argc, char *argv[])
             (math::Quat::angleAxis(0, math::up) *
             math::Quat::angleAxis(-math::pi / 2.f, math::right)).normalize();
 
+    std::string path = "/home/aarav/gpudrive/nocturne_data";
+    std::string mapPath;
+    for (auto const &mapFile : std::filesystem::directory_iterator(path))
+    {
+        if (mapFile.path().extension() == ".json")
+        {
+            mapPath = mapFile.path().string();
+            break;
+        }
+    }
+    std::cout<<mapPath<<std::endl;
+    std::ifstream in(mapPath);
+    nlohmann::json rawJson;
+    in >> rawJson;
+    auto mean = calc_mean(rawJson);
+
     Viewer viewer(mgr.getRenderManager(), window.get(), {
         .numWorlds = num_worlds,
         .simTickRate = 20,
         .cameraMoveSpeed = 20.f,
-        .cameraPosition = 20.f * math::up,
+        .cameraPosition = mean.first * math::right + mean.second*math::fwd + 100.f * math::up,
         .cameraRotation = initial_camera_rotation,
     });