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

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

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,