111 lines
3.6 KiB
Python
111 lines
3.6 KiB
Python
import jax
|
|
import jax.numpy as jnp
|
|
|
|
from brittle_star_project import (
|
|
EnvConfig,
|
|
BrittleStarEnvFactory,
|
|
MorphologyConfig,
|
|
ArenaConfig,
|
|
Backend,
|
|
)
|
|
|
|
|
|
class BrittleStarJaxEnvWrapper:
|
|
def __init__(
|
|
self,
|
|
morphology: MorphologyConfig,
|
|
arena: ArenaConfig,
|
|
env_config: EnvConfig,
|
|
num_envs: int,
|
|
backend: Backend = Backend.MJX,
|
|
):
|
|
self._morphology = morphology
|
|
self._arena = arena
|
|
self._env_config = env_config
|
|
self._backend = backend
|
|
self._num_envs = num_envs
|
|
self._env = BrittleStarEnvFactory.create_environment(
|
|
self._backend, self._morphology, self._arena, self._env_config
|
|
)
|
|
|
|
from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks
|
|
|
|
self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm)
|
|
|
|
self._vectorized_reset = jax.jit(jax.vmap(self._env.reset))
|
|
self._vectorized_step = jax.jit(jax.vmap(self._env.step))
|
|
self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample))
|
|
|
|
self._action_rng = None
|
|
|
|
from experiment_logger import get_logger
|
|
|
|
self.logger = get_logger()
|
|
self.logger.info(
|
|
f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}"
|
|
)
|
|
|
|
@property
|
|
def backend(self):
|
|
return self._backend
|
|
|
|
@property
|
|
def raw(self):
|
|
return self._env
|
|
|
|
@property
|
|
def single_action_space(self):
|
|
return self._env.action_space
|
|
|
|
@property
|
|
def single_observation_space(self):
|
|
return self._env.observation_space
|
|
|
|
def reset(self, seed: int = 0):
|
|
self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
|
|
self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
|
|
env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
|
|
state = self._vectorized_reset(rng=env_rngs)
|
|
from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched
|
|
|
|
state = state.replace(
|
|
observations=pad_observations_batched(state.observations, self._padding_masks)
|
|
)
|
|
return state
|
|
|
|
def sample_actions(self):
|
|
assert self._action_rng is not None, "Call reset() before sample_actions()"
|
|
self._action_rng, *sub_rngs = jnp.array(
|
|
jax.random.split(self._action_rng, self._num_envs + 1)
|
|
)
|
|
return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
|
|
|
|
def step(self, state, action):
|
|
next_state = self._vectorized_step(state=state, action=action)
|
|
from brittle_star_project.environment.padded_obs_wrapper import pad_observations_batched
|
|
|
|
next_state = next_state.replace(
|
|
observations=pad_observations_batched(next_state.observations, self._padding_masks)
|
|
)
|
|
return next_state
|
|
|
|
def close(self):
|
|
self._env.close()
|
|
|
|
@staticmethod
|
|
def default(num_envs: int, backend: Backend = Backend.MJX) -> "BrittleStarJaxEnvWrapper":
|
|
morphology = MorphologyConfig()
|
|
arena = ArenaConfig()
|
|
env_config = EnvConfig()
|
|
return BrittleStarJaxEnvWrapper(
|
|
morphology, arena, env_config, num_envs=num_envs, backend=backend
|
|
)
|
|
|
|
def __str__(self):
|
|
morphology_str = str(self._morphology)
|
|
arena_str = str(self._arena)
|
|
env_config_str = str(self._env_config)
|
|
return (
|
|
f"BrittleStarJaxEnvWrapper(backend={self._backend}, num_envs={self._num_envs}, "
|
|
+ f"morphology={morphology_str}, arena={arena_str}, env_config={env_config_str})"
|
|
)
|