Getting always the same action on an A2C from stable_baselines3

Question

I'm quite new to RL and have been trying to train an A2C model from stable_baselines3 to derive an integer sequence based on 3 other input sequences of floats. I have a custom gym environment that computes the agent rewards based on the actions taken on the step method. Spaces are:

self.observation_space = spaces.Dict({
    "s1": spaces.Box(
        low=0,
        high=float("inf"),
        shape=(48,),
        dtype=np.float64
    ),
    "s2": spaces.Box(
        low=0,
        high=float("inf"),
        shape=(48,),
        dtype=np.float64
    ),
    "s3": spaces.Box(
        low=0,
        high=float("inf"),
        shape=(48,),
        dtype=np.float64
    ),
})
self.action_space = spaces.MultiDiscrete([5]*48)

The observables do not change based on the action, so there is not state carried over between step calls; for each obs a given action will maximise the reward. Then the environment would transition to the next state with a completely new obs that will have an unrelated new optimal action.

For training I run:

env_train = Monitor(gym.make("gym/MyCustomEnvironment-v0", data=data_train), log_dir)
model = A2C("MultiInputPolicy", env_train, verbose=0, device="cpu")    
model.learn(total_timesteps=timesteps)

The rewards/timesteps plot converges towards a value in a logarithmic curve. Even though when I evaluate the trained model on a new observation I always get the same set of actions.

env = Monitor(gym.make("gym/MyCustomEnvironment-v0", data=data_test), log_dir)
obs, _ = env.reset()
action, _ = model.predict(obs, deterministic=True)

Here the value of action will be the same or very similar regardless of the data used to create the environment (which is then used to derive obs). This is very surprising because the optimal policy that would maximise training rewards would be to choose different output sequences depending on the input sequence.

I have tried multiple things, including training with large datasets, changing learning rates, gamma and network topology, playing with episode lengths, etc. I have even changed the reward function to one that produces arbitrary numbers and even in that case the "trained" model will always output the same actions regardless of obs.

Am I just misusing A2C for a use-case it's not intended for?

Answer 1

The environment you described fits a multi-armed bandit problem where the action does not influence the next state encountered. Algorithms such as A2C are designed for situations where the state transition is dependent on the action. If this is not the case, you should look for multi-armed bandit type algorithms that are more suited for the job.