Environment fails in between the process with PPO training.

[jaygupta907]

The model trains fine for some steps but after some steps the environment stops returning the observation.Please help. @jcallaham

[jcallaham]

When reporting issues, please give enough information to fully reproduce the issue, such as:

• Are you using a Docker configuration or local install?
• What are your config files?
• What specific command are you running?
• Also copy paste (not screenshot) a full error trace

Without this information it is near impossible to figure out how to help you. Same goes for #201 (and any other open source project you raise an issue with).

Then one of us can try to take a look (Edit: didn't mean to close, just pressed the wrong button).

[jaygupta907]

(RolloutWorker pid=42222) (tensor([0.3506], device='cuda:0'), tensor([1.4527], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.4126], device='cuda:0'), tensor([1.4473], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3709], device='cuda:0'), tensor([1.4498], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.4735], device='cuda:0'), tensor([1.4432], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3288], device='cuda:0'), tensor([1.4593], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3631], device='cuda:0'), tensor([1.4540], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.2681], device='cuda:0'), tensor([1.4552], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.4095], device='cuda:0'), tensor([1.4376], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.4399], device='cuda:0'), tensor([1.4402], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3800], device='cuda:0'), tensor([1.4540], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3941], device='cuda:0'), tensor([1.4573], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.4109], device='cuda:0'), tensor([1.4585], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3660], device='cuda:0'), tensor([1.4648], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.2296], device='cuda:0'), tensor([1.4724], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3896], device='cuda:0'), tensor([1.4514], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3631], device='cuda:0'), tensor([1.4557], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.3489], device='cuda:0'), tensor([1.4627], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.1572], device='cuda:0'), tensor([1.4760], device='cuda:0'))
(RolloutWorker pid=42222) (tensor([0.2619], device='cuda:0'), tensor([1.4538], device='cuda:0'))
2025-03-11 02:32:07,260 WARNING rollout_ops.py:115 -- No samples returned from remote workers. If you have a slow environment or model, consider increasing the sample_timeout_s or decreasing the rollout_fragment_length in AlgorithmConfig.env_runners(). 2025-03-11 02:32:07,260 WARNING deprecation.py:50 -- DeprecationWarning: ray.rllib.execution.train_ops.multi_gpu_train_one_step` has been deprecated. This will raise an error in the future!
2025-03-11 02:32:07,260 DEBUG train_ops.py:169 -- == sgd epochs for default_policy ==
Traceback (most recent call last):
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py", line 1348, in _worker
self.loss(model, self.dist_class, sample_batch)
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_torch_policy.py", line 84, in loss
logits, state = model(train_batch)
^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/models/modelv2.py", line 244, in call
input_dict["obs"], self.obs_space, self.framework
~~~~~~~~~~^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/sample_batch.py", line 950, in getitem
value = dict.getitem(self, key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'obs'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/jay/hydrogym/examples/cylinder/rllib/ppo_train.py", line 100, in
result = trainer.train()
^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/tune/trainable/trainable.py", line 331, in train
raise skipped from exception_cause(skipped)
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/tune/trainable/trainable.py", line 328, in train
result = self.step()
^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm.py", line 870, in step
train_results, train_iter_ctx = self._run_one_training_iteration()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/algorithms/algorithm.py", line 3160, in _run_one_training_iteration
training_step_results = self.training_step()
^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo.py", line 428, in training_step
return self._training_step_old_and_hybrid_api_stacks()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo.py", line 596, in _training_step_old_and_hybrid_api_stacks
train_results = multi_gpu_train_one_step(self, train_batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/execution/train_ops.py", line 176, in multi_gpu_train_one_step
results = policy.learn_on_loaded_batch(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py", line 838, in learn_on_loaded_batch
return self.learn_on_batch(batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/utils/threading.py", line 24, in wrapper
return func(self, *a, **k)
^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py", line 715, in learn_on_batch
grads, fetches = self.compute_gradients(postprocessed_batch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/utils/threading.py", line 24, in wrapper
return func(self, *a, **k)
^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py", line 933, in compute_gradients
tower_outputs = self._multi_gpu_parallel_grad_calc([postprocessed_batch])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py", line 1433, in _multi_gpu_parallel_grad_calc
raise last_result[0] from last_result[1]
ValueError: obs
tracebackTraceback (most recent call last):
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/torch_policy_v2.py", line 1348, in _worker
self.loss(model, self.dist_class, sample_batch)
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/algorithms/ppo/ppo_torch_policy.py", line 84, in loss
logits, state = model(train_batch)
^^^^^^^^^^^^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/models/modelv2.py", line 244, in call
input_dict["obs"], self.obs_space, self.framework
~~~~~~~~~~^^^^^^^
File "/home/jay/miniconda3/envs/hydrogym/lib/python3.12/site-packages/ray/rllib/policy/sample_batch.py", line 950, in getitem
value = dict.getitem(self, key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'obs'

In tower 0 on device cpu
PyOP2 Finalizing
Calling gc.collect()
STATE0
PYOP2 Communicator reference counts:
| Communicator name | Count |

Freeing PYOP2_COMM_WORLD
STATE1
PYOP2 Communicator reference counts:
| Communicator name | Count |

Freeing PYOP2_COMM_SELF
STATE2
PYOP2 Communicator reference counts:
| Communicator name | Count |

This is the complete error trace I am getting. I am running it in local install.

[jaygupta907]

import os
import ray
from common import *
from common import TorchCustomModel, parser
from firedrake import logging
from ray import tune
from ray.rllib.algorithms import ppo  # ray.rllib.algorithms in latest version
from ray.rllib.models import ModelCatalog
from ray.rllib.utils.test_utils import check_learning_achieved
from ray.tune.logger import pretty_print
import hydrogym as hgym
import numpy as np
import psutil



precomputed_data = "../output"
env_config = {
    "flow": hgym.firedrake.RotaryCylinder,
    "flow_config": {
        "Re" :100,
        "restart": f"{precomputed_data}/checkpoint.h5",
        "mesh":'medium',
        "velocity_order":1
    },
    "solver_config":{
        "dt":0.01,
    },
    "solver": hgym.firedrake.SemiImplicitBDF,
}

logging.set_log_level(logging.DEBUG)

if __name__ == "__main__":
    args = parser.parse_args()
    print(f"Running with following CLI options: {args}")

    tune.register_env("cylinder", lambda config: hgym.FlowEnv(env_config))
    ray.init(local_mode=args.local_mode)
  


    # Can also register the env creator function explicitly with:
    # register_env("corridor", lambda config: SimpleCorridor(config))
    ModelCatalog.register_custom_model("cav_actor", TorchCustomModel)


    def log_postprocess(flow):
        mem_usage = psutil.virtual_memory().percent
        CL, CD = flow.get_observations()
        return CL, CD, mem_usage
    
    # Set up the printing callback
    print_fmt = "t: {0:0.2f},\t\t CL: {1:0.3f},\t\t CD: {2:0.03f}\t\t Mem: {3:0.1f}"  # This will format the output
    log = hgym.firedrake.utils.io.LogCallback(
        postprocess=log_postprocess,
        nvals=3,
        interval=1,
        print_fmt=print_fmt,
        filename=None,
    )

    config = {
        "log_level": "DEBUG",
        "env": "cylinder",
        "env_config": {
            "Re": 100,
            "checkpoint": "./output/checkpoint.h5",
            "mesh": "coarse",
            "callbacks": [log],
            "max_steps": 1000,
        },
        # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
        "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")),
        "model": {
            "custom_model": "cav_actor",
            "vf_share_layers": True,
        },
        "num_workers": 1, 
    }

    stop = {
        "training_iteration": args.stop_iters,
        "timesteps_total": args.stop_timesteps,
        "episode_reward_mean": args.stop_reward,
    }

    if not args.tune:
        # manual training with train loop using PPO and fixed learning rate
        if args.run != "PPO":
            raise ValueError("Only support --run PPO with --no-tune.")
        print("Running manual train loop without Ray Tune.")
        ppo_config = ppo.PPOConfig()
        ppo_config = ppo_config.update_from_dict(config_dict=config)
        ppo_config["lr"] = 1e-3
        trainer = ppo.PPO(config=ppo_config, env="cylinder")
        # run manual training loop and print results after each iteration
        for i in range(args.stop_iters):
            result = trainer.train()
            print(pretty_print(result))
            trainer.save("./rllib_checkpoint")
            # stop training of the target train steps or reward are reached
            if (result["timesteps_total"] >= args.stop_timesteps or
                result["episode_reward_mean"] >= args.stop_reward):
                break
    else:
        # automated run with Tune and grid search and TensorBoard
        print("Training automatically with Ray Tune")
        results = tune.run(args.run, config=config, stop=stop)

        if args.as_test:
            print("Checking if learning goals were achieved")
            check_learning_achieved(results, args.stop_reward)

    ray.shutdown()

This the ppo_train.py file that I am running it's similar to the ppo_train provided for cavity environment in the github just for different environment.

[jaygupta907]

I did try increasing sample_timeout_s and decreaseing rollout_fragment_length as mentioned in the error still the error is same .
Please let me if any other information is needed.

[jcallaham]

@ludgerpaehler does anything jump out at you here? I haven't had a chance to run the script myself. Could it be a versioning thing with gymnasium or rllib?

[ludgerpaehler]

I think rllib is doing something funny here - I'll try to debug it

@jaygupta907 can you maybe provide the exact versions you are running? I.e. pip freeze > requirements.txt and then paste the requirements file in here? :)

[jaygupta907]

absl-py==2.1.0
aiosignal==1.3.2
asttokens @ file:///home/conda/feedstock_root/build_artifacts/asttokens_1733250440834/work
attrs==25.1.0
cachetools==5.5.2
certifi==2025.1.31
cgen==2020.1
charset-normalizer==3.4.1
checkpoint_schedules==1.0.4
click==8.1.8
cloudpickle==3.1.1
codepy==2023.1
colorama==0.4.6
comm @ file:///home/conda/feedstock_root/build_artifacts/comm_1733502965406/work
constantdict==2025.1.1
contourpy==1.3.1
control==0.9.4
cycler==0.12.1
debugpy @ file:///croot/debugpy_1736267418885/work
decorator==4.4.2
dm-tree==0.1.9
dmsuite==0.1.1
exceptiongroup @ file:///home/conda/feedstock_root/build_artifacts/exceptiongroup_1733208806608/work
executing @ file:///home/conda/feedstock_root/build_artifacts/executing_1733569351617/work
Farama-Notifications==0.0.4
fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@637894a00d0257e00f9ce3254e4dfb4ec3e7aa42
fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git@dcad2e4a9115adbe9042c7fdb3ed8c43ff28fc14
filelock==3.17.0
firedrake @ git+https://github.com/firedrakeproject/firedrake.git@b35aab5a94736d092bb76bcec5de255994aaa05e
fonttools==4.56.0
frozenlist==1.5.0
fsspec==2025.2.0
genpy==2022.1
gmsh==4.13.1
GPUtil==1.4.0
gutils==0.1.9
gym==0.26.2
gym-notices==0.0.8
gymnasium==0.26.1
gymnasium-notices==0.0.1
h5py==3.13.0
hydrogym @ file:///home/jay/hydrogym
idna==3.10
imageio==2.37.0
immutabledict==4.2.1
importlib_metadata @ file:///home/conda/feedstock_root/build_artifacts/importlib-metadata_1737420181517/work
iniconfig==2.0.0
ipykernel @ file:///home/conda/feedstock_root/build_artifacts/ipykernel_1719845459717/work
ipython @ file:///home/conda/feedstock_root/build_artifacts/bld/rattler-build_ipython_1741120654/work
ipython_pygments_lexers @ file:///home/conda/feedstock_root/build_artifacts/ipython_pygments_lexers_1737123620466/work
islpy==2025.1.2
jedi @ file:///home/conda/feedstock_root/build_artifacts/jedi_1733300866624/work
Jinja2==3.1.6
joblib==1.4.2
jsonschema==4.23.0
jsonschema-specifications==2024.10.1
jupyter_client @ file:///home/conda/feedstock_root/build_artifacts/jupyter_client_1733440914442/work
jupyter_core @ file:///home/conda/feedstock_root/build_artifacts/jupyter_core_1727163409502/work
kiwisolver==1.4.8
lazy_loader==0.4
libsupermesh==2025.3.0
loopy @ git+https://github.com/firedrakeproject/loopy.git@27aead574b5188de8e0d98518c93f1863951fc4b
lz4==4.4.3
Mako==1.3.9
markdown-it-py==3.0.0
MarkupSafe==3.0.2
matplotlib==3.10.1
matplotlib-inline @ file:///home/conda/feedstock_root/build_artifacts/matplotlib-inline_1733416936468/work
mdurl==0.1.2
mpi-pytest==2025.2.0
mpi4py==4.0.3
mpmath==1.3.0
msgpack==1.1.0
nest_asyncio @ file:///home/conda/feedstock_root/build_artifacts/nest-asyncio_1733325553580/work
networkx==3.4.2
numpy==2.2.3
nvidia-cublas-cu12==12.4.5.8
nvidia-cuda-cupti-cu12==12.4.127
nvidia-cuda-nvrtc-cu12==12.4.127
nvidia-cuda-runtime-cu12==12.4.127
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.2.1.3
nvidia-curand-cu12==10.3.5.147
nvidia-cusolver-cu12==11.6.1.9
nvidia-cusparse-cu12==12.3.1.170
nvidia-cusparselt-cu12==0.6.2
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.4.127
nvidia-nvtx-cu12==12.4.127
packaging==24.2
pandas==2.2.3
parso @ file:///home/conda/feedstock_root/build_artifacts/parso_1733271261340/work
petsc4py==3.22.3
pexpect @ file:///home/conda/feedstock_root/build_artifacts/pexpect_1733301927746/work
pickleshare @ file:///home/conda/feedstock_root/build_artifacts/pickleshare_1733327343728/work
pillow==11.1.0
pkgconfig==1.5.5
platformdirs==4.3.6
pluggy==1.5.0
progress==1.6
prompt_toolkit @ file:///home/conda/feedstock_root/build_artifacts/prompt-toolkit_1737453357274/work
protobuf==6.30.0
psutil==7.0.0
ptyprocess @ file:///home/conda/feedstock_root/build_artifacts/ptyprocess_1733302279685/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl#sha256=92c32ff62b5fd8cf325bec5ab90d7be3d2a8ca8c8a3813ff487a8d2002630d1f
pure_eval @ file:///home/conda/feedstock_root/build_artifacts/pure_eval_1733569405015/work
pyadjoint-ad @ git+https://github.com/dolfin-adjoint/pyadjoint.git@dbf923e656b37333e45267da2e0264a84f160be5
pyarrow==19.0.1
pycparser==2.22
Pygments @ file:///home/conda/feedstock_root/build_artifacts/pygments_1736243443484/work
pymbolic==2024.2.2
pyparsing==3.2.1
pytest==8.3.5
python-dateutil==2.9.0.post0
pytools==2025.1.1
pytz==2025.1
PyYAML==6.0.2
pyzmq @ file:///croot/pyzmq_1734687138743/work
ray==2.31.0
recursivenodes==0.2.0
referencing==0.36.2
requests==2.32.3
rich==13.9.4
rpds-py==0.23.1
rtree==1.4.0
scikit-image==0.25.2
scikit-learn==1.6.1
scipy==1.15.2
seaborn==0.13.2
setuptools==75.8.0
shellingham==1.5.4
siphash24==1.7
six==1.17.0
stack_data @ file:///home/conda/feedstock_root/build_artifacts/stack_data_1733569443808/work
svgwrite==1.4.3
symengine==0.14.0
sympy==1.13.1
tensorboardX==2.6.2.2
threadpoolctl==3.5.0
tifffile==2025.2.18
torch==2.6.0
tornado @ file:///croot/tornado_1733960490606/work
traitlets @ file:///home/conda/feedstock_root/build_artifacts/traitlets_1733367359838/work
Tree==0.2.4
triton==3.2.0
typer==0.15.2
typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1733188668063/work
tzdata==2025.1
urllib3==2.3.0
vtk==9.4.1
wcwidth @ file:///home/conda/feedstock_root/build_artifacts/wcwidth_1733231326287/work
wheel==0.45.1
wrapt==1.17.2
zipp @ file:///home/conda/feedstock_root/build_artifacts/zipp_1732827521216/work

[jaygupta907]

Hi @ludgerpaehler are you able to find the issue here?

[ludgerpaehler]

Still on it, but I think I am managing to pinpoint it by now