From 1ad585e4d7b1781c1425e281940b09ff5e707ae5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Antonio=20Serrano=20Mu=C3=B1oz?= Date: Sat, 22 Jun 2024 18:29:45 -0400 Subject: [PATCH] Get JAX device from string --- skrl/agents/jax/base.py | 5 ++- skrl/envs/wrappers/jax/base.py | 33 +++++++++++++------ skrl/memories/jax/base.py | 5 ++- skrl/models/jax/base.py | 5 ++- skrl/multi_agents/jax/base.py | 5 ++- skrl/resources/noises/jax/base.py | 5 ++- .../jax/running_standard_scaler.py | 5 ++- 7 files changed, 47 insertions(+), 16 deletions(-) diff --git a/skrl/agents/jax/base.py b/skrl/agents/jax/base.py index 561d2701..fbcb9fa5 100644 --- a/skrl/agents/jax/base.py +++ b/skrl/agents/jax/base.py @@ -53,7 +53,10 @@ def __init__(self, if device is None: self.device = jax.devices()[0] else: - self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0] + self.device = device + if type(device) == str: + device_type, device_index = f"{device}:0".split(':')[:2] + self.device = jax.devices(device_type)[int(device_index)] if type(memory) is list: self.memory = memory[0] diff --git a/skrl/envs/wrappers/jax/base.py b/skrl/envs/wrappers/jax/base.py index 0be8b742..b1136a75 100644 --- a/skrl/envs/wrappers/jax/base.py +++ b/skrl/envs/wrappers/jax/base.py @@ -20,12 +20,19 @@ def __init__(self, env: Any) -> None: self._env = env # device (faster than @property) - self.device = jax.devices()[0] + self.device = None if hasattr(self._env, "device"): - try: - self.device = jax.devices(self._env.device.split(':')[0] if type(self._env.device) == str else self._env.device.type)[0] - except RuntimeError: - pass + if type(self._env.device) == str: + device_type, device_index = f"{self._env.device}:0".split(':')[:2] + try: + self.device = jax.devices(device_type)[int(device_index)] + except RuntimeError: + self.device = None + else: + self.device = self._env.device + if self.device is None: + self.device = jax.devices()[0] + # spaces try: self._action_space = self._env.single_action_space @@ -135,12 +142,18 @@ def __init__(self, env: Any) -> None: self._env = env # device (faster than @property) - self.device = jax.devices()[0] + self.device = None if hasattr(self._env, "device"): - try: - self.device = jax.devices(self._env.device.split(':')[0] if type(self._env.device) == str else self._env.device.type)[0] - except RuntimeError: - pass + if type(self._env.device) == str: + device_type, device_index = f"{self._env.device}:0".split(':')[:2] + try: + self.device = jax.devices(device_type)[int(device_index)] + except RuntimeError: + self.device = None + else: + self.device = self._env.device + if self.device is None: + self.device = jax.devices()[0] self.possible_agents = [] diff --git a/skrl/memories/jax/base.py b/skrl/memories/jax/base.py index 41cea1b1..90783836 100644 --- a/skrl/memories/jax/base.py +++ b/skrl/memories/jax/base.py @@ -70,7 +70,10 @@ def __init__(self, if device is None: self.device = jax.devices()[0] else: - self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0] + self.device = device + if type(device) == str: + device_type, device_index = f"{device}:0".split(':')[:2] + self.device = jax.devices(device_type)[int(device_index)] # internal variables self.filled = False diff --git a/skrl/models/jax/base.py b/skrl/models/jax/base.py index 30db0e83..7d121e25 100644 --- a/skrl/models/jax/base.py +++ b/skrl/models/jax/base.py @@ -79,7 +79,10 @@ def __call__(self, inputs, role): if device is None: self.device = jax.devices()[0] else: - self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0] + self.device = device + if type(device) == str: + device_type, device_index = f"{device}:0".split(':')[:2] + self.device = jax.devices(device_type)[int(device_index)] self.observation_space = observation_space self.action_space = action_space diff --git a/skrl/multi_agents/jax/base.py b/skrl/multi_agents/jax/base.py index 77619b81..e177610d 100644 --- a/skrl/multi_agents/jax/base.py +++ b/skrl/multi_agents/jax/base.py @@ -60,7 +60,10 @@ def __init__(self, if device is None: self.device = jax.devices()[0] else: - self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0] + self.device = device + if type(device) == str: + device_type, device_index = f"{device}:0".split(':')[:2] + self.device = jax.devices(device_type)[int(device_index)] # convert the models to their respective device for _models in self.models.values(): diff --git a/skrl/resources/noises/jax/base.py b/skrl/resources/noises/jax/base.py index 2df54e20..e51e6fd3 100644 --- a/skrl/resources/noises/jax/base.py +++ b/skrl/resources/noises/jax/base.py @@ -31,7 +31,10 @@ def sample(self, size): if device is None: self.device = jax.devices()[0] else: - self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0] + self.device = device + if type(device) == str: + device_type, device_index = f"{device}:0".split(':')[:2] + self.device = jax.devices(device_type)[int(device_index)] def sample_like(self, tensor: Union[np.ndarray, jax.Array]) -> Union[np.ndarray, jax.Array]: """Sample a noise with the same size (shape) as the input tensor diff --git a/skrl/resources/preprocessors/jax/running_standard_scaler.py b/skrl/resources/preprocessors/jax/running_standard_scaler.py index 1d319605..ace65f06 100644 --- a/skrl/resources/preprocessors/jax/running_standard_scaler.py +++ b/skrl/resources/preprocessors/jax/running_standard_scaler.py @@ -95,7 +95,10 @@ def __init__(self, if device is None: self.device = jax.devices()[0] else: - self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0] + self.device = device + if type(device) == str: + device_type, device_index = f"{device}:0".split(':')[:2] + self.device = jax.devices(device_type)[int(device_index)] size = self._get_space_size(size)