mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-20 22:12:03 +00:00
c6a97e17ee
* Initial version of vectorized environments * Raise an exception in the main process if child process raises an exception * Add list of exposed functions in vector module * Use deepcopy instead of np.copy * Add documentation for vector utils * Add tests for copy in AsyncVectorEnv * Add example in documentation for batch_space * Add cloudpickle dependency in setup.py * Fix __del__ in VectorEnv * Check if all observation spaces are equal in AsyncVectorEnv * Check if all observation spaces are equal in SyncVectorEnv * Fix spaces non equality in SyncVectorEnv for Python 2 * Handle None parameter in create_empty_array * Fix check_observation_space with spaces equality * Raise an exception when operations are out of order in AsyncVectorEnv * Add version requirement for cloudpickle * Use a state instead of binary flags in AsyncVectorEnv * Use numpy.zeros when initializing observations in vectorized environments * Remove poll from public API in AsyncVectorEnv * Remove close_extras from VectorEnv * Add test between AsyncVectorEnv and SyncVectorEnv * Remove close in check_observation_space * Add documentation for seed and close * Refactor exceptions for AsyncVectorEnv * Close pipes if the environment raises an error * Add tests for out of order operations * Change default argument in create_empty_array to np.zeros * Add get_attr and set_attr methods to VectorEnv * Improve consistency in SyncVectorEnv
41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
import contextlib
|
|
import os
|
|
|
|
__all__ = ['CloudpickleWrapper', 'clear_mpi_env_vars']
|
|
|
|
class CloudpickleWrapper(object):
|
|
def __init__(self, fn):
|
|
self.fn = fn
|
|
|
|
def __getstate__(self):
|
|
import cloudpickle
|
|
return cloudpickle.dumps(self.fn)
|
|
|
|
def __setstate__(self, ob):
|
|
import pickle
|
|
self.fn = pickle.loads(ob)
|
|
|
|
def __call__(self):
|
|
return self.fn()
|
|
|
|
@contextlib.contextmanager
|
|
def clear_mpi_env_vars():
|
|
"""
|
|
`from mpi4py import MPI` will call `MPI_Init` by default. If the child
|
|
process has MPI environment variables, MPI will think that the child process
|
|
is an MPI process just like the parent and do bad things such as hang.
|
|
|
|
This context manager is a hacky way to clear those environment variables
|
|
temporarily such as when we are starting multiprocessing Processes.
|
|
"""
|
|
removed_environment = {}
|
|
for k, v in list(os.environ.items()):
|
|
for prefix in ['OMPI_', 'PMI_']:
|
|
if k.startswith(prefix):
|
|
removed_environment[k] = v
|
|
del os.environ[k]
|
|
try:
|
|
yield
|
|
finally:
|
|
os.environ.update(removed_environment)
|