-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
0889fa0
commit bb01f48
Showing
5 changed files
with
164 additions
and
68 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
import threading | ||
import time | ||
|
||
from jax import random | ||
|
||
class SimulationWrapper: | ||
|
||
def __init__(self, simulation, state, key, step_delay=0.1, print_data=False): | ||
self.running = False | ||
self.paused = False | ||
self.stop_requested = False | ||
self.update_thread = None | ||
self.print_data = print_data | ||
self.step_delay = step_delay | ||
# simulation_dependent | ||
self.simulation = simulation | ||
self.state = state | ||
self.key = key | ||
|
||
def start(self): | ||
if not self.running: | ||
self.running = True | ||
self.stop_requested = False | ||
self.update_thread = threading.Thread(target=self.simulation_loop) | ||
self.update_thread.start() | ||
|
||
def pause(self): | ||
self.paused = True | ||
|
||
def resume(self): | ||
self.paused = False | ||
|
||
def stop(self): | ||
self.stop_requested = True | ||
|
||
def simulation_loop(self): | ||
while not self.stop_requested: | ||
if self.paused: | ||
time.sleep(0.1) | ||
continue | ||
|
||
self.state = self._update_simulation() | ||
if self.print_data: | ||
print(f"{self.state = }") | ||
|
||
time.sleep(self.step_delay) | ||
|
||
def _update_simulation(self): | ||
self.key, a_key, step_key = random.split(self.key, 3) | ||
actions = self.simulation.choose_action(self.state.obs, a_key) | ||
return self.simulation.step(self.state, actions, step_key) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
import time | ||
|
||
import jax | ||
import jax.numpy as jnp | ||
from jax import random | ||
|
||
from simulationsandbox.two_d_simulation import SimpleSimulation | ||
from simulationsandbox.three_d_simulation import ThreeDSimulation | ||
|
||
NUM_AGENTS = 5 | ||
MAX_AGENTS = 10 | ||
NUM_OBS = 3 | ||
GRID_SIZE = 20 | ||
NUM_STEPS = 50 | ||
VIZUALIZE = True | ||
STEP_DELAY = 0.000001 | ||
SEED = 0 | ||
|
||
|
||
def test_simulation_init(): | ||
key = random.PRNGKey(SEED) | ||
|
||
sim = SimpleSimulation(MAX_AGENTS, GRID_SIZE) | ||
state = sim.init_state(NUM_AGENTS, NUM_OBS, key) | ||
|
||
assert sim.max_agents == MAX_AGENTS | ||
assert sim.grid_size == GRID_SIZE | ||
assert state.x_pos.shape == (MAX_AGENTS,) | ||
assert jnp.sum(state.alive) == NUM_AGENTS | ||
assert state.grid.shape == (GRID_SIZE, GRID_SIZE) | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,74 @@ | ||
import time | ||
|
||
import jax | ||
import jax.numpy as jnp | ||
from jax import random | ||
|
||
from simulationsandbox.two_d_simulation import SimpleSimulation | ||
from simulationsandbox.three_d_simulation import ThreeDSimulation | ||
from simulationsandbox.wrapper import SimulationWrapper | ||
|
||
NUM_AGENTS = 5 | ||
MAX_AGENTS = 10 | ||
NUM_OBS = 3 | ||
GRID_SIZE = 20 | ||
NUM_STEPS = 50 | ||
VIZUALIZE = True | ||
SLEEP_TIME = 2 | ||
STEP_DELAY = 0.01 | ||
PRINT_DATA = True | ||
SEED = 0 | ||
|
||
|
||
def test_wrapper_two_d_sim(): | ||
key = random.PRNGKey(SEED) | ||
env = SimpleSimulation(MAX_AGENTS, GRID_SIZE) | ||
state = env.init_state(NUM_AGENTS, NUM_OBS, key) | ||
|
||
# Example usage: | ||
sim = SimulationWrapper(env, state, key, step_delay=STEP_DELAY, print_data=PRINT_DATA) | ||
|
||
print('Started') | ||
sim.start() | ||
time.sleep(SLEEP_TIME) | ||
assert sim.running == True | ||
|
||
sim.pause() | ||
print('Paused') | ||
time.sleep(SLEEP_TIME) | ||
assert sim.paused == True | ||
|
||
print('Resumed') | ||
sim.resume() | ||
time.sleep(SLEEP_TIME) | ||
assert sim.running == True | ||
|
||
sim.stop() | ||
print('stopped') | ||
|
||
|
||
# def test_wrapper_three_d_sim(): | ||
# key = random.PRNGKey(SEED) | ||
# env = ThreeDSimulation(MAX_AGENTS, GRID_SIZE) | ||
# state = env.init_state(NUM_AGENTS, NUM_OBS, key) | ||
|
||
# # Example usage: | ||
# sim = SimulationWrapper(env, state, key, print_data=True) | ||
|
||
# print('Started') | ||
# sim.start() | ||
# time.sleep(SLEEP_TIME) | ||
# assert sim.running == True | ||
|
||
# sim.pause() | ||
# print('Paused') | ||
# time.sleep(SLEEP_TIME) | ||
# assert sim.paused == True | ||
|
||
# print('Resumed') | ||
# sim.resume() | ||
# time.sleep(SLEEP_TIME) | ||
# assert sim.running == True | ||
|
||
# sim.stop() | ||
# print('stopped') |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters