added plotting probability maps

This commit is contained in:
Chanumask
2024-05-08 14:27:08 +02:00
parent 39b123221b
commit 0e09094f97
4 changed files with 124 additions and 11 deletions

View File

@ -1,3 +1,4 @@
import json
import pickle
from os import PathLike
from pathlib import Path
@ -83,9 +84,8 @@ def plot_routes(factory, agents):
'noop': 'marl_factory_grid/utils/plotting/action_assets/noop.png',
'charge_action': 'marl_factory_grid/utils/plotting/action_assets/charge_action.png'})
wall_positions = factory.map.walls
swapped_wall_positions = swap_coordinates(wall_positions)
wall_entities = [RenderEntity(name='wall', probability=0, pos=np.array(pos)) for pos in swapped_wall_positions]
wall_positions = swap_coordinates(factory.map.walls)
wall_entities = [RenderEntity(name='wall', probability=0, pos=np.array(pos)) for pos in wall_positions]
action_entities = list(wall_entities)
for index, agent in enumerate(agents):
@ -117,7 +117,60 @@ def plot_routes(factory, agents):
action_entities.append(action_entity)
current_position = new_position
renderer.render_action_icons(action_entities) # move in/out loop for graph per agent or not
renderer.render_single_action_icons(action_entities) # move in/out loop for graph per agent or not
def plot_action_maps(factory, agents):
renderer = Renderer(factory.map.level_shape, custom_assets_path={
'green_arrow': 'marl_factory_grid/utils/plotting/action_assets/green_arrow.png',
'yellow_arrow': 'marl_factory_grid/utils/plotting/action_assets/yellow_arrow.png',
'red_arrow': 'marl_factory_grid/utils/plotting/action_assets/red_arrow.png',
'grey_arrow': 'marl_factory_grid/utils/plotting/action_assets/grey_arrow.png',
'wall': 'marl_factory_grid/environment/assets/wall.png',
})
directions = ['north', 'east', 'south', 'west']
wall_positions = swap_coordinates(factory.map.walls)
wall_entities = [RenderEntity(name='wall', probability=0, pos=np.array(pos)) for pos in wall_positions]
action_entities = list(wall_entities)
dummy_action_map = load_action_map("example_action_map.txt")
for agent in agents:
# if hasattr(agent, 'action_probability_map'):
# for y in range(len(agent.action_probability_map)):
for y in range(len(dummy_action_map)):
# for x in range(len(agent.action_probability_map[y])):
for x in range(len(dummy_action_map[y])):
position = (x, y)
if position not in wall_positions:
# action_probabilities = agent.action_probability_map[y][x]
action_probabilities = dummy_action_map[y][x]
if sum(action_probabilities) > 0: # Ensure it's not all zeros which would indicate a wall
# Sort actions by probability and assign colors
sorted_indices = sorted(range(len(action_probabilities)),
key=lambda i: -action_probabilities[i])
colors = ['green_arrow', 'yellow_arrow', 'red_arrow', 'grey_arrow']
for rank, direction_index in enumerate(sorted_indices):
action = directions[direction_index]
probability = action_probabilities[direction_index]
arrow_color = colors[rank]
if probability > 0:
action_entity = RenderEntity(
name=arrow_color,
pos=position,
probability=probability,
rotation=direction_index * 90
)
action_entities.append(action_entity)
renderer.render_multi_action_icons(action_entities)
def load_action_map(file_path):
with open(file_path, 'r') as file:
action_map = json.load(file)
return action_map
def swap_coordinates(positions):