cloudpickle cannot pickle '_jpype._JField' objects

[framepixel]

So, I've been working on a project which involves implementing reinforcement learning in a server-client app. The server is written in Java and the client is in Python, which is why I use JPype to import some server classes.

After importing the necessary packages and creating the environment using PettingZoo, it is time to create the model and train it using Stable-Baselines3, but the problem is that when I use Supersuit, it needs to pickle and unpickle the environment, and because the environment contains many Java objects, an error is thrown: TypeError: cannot pickle '_jpype._JField' object.

The normal Pickle package does not support JField objects, but in the JPype library, there is a JPickle version that supports JField objects. I tried to modify the cloudpickle_fast.py to add the JPickle package but I end up having a problem with the cloudpickle.loads()

Here is what I modified in cloudpickle_fast.py:

from jpype.pickle import JPickler, JUnpickler

    def dump(self, obj):
        try:
            return Pickler.dump(self, obj)
        except RuntimeError as e:
            if "recursion" in e.args[0]:
                msg = (
                    "Could not pickle object as excessively deep recursion "
                    "required."
                )
                raise pickle.PicklingError(msg) from e
            else:
                raise
        except TypeError as e:
            return JPickler.dump(self, obj)

And here is the full stacktrace I get:

---------------------------------------------------------------------------
UnpicklingError                           Traceback (most recent call last)
Input In [11], in <cell line: 6>()
      1 env = MARL_Env_Parallel(4)
      5 env = ss.pettingzoo_env_to_vec_env_v1(env)
----> 6 env = ss.concat_vec_envs_v1(env, 1, num_cpus=1, base_class='stable_baselines3')

File ~\anaconda3\envs\gym\lib\site-packages\supersuit\vector\vector_constructors.py:61, in concat_vec_envs_v1(vec_env, num_vec_envs, num_cpus, base_class)
     59 def concat_vec_envs_v1(vec_env, num_vec_envs, num_cpus=0, base_class="gymnasium"):
     60     num_cpus = min(num_cpus, num_vec_envs)
---> 61     vec_env = MakeCPUAsyncConstructor(num_cpus)(*vec_env_args(vec_env, num_vec_envs))
     63     if base_class == "gymnasium":
     64         return vec_env

File ~\anaconda3\envs\gym\lib\site-packages\supersuit\vector\concat_vec_env.py:22, in ConcatVecEnv.__init__(self, vec_env_fns, obs_space, act_space)
     21 def __init__(self, vec_env_fns, obs_space=None, act_space=None):
---> 22     self.vec_envs = vec_envs = [vec_env_fn() for vec_env_fn in vec_env_fns]
     23     for i in range(len(vec_envs)):
     24         if not hasattr(vec_envs[i], "num_envs"):

File ~\anaconda3\envs\gym\lib\site-packages\supersuit\vector\concat_vec_env.py:22, in <listcomp>(.0)
     21 def __init__(self, vec_env_fns, obs_space=None, act_space=None):
---> 22     self.vec_envs = vec_envs = [vec_env_fn() for vec_env_fn in vec_env_fns]
     23     for i in range(len(vec_envs)):
     24         if not hasattr(vec_envs[i], "num_envs"):

File ~\anaconda3\envs\gym\lib\site-packages\supersuit\vector\vector_constructors.py:11, in vec_env_args.<locals>.env_fn()
     10 def env_fn():
---> 11     env_copy = cloudpickle.loads(cloudpickle.dumps(env))
     12     return env_copy

UnpicklingError: Memo value not found at index 3

I don't have a lot of experience with Pickle, so any advice would be welcome, thanks.