diff --git a/flatland/cli.py b/flatland/cli.py index f4d4317676d2bf3de0e2f6fe522d5a6e106741f0..46bf5aecbb47f04bd00f2772897650aa4c47c12e 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 22ee05ddb49972e21dda7e8df45ee02dcebae768..60fd120984176742ecf3d82ab094ff92e6f0310c 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,