178 lines
6.9 KiB
Python
178 lines
6.9 KiB
Python
"""Adapter: wrap a Gym-style env so it implements DiscoRL's Environment API.
|
|
|
|
This is a minimal adapter intended for evaluation / inference (batching=1
|
|
or small batches). It converts Gym observations/rewards/dones into the
|
|
`types.EnvironmentTimestep` structure expected by DiscoRL and keeps a
|
|
Python-side list of env instances for the batch.
|
|
|
|
Notes:
|
|
- This adapter does not attempt to JIT or vectorize with JAX. It simply
|
|
converts numpy -> jax arrays before returning timesteps so the DiscoRL
|
|
agent (Haiku/JAX) can consume them.
|
|
- For training at scale you can rework this into a true batched env that
|
|
runs multiple envs in parallel / in subprocesses.
|
|
"""
|
|
|
|
from typing import Any, Tuple
|
|
|
|
import numpy as np
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
from disco_rl.environments import base
|
|
from disco_rl import types
|
|
try:
|
|
from dm_env import StepType
|
|
except ImportError:
|
|
# Fallback if dm_env not available
|
|
class StepType:
|
|
MID = 0
|
|
LAST = 1
|
|
|
|
|
|
class GymToDiscoEnv(base.Environment):
|
|
"""Wrap a Gym-compatible environment class.
|
|
|
|
The wrapped `gym_env_cls` must follow the Gym API (reset() -> obs, info,
|
|
step(action) -> obs, reward, terminated, truncated, info).
|
|
|
|
Args:
|
|
gym_env_cls: factory/class that creates Gym environment instances.
|
|
batch_size: number of parallel env instances to manage.
|
|
env_settings: dict of kwargs to pass to gym_env_cls.
|
|
discrete_actions: optional array of shape (num_actions, action_dim) mapping
|
|
discrete action indices to continuous action vectors. If provided, the
|
|
action_spec becomes discrete (int32 scalar indices) and step() will
|
|
map indices to continuous actions before sending to underlying env.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
gym_env_cls: Any,
|
|
batch_size: int = 1,
|
|
env_settings=None,
|
|
discrete_actions: np.ndarray | None = None,
|
|
):
|
|
self.batch_size = batch_size
|
|
env_settings = {} if env_settings is None else env_settings
|
|
# Create multiple env instances for simple batching.
|
|
self._envs = [gym_env_cls(**env_settings) for _ in range(batch_size)]
|
|
self._discrete_actions = (
|
|
np.asarray(discrete_actions) if discrete_actions is not None else None
|
|
)
|
|
|
|
# Build single action/observation specs in the simple form expected by
|
|
# DiscoRL (a mapping with key 'observation'). We keep dtype/shape simple.
|
|
obs_space = self._envs[0].observation_space
|
|
act_space = self._envs[0].action_space
|
|
|
|
# Use dm_env-like BoundedArray for actions if available, else a simple
|
|
# placeholder (the agent only queries shape/dtype in most places).
|
|
try:
|
|
from dm_env import specs as dm_specs
|
|
|
|
# If discrete_actions is provided, create a discrete action spec;
|
|
# otherwise use the original continuous spec.
|
|
if self._discrete_actions is not None:
|
|
num_actions = len(self._discrete_actions)
|
|
self._single_action_spec = dm_specs.BoundedArray(
|
|
shape=(), dtype=np.int32, minimum=0, maximum=num_actions - 1
|
|
)
|
|
else:
|
|
self._single_action_spec = dm_specs.BoundedArray(
|
|
act_space.shape, act_space.dtype, act_space.low, act_space.high
|
|
)
|
|
self._single_observation_spec = {
|
|
'observation': dm_specs.Array(shape=obs_space.shape, dtype=obs_space.dtype)
|
|
}
|
|
except Exception:
|
|
# Fallback to simple numpy-shape descriptors.
|
|
if self._discrete_actions is not None:
|
|
num_actions = len(self._discrete_actions)
|
|
self._single_action_spec = type('ActionSpec', (), {
|
|
'shape': (),
|
|
'dtype': np.int32,
|
|
'low': 0,
|
|
'high': num_actions - 1,
|
|
})
|
|
else:
|
|
self._single_action_spec = act_space
|
|
self._single_observation_spec = {'observation': obs_space}
|
|
|
|
# Keep last observations / states for each env
|
|
self._last_obs = [None] * batch_size
|
|
self._dones = [True] * batch_size
|
|
|
|
def single_action_spec(self):
|
|
return self._single_action_spec
|
|
|
|
def single_observation_spec(self):
|
|
return self._single_observation_spec
|
|
|
|
def _obs_to_timestep(self, obs, reward, done):
|
|
# Convert a single env's raw outputs into types.EnvironmentTimestep
|
|
# DiscoRL expects a mapping for observation (e.g. {'observation': ...}).
|
|
obs_map = {'observation': jnp.asarray(obs, dtype=jnp.float32)}
|
|
step_type = jnp.array(StepType.LAST if done else StepType.MID, dtype=jnp.int32)
|
|
return types.EnvironmentTimestep(observation=obs_map, step_type=step_type, reward=jnp.array(float(reward), dtype=jnp.float32))
|
|
|
|
def step(self, state_unused, actions) -> Tuple[Any, types.EnvironmentTimestep]:
|
|
# actions expected to be a batched array with shape (batch_size, ...)
|
|
# For simplicity we iterate over envs sequentially.
|
|
# Support actions provided as numpy/jax arrays.
|
|
actions = [np.array(a) for a in list(actions)] if hasattr(actions, '__iter__') else [np.array(actions)]
|
|
|
|
obs_batch = []
|
|
reward_batch = []
|
|
done_batch = []
|
|
|
|
for i, env in enumerate(self._envs):
|
|
act = actions[i] if i < len(actions) else actions[0]
|
|
# If discrete_actions is provided, map the action index to continuous action.
|
|
if self._discrete_actions is not None:
|
|
act = self._discrete_actions[int(act)]
|
|
# Convert to python scalar if necessary
|
|
obs, reward, terminated, truncated, info = env.step(act)
|
|
done = bool(terminated or truncated)
|
|
self._last_obs[i] = obs
|
|
self._dones[i] = done
|
|
obs_batch.append(jnp.asarray(obs, dtype=jnp.float32))
|
|
reward_batch.append(float(reward))
|
|
done_batch.append(done)
|
|
|
|
# Stack to produce batched structures. DiscoRL typically expects
|
|
# observations to be a mapping of arrays with leading batch dimension.
|
|
obs_map = {'observation': jnp.stack(obs_batch)}
|
|
rewards = jnp.asarray(reward_batch, dtype=jnp.float32)
|
|
is_terminal = jnp.asarray(done_batch)
|
|
# Use jnp.where so scalar StepType values broadcast to the array shape.
|
|
step_type = jnp.where(is_terminal, StepType.LAST, StepType.MID)
|
|
|
|
timestep = types.EnvironmentTimestep(observation=obs_map, step_type=step_type, reward=rewards)
|
|
return None, timestep
|
|
|
|
def reset(self, rng_key=None) -> Tuple[Any, types.EnvironmentTimestep]:
|
|
obs_batch = []
|
|
reward_batch = []
|
|
done_batch = []
|
|
for i, env in enumerate(self._envs):
|
|
# Gym reset returns (obs, info) in Gymnasium; support both
|
|
out = env.reset()
|
|
if isinstance(out, tuple) and len(out) >= 1:
|
|
obs = out[0]
|
|
else:
|
|
obs = out
|
|
self._last_obs[i] = obs
|
|
self._dones[i] = False
|
|
obs_batch.append(jnp.asarray(obs, dtype=jnp.float32))
|
|
reward_batch.append(0.0)
|
|
done_batch.append(False)
|
|
|
|
obs_map = {'observation': jnp.stack(obs_batch)}
|
|
rewards = jnp.asarray(reward_batch, dtype=jnp.float32)
|
|
is_terminal = jnp.asarray(done_batch)
|
|
# Use jnp.where so scalar StepType values broadcast to the array shape.
|
|
step_type = jnp.where(is_terminal, StepType.LAST, StepType.MID)
|
|
timestep = types.EnvironmentTimestep(observation=obs_map, step_type=step_type, reward=rewards)
|
|
return None, timestep
|