How do I train a model for an environment with variable features using openai gym? - machine-learning

Question

I have a grid network (like a maze) and I want to find the shortest route from the source node to the destination node.
I've coded that and it works well. It finds the shortest route for a network of dimensions x*y and with the specific nodes as the source and destination nodes which are determined as inputs to the environment. env = netEnv(3, 4, 1, 8) means my env is a 3*4 grid network with node 1 as the source node and node 8 as the destination node.
The question is how to train the model for any destination? I mean the model works well if the destination is say node 8. But I want it to be trained so that it will find the answer if the destination is a different node in the environment. I don't want the destination to be always fixed.
How can I do that?
Here is my code:
#import dependecies
import gym
from gym import Env
from gym.spaces import Discrete
from stable_baselines3 import DQN
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.evaluation import evaluate_policy
import os
#building the environment
class netEnv(Env):
def __init__(self, x, y, sourceNode, destinationNode):
self.action_space = Discrete(4) #0, 1, 2, 3 => up, down, right, left
self.observation_space = Discrete(x*y)
self.sourceNode = sourceNode
self.destinationNode = destinationNode
self.x = x
self.y = y
self.state = self.sourceNode
self.episodeLength = 20
def step(self, action):
#set actions and rewards
if action == 0: #up
if self.state >= self.y:
self.state = self.state - self.y
#print(self.state)
reward = 0.1
else:
reward = 0
elif action == 1: #down
if self.state < (self.x-1)*self.y:
self.state = self.state + self.y
#print(self.state)
reward = 0.1
else:
reward = 0
elif action == 2: #right
if (self.state % self.y) != self.y-1:
self.state += 1
#print(self.state)
reward = 0.1
else:
reward = 0
else: #left
if (self.state % self.y) != 0:
self.state -= 1
#print(self.state)
reward = 0.1
else:
reward = 0
#decrease the remaining time
self.episodeLength -=1
#reward in termination state
if self.state == self.destinationNode:
reward += self.episodeLength
#set done
if self.episodeLength <= 0 or self.state == self.destinationNode:
done = True
else:
done = False
#set info
info = {}
return self.state, reward, done, info
def render(self):
pass
def reset(self):
self.state = self.sourceNode
self.episodeLength = 20
return self.state
env = netEnv(3, 4, 1, 8)
#training
log_path = os.path.join('Training', 'Logs')
model = PPO('MlpPolicy', env, verbose = 1, tensorboard_log = log_path)
model.learn(total_timesteps=15000)
net_path = os.path.join('Training', 'Saved_models', 'net_PPO_15000')
model.save(net_path)
#delete and reload the model
del model
model = PPO.load(net_path, env)
#evaluate the model
print(evaluate_policy(model, env, n_eval_episodes = 10, render = False))