CleanDirty.py

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import time

env={
    "Room1":"Clean",
    "Room2":"Dirty",
    "Room3":"Clean",
    "Room4":"Dirty"
}
room_pos={
   "Room1":(0,1),
    "Room2":(1,1),
    "Room3":(1,0,) ,
    "Room4":(0,0)
}
rooms=list(env.keys())
agent=0

def ref(state):
    if state=="Dirty":
        return "Clean"
    else:
        return "Move"
    
def draw(env,agent,step):
    fig,ax=plt.subplots()
    ax.set_xlim(0,2)
    ax.set_ylim(0,2)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title(f"Step {step} - Agent in {rooms[agent]} ")

    for rm,pos in room_pos.items():
        x,y=pos
        color='red' if env[rm]=="Dirty" else "green"
        rect=patches.Rectangle((x,y),1,1,facecolor=color,edgecolor='black')
        ax.add_patch(rect)
        ax.text(x+0.5,y+0.5,rm,ha='center',va='center',color='white',fontsize=10)
    agx,agy=room_pos[rooms[agent]]
    agent_patch=patches.Circle((agx+0.5,agy+0.5),0.1,color='blue')
    ax.add_patch(agent_patch)
    plt.pause(3)
    plt.close()

#Run simulation
plt.ion() #Interaction mode on
step=8
for s in range(step):
    curr=rooms[agent]
    state=env[curr]
    action=ref(state)
     
    draw(env,agent,s+1)

    if action=="Clean":
        env[curr]="Clean"
    else:
        agent=(agent+1)%len(rooms)

plt.ioff()
print("Simulation completed")