Action masking Problem

Hello,

I am new to Ray and I am trying to simulate a vehicule test environment. The agent has to pick a test at each time step and depending on the agent’s advancement some tests cannot be picked.
I have tried to adapt the action masking available on the net to my problem but nothing seems to work.
I keep getting this error:

Failure # 1 (occurred at 2022-07-11_13-14-05)
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/ray/tune/ray_trial_executor.py", line 934, in get_next_executor_event
    future_result = ray.get(ready_future)
  File "/usr/local/lib/python3.7/dist-packages/ray/_private/client_mode_hook.py", line 105, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/ray/worker.py", line 1833, in get
    raise value
  File "python/ray/_raylet.pyx", line 799, in ray._raylet.task_execution_handler
  File "python/ray/_raylet.pyx", line 618, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 762, in ray._raylet.execute_task
ray.exceptions.RayActorError: The actor died because of an error raised in its creation task, e[36mray::PPOTrainer.__init__()e[39m (pid=1349, ip=172.28.0.2, repr=PPOTrainer)
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/agents/trainer.py", line 871, in __init__
    config, logger_creator, remote_checkpoint_dir, sync_function_tpl
  File "/usr/local/lib/python3.7/dist-packages/ray/tune/trainable.py", line 156, in __init__
    self.setup(copy.deepcopy(self.config))
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/agents/trainer.py", line 957, in setup
    logdir=self.logdir,
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/evaluation/worker_set.py", line 144, in __init__
    lambda p, pid: (pid, p.observation_space, p.action_space)
ray.exceptions.RayActorError: The actor died because of an error raised in its creation task, e[36mray::RolloutWorker.__init__()e[39m (pid=1396, ip=172.28.0.2, repr=<ray.rllib.evaluation.rollout_worker.RolloutWorker object at 0x7f27c0f0db90>)
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/evaluation/rollout_worker.py", line 634, in __init__
    seed=seed,
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/evaluation/rollout_worker.py", line 1789, in _build_policy_map
    name, orig_cls, obs_space, act_space, conf, merged_conf
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/policy/policy_map.py", line 147, in create_policy
    observation_space, action_space, merged_config
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/policy/eager_tf_policy.py", line 450, in __init__
    stats_fn=stats_fn,
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/policy/policy.py", line 905, in _initialize_loss_from_dummy_batch
    self._dummy_batch, explore=False
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/policy/eager_tf_policy.py", line 502, in compute_actions_from_input_dict
    timestep,
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/utils/threading.py", line 21, in wrapper
    return func(self, *a, **k)
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/policy/eager_tf_policy.py", line 862, in _compute_actions_helper
    extra_fetches.update(extra_action_out_fn(self))
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/agents/ppo/ppo_tf_policy.py", line 198, in vf_preds_fetches
    SampleBatch.VF_PREDS: policy.model.value_function(),
  File "/usr/local/lib/python3.7/dist-packages/ray/rllib/models/modelv2.py", line 152, in value_function
    raise NotImplementedError
NotImplementedError

You can find my full code here:

Any suggestion or advice is more than welcome! Thank you.