From 62badcc380e2192acc5098cafc6e8163664b3216 Mon Sep 17 00:00:00 2001
From: SP Mohanty <spmohanty91@gmail.com>
Date: Thu, 25 Jul 2019 01:57:31 +0200
Subject: [PATCH] Addresses #76 - implements a simple flatland-demo script to
test installation
---
flatland/cli.py | 48 +++++++++++++++++++++++++++++++++++++++++-------
setup.py | 2 +-
2 files changed, 42 insertions(+), 8 deletions(-)
diff --git a/flatland/cli.py b/flatland/cli.py
index f4d43176..46bf5aec 100644
--- a/flatland/cli.py
+++ b/flatland/cli.py
@@ -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
diff --git a/setup.py b/setup.py
index 22ee05dd..60fd1209 100644
--- a/setup.py
+++ b/setup.py
@@ -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,
--
GitLab