BSK_RL: Issue with Serialization of Custom Environment When Using RLlib APPO with Multiple Workers

[Mark2000]

From @clementejuanoliver:

When attempting to use custom environments based on GeneralSatelliteTasking class with Ray/RLlib, I encounter a TypeError: cannot pickle 'SwigPyObject' object during the environment's initialization phase. This issue arises when I attempt to build an APPO algorithm trainer with a configuration that includes the GeneralSatelliteTasking environment (with both Gym and RLlib) and uses >0 workers.

I managed to train several agents with 0 workers (RLlib creates num_workers + 1 copies of the environment) and several environments per worker, but when using a number greater than 0 workers, config.build(), which calls the deepcopy function (which internally uses pickle function), gives the error.

The custom environment GeneralSatelliteTasking relies on external libraries that use SWIG for Python bindings, so I guess it would be better to change the env_config to serializable objects, but I would really like to use these ones since they provide with more flexibility and automatised parameters.

Any help or references to documentation/examples that tackle similar challenges would be greatly appreciated.

Thanks in advance!

env_config = {
    "satellites": satellites,
    "env_type": environment.GroundStationEnvModel,
    "env_args": environment.GroundStationEnvModel.default_env_args(),
    "env_features": env_features,
    "data_manager": data_manager,
    "communicator": communicator,
    "sim_rate": 0.5, # Integration frequency in seconds
    "max_step_duration": 600.0,
    "time_limit": 95 * 60, 
    "log_level": "INFO",
}

# Configure APPO training
config = APPOConfig()
config = config.environment(
    env="GeneralSatelliteTasking-v1",
    env_config=env_config,
    disable_env_checking=True,
)
config = config.framework("tf2") # tf2, tf, torch
config = config.resources(
    num_gpus=0
)
config = config.rollouts(
    num_rollout_workers=1,
    num_envs_per_worker=2,
)
config = config.training(
    train_batch_size=500,
    num_sgd_iter=3,
    lr=0.01,
    model={
        "fcnet_hiddens": [64, 64],
        "fcnet_activation": "relu"
    },
)

# Register the environment
# For Gym
env = gym.make(
        "GeneralSatelliteTasking-v1",
        **env_config
    )

# For RLlib:
register_env("GeneralSatelliteTasking-v1", lambda config: GeneralSatelliteTasking(**config))

# Create and build the trainer
trainer = config.build()

[Mark2000]

There seems to be an issue with how gym's registration function works in some multiprocessed cases due to SWIGed objects. This seems to be worst when you try to register and "incomplete" environments and then pass other environment parameters when making the environment.

The main workaround that I've found is to not use environment registration, which is a feature of gym that seems mostly useless imo. With many RL frameworks, you can pass the environment class and some kwargs dict that then gets unpacked and passed to the environment class constructor. However, RLLib in particular (I currently am using 2.6.3) has a strange API where the environment can only take one argument called env_config. As a result my code looks like this:

    config = APPOConfig().training(**ppo_args)
    config = config.environment(
        env=unpack_config(SingleSatelliteTasking),
        env_config=env_args,
        auto_wrap_old_gym_envs=False,
        disable_env_checking=True,
    )

where unpack_config is

def unpack_config(env):
    """Create a wrapped version of an env class that unpacks env_config from Ray into kwargs"""

    class UnpackedEnv(env):
        def __init__(self, env_config):
            super().__init__(**env_config)

    UnpackedEnv.__name__ = f"{env.__name__}_Unpacked"

    return UnpackedEnv

I plan to add some example training scripts when we get a major release of the repo together this summer. Hope this helps for now!

[clementejuanoliver]

Hi Mark,
It worked perfectly, even with multiple GeneralSatelliteTasking class satellites. It definitely helps, training is way more scalable now.
Thank you a lot!

[hithxh]

Hi, Mark and clementejuanoliver,
I tried appo with the following configuration, but the episode mean reward stops increasing after about 5 steps.
Did I miss some important item?

config = (
   APPOConfig()
   .environment(
       env=unpack_config(SingleSatelliteTasking),
       env_config=env_config,
   )
   .framework("torch") # tf2, tf, torch
   .resources(
       num_gpus=1
   )   
   .rollouts(
       num_rollout_workers=30,
       num_envs_per_worker=2,
       )
   .training(
       train_batch_size=10000,
       grad_clip=0.1,
       gamma=0.999,
       lr=0.01,
       model={
           "fcnet_hiddens": [640,640],
           "fcnet_activation": "relu",
       },
   )
)
trainer = config.build()
for _ in range(100000):
   result = trainer.train()
   print("episode mean reward: ", result["episode_reward_mean"] ) 

[Mark2000]

I've not had success with the torch backend (mainly when it comes to loading trained policies), but I don't think that would be the issue. I would double check that you're setting terminate_on_time_limit=True and that your satellite is subclassing sat_observations.TimeState (i.e. time is included in the observation), otherwise learning often doesn't work. None of the hyperparameters look obviously wrong to me.

[hithxh]

Here is the whole file, it doses set terminate_on_time_limit to True：

env_features = StaticTargets(n_targets=1000)
data_manager = data.UniqueImagingManager(env_features)
sat_type = sats.FullFeaturedSatellite

oes = leo_orbit.walker_delta(
    n_spacecraft=3,  # Number of satellites
    n_planes=3,
    rel_phasing=0,
    altitude=500 * 1e3,
    inc=45,
    clustersize=3,  # Cluster all 3 satellites together
    clusterspacing=30,  # Space satellites by a true anomaly of 30 degrees
)

satellites = []

for i, oe in enumerate(oes):
    sat_args = sat_type.default_sat_args(
        oe=oe,
        imageAttErrorRequirement=0.01,  # Change a default parameter
        imageRateErrorRequirement=0.01,
        panelEfficiency=lambda: 0.2 + np.random.uniform(-0.01, 0.01),
    )
    satellite = sat_type(
        "EO" + str(i + 1), sat_args, n_ahead_observe=15, n_ahead_act=15
    )
    satellites.append(satellite)

communicator = communication.LOSMultiCommunication(satellites)

env_config = {
    "satellites": satellites,
    "env_type": environment.GroundStationEnvModel,
    "env_args": environment.GroundStationEnvModel.default_env_args(),
    "env_features": env_features,
    "data_manager": data_manager,
    "communicator": communicator,
    "sim_rate": 0.5, # Integration frequency in seconds
    "max_step_duration": 600.0,
    "time_limit": 95 * 60 * 3,  
    "terminate_on_time_limit": True,
    "log_level": "warning",
    }

config = (
    APPOConfig()
    .environment(
        env=unpack_config(GeneralSatelliteTasking),
        env_config=env_config,
    )
    .framework("torch") # tf2, tf, torch
    .resources(
        num_gpus=1
    )   
    .rollouts(
        num_rollout_workers=32,
        num_envs_per_worker=2,
        )
    .training(
        train_batch_size=10000,
        lr=0.0003,
        clip_param = 0.1,
        use_critic=True,
        model={
            "fcnet_hiddens": [640,640],
            "fcnet_activation": "relu",
        },
    )
)

trainer = config.build()
for _ in range(100000):
    result = trainer.train()
    # print(result)

    print("episode len mean: ", result["episode_len_mean"],"episode mean reward: ", result["episode_reward_mean"] ) 
trainer.evaluate()

and the FullFeaturedSatellite class contains the TimeState for sure.

class FullFeaturedSatellite(
    sa.ImagingActions,
    sa.DesatAction,
    sa.ChargingAction,
    so.TimeState,
    so.EclipseState,
    so.TargetState.configure(n_ahead_observe=3),
    so.NormdPropertyState.configure(
        obs_properties=[
            dict(prop="omega_BP_P", norm=0.03),
            dict(prop="c_hat_P"),
            dict(prop="r_BN_P", norm=orbitalMotion.REQ_EARTH * 1e3),
            dict(prop="v_BN_P", norm=7616.5),
            dict(prop="battery_charge_fraction"),
        ]
    ),
    
    SteeringImagerSatellite,
):
    """Convenience type for a satellite with common features enabled."""

    pass