# Copyright (c) 2022-2024.
# ProrokLab (https://www.proroklab.org/)
# All rights reserved.
"""
Use this script to interactively play with scenarios
You can change agent by pressing TAB
You can reset the environment by pressing R
You can move agents with the arrow keys
If you have more than 1 agent, you can control another one with W,A,S,D
and switch the agent with these controls using LSHIFT
"""
from argparse import ArgumentParser, BooleanOptionalAction
from operator import add
from typing import Dict, Union
import numpy as np
from torch import Tensor
from vmas.make_env import make_env
from vmas.simulator.environment.gym import GymWrapper
from vmas.simulator.scenario import BaseScenario
from vmas.simulator.utils import save_video
N_TEXT_LINES_INTERACTIVE = 6
class InteractiveEnv:
"""
Use this script to interactively play with scenarios
You can change agent by pressing TAB
You can reset the environment by pressing R
You can control agent actions with the arrow keys and M/N (left/right control the first action, up/down control the second, M/N controls the third)
If you have more than 1 agent, you can control another one with W,A,S,D and Q,E in the same way.
and switch the agent with these controls using LSHIFT
"""
def __init__(
self,
env: GymWrapper,
control_two_agents: bool = False,
display_info: bool = True,
save_render: bool = False,
render_name: str = "interactive",
):
self.env = env
self.control_two_agents = control_two_agents
# hard-coded keyboard events
self.current_agent_index = 0
self.current_agent_index2 = 1
self.n_agents = self.env.unwrapped.n_agents
self.agents = self.env.unwrapped.agents
self.continuous = self.env.unwrapped.continuous_actions
self.reset = False
self.keys = np.array(
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
) # up, down, left, right, rot+, rot-
self.keys2 = np.array(
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
) # up, down, left, right, rot+, rot-
self.u = [0] * (3 if self.continuous else 2)
self.u2 = [0] * (3 if self.continuous else 2)
self.frame_list = []
self.display_info = display_info
self.save_render = save_render
self.render_name = render_name
if self.control_two_agents:
assert (
self.n_agents >= 2
), "Control_two_agents is true but not enough agents in scenario"
self.text_lines = []
self.font_size = 15
self.env.render()
self.text_idx = len(self.env.unwrapped.text_lines)
self._init_text()
self.env.unwrapped.viewer.window.on_key_press = self._key_press
self.env.unwrapped.viewer.window.on_key_release = self._key_release
self._cycle()
def _increment_selected_agent_index(self, index: int):
index += 1
if index == self.n_agents:
index = 0
return index
def _cycle(self):
total_rew = [0] * self.n_agents
while True:
if self.reset:
if self.save_render:
save_video(
self.render_name,
self.frame_list,
fps=1 / self.env.unwrapped.world.dt,
)
self.env.reset()
self.reset = False
total_rew = [0] * self.n_agents
if self.n_agents > 0:
action_list = [[0.0] * agent.action_size for agent in self.agents]
action_list[self.current_agent_index][
: self.agents[self.current_agent_index].dynamics.needed_action_size
] = self.u[
: self.agents[self.current_agent_index].dynamics.needed_action_size
]
else:
action_list = []
if self.n_agents > 1 and self.control_two_agents:
action_list[self.current_agent_index2][
: self.agents[self.current_agent_index2].dynamics.needed_action_size
] = self.u2[
: self.agents[self.current_agent_index2].dynamics.needed_action_size
]
obs, rew, done, info = self.env.step(action_list)
if self.display_info and self.n_agents > 0:
# TODO: Determine number of lines of obs_str and render accordingly
obs_str = str(InteractiveEnv.format_obs(obs[self.current_agent_index]))
message = f"\t\t{obs_str[len(obs_str) // 2:]}"
self._write_values(0, message)
message = f"Obs: {obs_str[:len(obs_str) // 2]}"
self._write_values(1, message)
message = f"Rew: {round(rew[self.current_agent_index],3)}"
self._write_values(2, message)
total_rew = list(map(add, total_rew, rew))
message = f"Total rew: {round(total_rew[self.current_agent_index], 3)}"
self._write_values(3, message)
message = f"Done: {done}"
self._write_values(4, message)
message = f"Selected: {self.env.unwrapped.agents[self.current_agent_index].name}"
self._write_values(5, message)
frame = self.env.render(
mode="rgb_array" if self.save_render else "human",
visualize_when_rgb=True,
)
if self.save_render:
self.frame_list.append(frame)
if done:
self.reset = True
def _init_text(self):
from vmas.simulator import rendering
for i in range(N_TEXT_LINES_INTERACTIVE):
text_line = rendering.TextLine(
y=(self.text_idx + i) * 40, font_size=self.font_size
)
self.env.unwrapped.viewer.add_geom(text_line)
self.text_lines.append(text_line)
def _write_values(self, index: int, message: str):
self.text_lines[index].set_text(message)
# keyboard event callbacks
def _key_press(self, k, mod):
from pyglet.window import key
agent_range = self.agents[self.current_agent_index].action.u_range_tensor
try:
if k == key.LEFT:
self.keys[0] = agent_range[0]
elif k == key.RIGHT:
self.keys[1] = agent_range[0]
elif k == key.DOWN:
self.keys[2] = agent_range[1]
elif k == key.UP:
self.keys[3] = agent_range[1]
elif k == key.M:
self.keys[4] = agent_range[2]
elif k == key.N:
self.keys[5] = agent_range[2]
elif k == key.TAB:
self.current_agent_index = self._increment_selected_agent_index(
self.current_agent_index
)
if self.control_two_agents:
while self.current_agent_index == self.current_agent_index2:
self.current_agent_index = self._increment_selected_agent_index(
self.current_agent_index
)
if self.control_two_agents:
agent2_range = self.agents[
self.current_agent_index2
].action.u_range_tensor
if k == key.A:
self.keys2[0] = agent2_range[0]
elif k == key.D:
self.keys2[1] = agent2_range[0]
elif k == key.S:
self.keys2[2] = agent2_range[1]
elif k == key.W:
self.keys2[3] = agent2_range[1]
elif k == key.E:
self.keys2[4] = agent2_range[2]
elif k == key.Q:
self.keys2[5] = agent2_range[2]
elif k == key.LSHIFT:
self.current_agent_index2 = self._increment_selected_agent_index(
self.current_agent_index2
)
while self.current_agent_index == self.current_agent_index2:
self.current_agent_index2 = (
self._increment_selected_agent_index(
self.current_agent_index2
)
)
except IndexError:
print("Action not available")
if k == key.R:
self.reset = True
self.set_u()
def _key_release(self, k, mod):
from pyglet.window import key
if k == key.LEFT:
self.keys[0] = 0
elif k == key.RIGHT:
self.keys[1] = 0
elif k == key.DOWN:
self.keys[2] = 0
elif k == key.UP:
self.keys[3] = 0
elif k == key.M:
self.keys[4] = 0
elif k == key.N:
self.keys[5] = 0
if self.control_two_agents:
if k == key.A:
self.keys2[0] = 0
elif k == key.D:
self.keys2[1] = 0
elif k == key.S:
self.keys2[2] = 0
elif k == key.W:
self.keys2[3] = 0
elif k == key.E:
self.keys2[4] = 0
elif k == key.Q:
self.keys2[5] = 0
self.set_u()
def set_u(self):
if self.continuous:
self.u = [
self.keys[1] - self.keys[0],
self.keys[3] - self.keys[2],
self.keys[4] - self.keys[5],
]
self.u2 = [
self.keys2[1] - self.keys2[0],
self.keys2[3] - self.keys2[2],
self.keys2[4] - self.keys2[5],
]
else:
if np.sum(self.keys[:4]) >= 1:
self.u[0] = np.argmax(self.keys[:4]) + 1
else:
self.u[0] = 0
if np.sum(self.keys[4:]) >= 1:
self.u[1] = np.argmax(self.keys[4:]) + 1
else:
self.u[1] = 0
if np.sum(self.keys2[:4]) >= 1:
self.u2[0] = np.argmax(self.keys2[:4]) + 1
else:
self.u2[0] = 0
if np.sum(self.keys2[4:]) >= 1:
self.u2[1] = np.argmax(self.keys2[4:]) + 1
else:
self.u2[1] = 0
@staticmethod
def format_obs(obs):
if isinstance(obs, (Tensor, np.ndarray)):
return list(np.around(obs.tolist(), decimals=2))
elif isinstance(obs, Dict):
return {key: InteractiveEnv.format_obs(value) for key, value in obs.items()}
else:
raise NotImplementedError(f"Invalid type of observation {obs}")
[docs]
def render_interactively(
scenario: Union[str, BaseScenario],
control_two_agents: bool = False,
display_info: bool = True,
save_render: bool = False,
**kwargs,
):
"""Executes a scenario and renders it so that you can debug and control agents interactively.
You can change the agent to control by pressing TAB.
You can reset the environment by pressing R.
You can control agent actions with the arrow keys and M/N (left/right control the first action, up/down control the second, M/N controls the third)
If you have more than 1 agent, you can control another one with W,A,S,D and Q,E in the same way.
and switch the agent using LSHIFT.
Args:
scenario (Union[str, BaseScenario]): Scenario to load.
Can be the name of a file in `vmas.scenarios` folder or a :class:`~vmas.simulator.scenario.BaseScenario` class
control_two_agents (bool, optional): Whether to control two agents or just one. Defaults to ``False``.
display_info (bool, optional): Whether to display on the screen the following info from the first controlled agent:
name, reward, total reward, done, and observation. Defaults to ``True``.
save_render (bool, optional): Whether to save a video of the render up to the first reset.
The video will be saved in the directory of this file with the name ``{scenario}_interactive``.
Defaults to ``False``.
Examples:
>>> from vmas import render_interactively
>>> render_interactively(
... "waterfall",
... control_two_agents=True,
... save_render=False,
... display_info=True,
... )
"""
InteractiveEnv(
make_env(
scenario=scenario,
num_envs=1,
device="cpu",
continuous_actions=True,
wrapper="gym",
seed=0,
wrapper_kwargs={"return_numpy": False},
# Environment specific variables
**kwargs,
),
control_two_agents=control_two_agents,
display_info=display_info,
save_render=save_render,
render_name=f"{scenario}_interactive"
if isinstance(scenario, str)
else "interactive",
)
def parse_args():
parser = ArgumentParser(description="Interactive rendering")
parser.add_argument(
"--scenario",
type=str,
default="waterfall",
help="Scenario to load. Can be the name of a file in `vmas.scenarios` folder or a :class:`~vmas.simulator.scenario.BaseScenario` class",
)
parser.add_argument(
"--control_two_agents",
action=BooleanOptionalAction,
default=True,
help="Whether to control two agents or just one",
)
parser.add_argument(
"--display_info",
action=BooleanOptionalAction,
default=True,
help="Whether to display on the screen the following info from the first controlled agent: name, reward, total reward, done, and observation",
)
parser.add_argument(
"--save_render",
action="store_true",
help="Whether to save a video of the render up to the first reset",
)
return parser.parse_args()
if __name__ == "__main__":
# Use this script to interactively play with scenarios
#
# You can change agent by pressing TAB
# You can reset the environment by pressing R
# You can control agent actions with the arrow keys and M/N (left/right control the first action, up/down control the second, M/N controls the third)
# If you have more than 1 agent, you can control another one with W,A,S,D and Q,E in the same way.
# and switch the agent with these controls using LSHIFT
args = parse_args()
render_interactively(
scenario=args.scenario,
control_two_agents=args.control_two_agents,
save_render=args.save_render,
display_info=args.display_info,
)