custom_railmap_example.py 1014 Bytes
Newer Older
1
2
import random

3
4
import numpy as np

5
from flatland.core.transition_map import GridTransitionMap
u214892's avatar
u214892 committed
6
from flatland.core.grid.rail_env_grid import RailEnvTransitions
7
from flatland.envs.rail_env import RailEnv
8
9
10
11
12
from flatland.utils.rendertools import RenderTool

random.seed(100)
np.random.seed(100)

spiglerg's avatar
spiglerg committed
13

14
15
16
17
18
19
20
21
22
23
24
25
def custom_rail_generator():
    def generator(width, height, num_agents=0, num_resets=0):
        rail_trans = RailEnvTransitions()
        grid_map = GridTransitionMap(width=width, height=height, transitions=rail_trans)
        rail_array = grid_map.grid
        rail_array.fill(0)

        agents_positions = []
        agents_direction = []
        agents_target = []

        return grid_map, agents_positions, agents_direction, agents_target
26

27
28
    return generator

spiglerg's avatar
spiglerg committed
29

30
31
32
33
34
35
36
env = RailEnv(width=6,
              height=4,
              rail_generator=custom_rail_generator(),
              number_of_agents=1)

env.reset()

37
env_renderer = RenderTool(env)
38
39
40
env_renderer.renderEnv(show=True)

input("Press Enter to continue...")