Commit 62badcc3 authored by spmohanty's avatar spmohanty
Browse files

Addresses #76 - implements a simple flatland-demo script to test installation

parent 8b2305c2
Pipeline #1588 passed with stages
in 21 minutes and 48 seconds
......@@ -2,18 +2,52 @@
"""Console script for flatland."""
import sys
import click
import numpy as np
import time
from flatland.envs.generators import complex_rail_generator
from flatland.envs.rail_env import RailEnv
from flatland.utils.rendertools import RenderTool
@click.command()
def main(args=None):
"""Console script for flatland."""
click.echo("Replace this message by putting your code into "
"flatland.cli.main")
click.echo("See click documentation at http://click.pocoo.org/")
def demo(args=None):
"""Demo script to check installation"""
env = RailEnv(
width=15,
height=15,
rail_generator=complex_rail_generator(
nr_start_goal=10,
nr_extra=1,
min_dist=8,
max_dist=99999),
number_of_agents=5)
env._max_episode_steps = int(15 * (env.width + env.height))
env_renderer = RenderTool(env)
while True:
obs = env.reset()
_done = False
# Run a single episode here
step = 0
while not _done:
# Compute Action
_action = {}
for _idx, _ in enumerate(env.agents):
_action[_idx] = np.random.randint(0, 5)
obs, all_rewards, done, _ = env.step(_action)
_done = done['__all__']
step += 1
env_renderer.render_env(
show=True,
frames=False,
show_observations=False,
show_predictions=False
)
time.sleep(0.3)
return 0
if __name__ == "__main__":
sys.exit(main()) # pragma: no cover
sys.exit(demo()) # pragma: no cover
......@@ -62,7 +62,7 @@ setup(
description="Multi Agent Reinforcement Learning on Trains",
entry_points={
'console_scripts': [
'flatland=flatland.cli:main',
'flatland-demo=flatland.cli:demo',
],
},
install_requires=requirements,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment