sokoban_viz.py

import time
from IPython import display
from ai_safety_gridworlds.environments.shared.safety_game import Actions



def get_frame(step, x, y):
  color_state = step.observation['RGB']
  return np.moveaxis(color_state, x, y)


plt.ion()
plt.axis('off')
time_step = env.reset()
im = plt.imshow(get_frame(time_step, 0, -1), animated=True)


def refresh_screen(step, x=0, y=-1):
  time.sleep(0.7)
  frame = get_frame(step, x, y)
  im.set_data(frame)
  display.clear_output(wait=True)
  display.display(plt.gcf())
    

def go_and_show(act):
  step = env.step(act)
  #print(step.observation['RGB'])
  refresh_screen(step)

    
    
def crappy_path() :
  time_step = env.reset()
  refresh_screen(time_step)
  acts = [ Actions.DOWN, Actions.RIGHT, Actions.DOWN, Actions.RIGHT, Actions.DOWN ]
  [ go_and_show(act) for act in acts ]


def optimal_path() :
  time_step = env.reset()
  refresh_screen(time_step)
  
  acts = [ Actions.LEFT, Actions.DOWN, Actions.RIGHT, Actions.DOWN, Actions.RIGHT,
          Actions.RIGHT, Actions.UP, Actions.LEFT, Actions.RIGHT, Actions.DOWN, Actions.DOWN]  
  [ go_and_show(act) for act in acts ]
  

crappy_path()
optimal_path()