Emerge-Lab / gpudrive

GPU-acceleration of Nocturne via Madrona
https://arxiv.org/abs/2408.01584
MIT License
233 stars 20 forks source link

Dynamic files world reset #265

Closed aaravpandya closed 1 month ago

aaravpandya commented 1 month ago

This PR adds support to dynamically reset the sim on runtime using different map files. We introduce new method set_maps(List: str) that takes in exactly the num_worlds file names and resets the sim to on these new maps.

Here is an example script to use the new functionality -

import gpudrive

scenes = ["data/examples/tfrecord-00012-of-01000_389.json",]

params = gpudrive.Parameters()
params.polylineReductionThreshold = 0.5

sim = gpudrive.SimManager(
    exec_mode=gpudrive.madrona.ExecMode.CPU,
    gpu_id=0,
    scenes=scenes,
    params=params,  # Environment parameters
)

shape_tensor = sim.shape_tensor().to_torch()

print(shape_tensor)
sim.set_maps(["data/examples/tfrecord-00001-of-01000_307.json"])
print(shape_tensor)

Expected output

(madrona) (base) aarav@emerge2-desktop:~/gpudrive$ python index.py
tensor([[ 18, 456]], dtype=torch.int32)
tensor([[ 10, 182]], dtype=torch.int32)
aaravpandya commented 1 month ago

LGTM

This PR still needs some work. I just fixed an issue where it was not compiling to CUDA, and Daphne is fixing some issues related to shapes during training.

eugenevinitsky commented 1 month ago

Ah can you mark it as WIP or a draft if it's not ready for review?