jax-envelope 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
envelope/__init__.py CHANGED
@@ -1,4 +1,4 @@
1
- from envelope.compat import create
1
+ from envelope.adapters import create
2
2
  from envelope.environment import Environment, Info, InfoContainer
3
3
  from envelope.spaces import BatchedSpace, Continuous, Discrete, PyTreeSpace, Space
4
4
  from envelope.struct import Container, FrozenPyTreeNode, field, static_field
@@ -4,14 +4,14 @@ from typing import Any, Protocol, Self
4
4
 
5
5
  # Lazy imports to avoid requiring all dependencies at once
6
6
  _env_module_map = {
7
- "gymnax": ("envelope.compat.gymnax_envelope", "GymnaxEnvelope"),
8
- "brax": ("envelope.compat.brax_envelope", "BraxEnvelope"),
9
- "navix": ("envelope.compat.navix_envelope", "NavixEnvelope"),
10
- "jumanji": ("envelope.compat.jumanji_envelope", "JumanjiEnvelope"),
11
- "kinetix": ("envelope.compat.kinetix_envelope", "KinetixEnvelope"),
12
- "craftax": ("envelope.compat.craftax_envelope", "CraftaxEnvelope"),
7
+ "gymnax": ("envelope.adapters.gymnax_envelope", "GymnaxEnvelope"),
8
+ "brax": ("envelope.adapters.brax_envelope", "BraxEnvelope"),
9
+ "navix": ("envelope.adapters.navix_envelope", "NavixEnvelope"),
10
+ "jumanji": ("envelope.adapters.jumanji_envelope", "JumanjiEnvelope"),
11
+ "kinetix": ("envelope.adapters.kinetix_envelope", "KinetixEnvelope"),
12
+ "craftax": ("envelope.adapters.craftax_envelope", "CraftaxEnvelope"),
13
13
  "mujoco_playground": (
14
- "envelope.compat.mujoco_playground_envelope",
14
+ "envelope.adapters.mujoco_playground_envelope",
15
15
  "MujocoPlaygroundEnvelope",
16
16
  ),
17
17
  }
@@ -69,7 +69,7 @@ class BraxEnvelope(Environment):
69
69
  info = InfoContainer(
70
70
  obs=brax_state.obs,
71
71
  reward=brax_state.reward,
72
- terminated=jnp.asarry(brax_state.done, dtype=bool).item(),
72
+ terminated=jnp.asarray(brax_state.done, dtype=bool),
73
73
  )
74
74
  info = info.update(**dataclasses.asdict(brax_state))
75
75
  return brax_state, info
@@ -10,7 +10,7 @@ from craftax.craftax_classic.envs.craftax_state import (
10
10
  from craftax.craftax_env import make_craftax_env_from_name
11
11
 
12
12
  from envelope import spaces as envelope_spaces
13
- from envelope.compat.gymnax_envelope import _convert_space as _convert_gymnax_space
13
+ from envelope.adapters.gymnax_envelope import _convert_space as _convert_gymnax_space
14
14
  from envelope.environment import Environment, Info, InfoContainer, State
15
15
  from envelope.struct import Container, static_field
16
16
  from envelope.typing import Key, PyTree, TypeAlias
@@ -81,7 +81,7 @@ class JumanjiEnvelope(Environment):
81
81
 
82
82
 
83
83
  def convert_jumanji_to_envelope_info(timestep: JumanjiTimeStep) -> InfoContainer:
84
- term = jnp.asarray(timestep.last(), dtype=bool).item()
84
+ term = jnp.asarray(timestep.last(), dtype=bool)
85
85
  info = InfoContainer(
86
86
  obs=timestep.observation, reward=timestep.reward, terminated=term
87
87
  ).update(**timestep.extras)
@@ -1,7 +1,7 @@
1
1
  """Kinetix compatibility wrapper.
2
2
 
3
3
  This module exposes Kinetix environments through the `envelope.environment.Environment`
4
- API. It mirrors envelope's compat philosophy:
4
+ API. It mirrors envelope's adapters philosophy:
5
5
  - prefer *no* environment-side auto-reset (use `AutoResetWrapper` in envelope)
6
6
  - prefer *no* fixed episode time-limits (use `TruncationWrapper` in envelope)
7
7
 
@@ -30,7 +30,7 @@ from kinetix.util.saving import load_from_json_file
30
30
 
31
31
  from envelope import field
32
32
  from envelope import spaces as envelope_spaces
33
- from envelope.compat.gymnax_envelope import _convert_space as _convert_gymnax_space
33
+ from envelope.adapters.gymnax_envelope import _convert_space as _convert_gymnax_space
34
34
  from envelope.environment import Environment, Info, InfoContainer, State
35
35
  from envelope.struct import Container, static_field
36
36
  from envelope.typing import Key, PyTree
envelope/environment.py CHANGED
@@ -42,7 +42,7 @@ class Environment(ABC, FrozenPyTreeNode):
42
42
 
43
43
  Two distinct lifecycle methods:
44
44
  init(key) - Initialize environment and all state from scratch.
45
- reset(key, state) - Reset the inner environment while preserving
45
+ reset(state, key) - Reset the inner environment while preserving
46
46
  episode-persistent state.
47
47
  """
48
48
 
@@ -51,7 +51,7 @@ class Environment(ABC, FrozenPyTreeNode):
51
51
  """Initialize environment and all state from scratch."""
52
52
  ...
53
53
 
54
- def reset(self, key: Key, state: State) -> tuple[State, Info]:
54
+ def reset(self, state: State, key: Key) -> tuple[State, Info]:
55
55
  """Reset the inner environment while preserving episode-persistent state."""
56
56
  return self.init(key)
57
57
 
envelope/spaces.py CHANGED
@@ -1,6 +1,6 @@
1
1
  from abc import ABC, abstractmethod
2
2
  from functools import cached_property
3
- from typing import cast, override
3
+ from typing import override
4
4
 
5
5
  import jax
6
6
  from jax import numpy as jnp
@@ -54,7 +54,7 @@ class AutoResetWrapper(Wrapper):
54
54
  return state, info.update(final=state.last_final)
55
55
 
56
56
  @override
57
- def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
57
+ def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
58
58
  raise NotImplementedError("Reset is not implemented for AutoResetWrapper")
59
59
 
60
60
  @override
@@ -63,7 +63,7 @@ class AutoResetWrapper(Wrapper):
63
63
  state = state.replace(reset_key=key)
64
64
 
65
65
  inner_state, info = self.env.step(state.inner_state, action)
66
- reset_inner_state, reset_info = self.env.reset(key_reset, inner_state)
66
+ reset_inner_state, reset_info = self.env.reset(inner_state, key_reset)
67
67
 
68
68
  # Select next state and info based on done
69
69
  done = info.terminated | info.truncated
@@ -35,8 +35,8 @@ class ContinuousObservationWrapper(Wrapper):
35
35
  return state, info
36
36
 
37
37
  @override
38
- def reset(self, key: Key, state: State) -> tuple[State, Info]:
39
- state, info = self.env.reset(key, state)
38
+ def reset(self, state: State, key: Key) -> tuple[State, Info]:
39
+ state, info = self.env.reset(state, key)
40
40
  info = info.update(obs=to_float(info.obs))
41
41
  return state, info
42
42
 
@@ -24,8 +24,8 @@ class EpisodeStatisticsWrapper(Wrapper):
24
24
  return state, info.update(stats=state.stats)
25
25
 
26
26
  @override
27
- def reset(self, key: Key, state: State) -> tuple[State, Info]:
28
- inner_state, info = self.env.reset(key, state.inner_state)
27
+ def reset(self, state: State, key: Key) -> tuple[State, Info]:
28
+ inner_state, info = self.env.reset(state.inner_state, key)
29
29
  state = state.replace(inner_state=inner_state)
30
30
  return state, info.update(stats=state.stats)
31
31
 
@@ -37,10 +37,9 @@ class FlattenObservationWrapper(Wrapper):
37
37
  return state, info
38
38
 
39
39
  @override
40
- def reset(self, key: Key, state: State) -> tuple[State, Info]:
41
- state, info = self.env.reset(key, state)
42
- info = info.update(obs=flatten_x(info.obs))
43
- return state, info
40
+ def reset(self, state: State, key: Key) -> tuple[State, Info]:
41
+ next_state, info = self.env.reset(state, key)
42
+ return next_state, info.update(obs=flatten_x(info.obs))
44
43
 
45
44
  @override
46
45
  def step(self, state: State, action: PyTree) -> tuple[State, Info]:
@@ -76,8 +76,8 @@ class ObservationNormalizationWrapper(Wrapper):
76
76
  return self._normalize_and_update(next_state, info)
77
77
 
78
78
  @override
79
- def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
80
- inner_state, info = self.env.reset(key, state.inner_state)
79
+ def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
80
+ inner_state, info = self.env.reset(state.inner_state, key)
81
81
  # Preserve running statistics across resets
82
82
  next_state = self.ObservationNormalizationState(
83
83
  inner_state=inner_state, rmv_state=state.rmv_state
@@ -36,7 +36,7 @@ class PooledInitVmapWrapper(Wrapper):
36
36
  return state, info.update(final=pholder_info)
37
37
 
38
38
  @override
39
- def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
39
+ def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
40
40
  # It's hard to support reset for this wrapper.
41
41
  # We would have to init the state of a pool of unwrapped environments, and then
42
42
  # somehow inject this into the stack of wrapped states. The current data
@@ -48,7 +48,7 @@ class PooledInitVmapWrapper(Wrapper):
48
48
  # episodes before vmapping, we will implement this later.
49
49
  keys = _split_or_keep_key(key, self.batch_size + 1)
50
50
  key_next, keys_pool = keys[0], keys[1:]
51
- inner_state, info = jax.vmap(self.env.reset)(keys_pool, state.inner_state)
51
+ inner_state, info = jax.vmap(self.env.reset)(state.inner_state, keys_pool)
52
52
  state = state.replace(inner_state=inner_state, init_key=key_next)
53
53
  return state, info.update(final=state.last_final)
54
54
 
@@ -69,13 +69,13 @@ class StateInjectionWrapper(Wrapper):
69
69
  return state, info
70
70
 
71
71
  @override
72
- def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
72
+ def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
73
73
  # If reset state is set, use it instead of resetting inner env
74
74
  if state.reset_state is not None and state.reset_obs is not None:
75
75
  inner_state = state.reset_state
76
76
  info = InfoContainer(obs=state.reset_obs, reward=0.0, terminated=False)
77
77
  elif state.reset_state is None and state.reset_obs is None:
78
- inner_state, info = self.env.reset(key, state.inner_state)
78
+ inner_state, info = self.env.reset(state.inner_state, key)
79
79
  else:
80
80
  raise ValueError("State must set both reset_state and reset_obs or neither")
81
81
 
@@ -21,8 +21,8 @@ class TruncationWrapper(Wrapper):
21
21
  return state, info.update(truncated=self.max_steps <= 0)
22
22
 
23
23
  @override
24
- def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
25
- inner_state, info = self.env.reset(key, state.inner_state)
24
+ def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
25
+ inner_state, info = self.env.reset(state.inner_state, key)
26
26
  state = state.replace(inner_state=inner_state, steps=0)
27
27
  return state, info.update(truncated=self.max_steps <= 0)
28
28
 
@@ -40,11 +40,9 @@ class VmapEnvsWrapper(Wrapper):
40
40
  return state, info
41
41
 
42
42
  @override
43
- def reset(self, key: Key, state: PyTree) -> tuple[WrappedState, Info]:
43
+ def reset(self, state: PyTree, key: Key) -> tuple[WrappedState, Info]:
44
44
  keys = self._split_keys(key)
45
- state, info = jax.vmap(lambda e, k, s: e.reset(k, s))(
46
- self.env, keys, state
47
- )
45
+ state, info = jax.vmap(lambda e, s, k: e.reset(s, k))(self.env, state, keys)
48
46
  return state, info
49
47
 
50
48
  @override
@@ -58,7 +56,9 @@ class VmapEnvsWrapper(Wrapper):
58
56
  @property
59
57
  def observation_space(self) -> spaces.Space:
60
58
  env0 = _index_env(self.env, 0, self.batch_size)
61
- return spaces.BatchedSpace(space=env0.observation_space, batch_size=self.batch_size)
59
+ return spaces.BatchedSpace(
60
+ space=env0.observation_space, batch_size=self.batch_size
61
+ )
62
62
 
63
63
  @override
64
64
  @cached_property
@@ -41,9 +41,9 @@ class VmapWrapper(Wrapper):
41
41
  return state, info
42
42
 
43
43
  @override
44
- def reset(self, key: Key, state: PyTree) -> tuple[WrappedState, Info]:
44
+ def reset(self, state: PyTree, key: Key) -> tuple[WrappedState, Info]:
45
45
  keys = _split_or_keep_key(key, self.batch_size)
46
- state, info = jax.vmap(self.env.reset)(keys, state)
46
+ state, info = jax.vmap(self.env.reset)(state, keys)
47
47
  return state, info
48
48
 
49
49
  @override
@@ -29,11 +29,11 @@ class Wrapper(Environment):
29
29
  return self.env.init(key)
30
30
 
31
31
  @override
32
- def reset(self, key: Key, state: State) -> tuple[State, Info]:
33
- return self.env.reset(key, state)
32
+ def reset(self, state: State, key: Key) -> tuple[State, Info]:
33
+ return self.env.reset(state, key)
34
34
 
35
35
  @override
36
- def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
36
+ def step(self, state: State, action: PyTree) -> tuple[State, Info]:
37
37
  return self.env.step(state, action)
38
38
 
39
39
  @override
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: jax-envelope
3
- Version: 0.2.0
3
+ Version: 0.4.0
4
4
  Summary: A JAX-native environment interface with powerful wrappers and adapters for popular RL environment suites
5
5
  Project-URL: Homepage, https://github.com/keraJLi/envelope
6
6
  Project-URL: Repository, https://github.com/keraJLi/envelope
@@ -25,12 +25,13 @@ Requires-Dist: jax>=0.5.0
25
25
  Description-Content-Type: text/markdown
26
26
 
27
27
  # 💌 Envelope: a JAX-native environment interface
28
+
28
29
  ```python
29
30
  # Create environments from JAX-native suites you have installed, ...
30
31
  env = envelope.create("gymnax::CartPole-v1")
31
32
 
32
33
  # ... interact with the environments using a simple interface, ...
33
- state, info = env.reset(key)
34
+ state, info = env.init(key)
34
35
  states, infos = jax.lax.scan(env.step, state, actions)
35
36
  plt.plot(infos.reward.cumsum())
36
37
 
@@ -41,47 +42,56 @@ env = envelope.wrappers.ObservationNormalizationWrapper(env)
41
42
  ```
42
43
 
43
44
  ## 🌍 Simple, expressive interaction!
44
- * **Environments are pytrees**. Squish them through JAX transformations and trace their parameters.
45
- * **Idiomatic jax-y interface** of `reset(key: Key) -> State, Info` and `step(state: State, action: PyTree) -> State, Info`. You can directly `jax.scan` over a `step(...)`!
46
- * **Spaces are super simple**. No `Tuple`, `Dict` nonsense! There are two spaces: `Continuous` and `Discrete`, which you can compose into a `PyTreeSpace`.
47
- * **Explicit episode truncation** supports correctly handling bootstrapping for value-function targets.
48
- * **No auto-reset** by default. Resetting every step can be expensive!
45
+
46
+ - **Environments are pytrees**. Squish them through JAX transformations and trace their parameters.
47
+ - **Idiomatic jax-y interface** of `init(key: Key) -> State, Info` and `step(state: State, action: PyTree) -> State, Info`. You can directly `jax.scan` over a `step(...)`!
48
+ - **Spaces are super simple**. No `Tuple`, `Dict` nonsense! There are two spaces: `Continuous` and `Discrete`, which you can compose into a `PyTreeSpace`.
49
+ - **Explicit episode truncation** supports correctly handling bootstrapping for value-function targets.
50
+ - **No auto-reset** by default. Resetting every step can be expensive!
49
51
 
50
52
  ## 💪 Powerful, composable wrappers!
51
- * **Carry state across episodes** to track running statistics, for example to normalize observations.
52
- * **Composable wrappers** can be stacked in any order. For example, `ObservationNormalizationWrapper` before vs. after `VmapWrapper` gives per-env vs. global normalization.
53
- <!-- TODO: Add auto-reset behavior (including state injection) and optimistic resets once I implement them. -->
53
+
54
+ - **Carry state across episodes** to track running statistics, for example to normalize observations.
55
+ - **Composable wrappers** can be stacked in any order. For example, `ObservationNormalizationWrapper` before vs. after `VmapWrapper` gives per-env vs. global normalization.
54
56
 
55
57
  ## 🔌 Adapters for existing suites
56
- | 📦 | # 🤖 | # 🌍 |
57
- |------|------|------|
58
- | [gymnax](https://github.com/RobertTLange/gymnax) | 🕺 | 24 |
59
- | [brax](https://github.com/google/brax) | 🕺 | 12 |
60
- | [jumanji](https://github.com/instadeepai/jumanji) | 🕺 / 👯 | 25 / 1 |
61
- | [kinetix](https://github.com/flairox/kinetix) | 🕺 | 74 |
62
- | [craftax](https://github.com/MichaelTMatthews/craftax) | 🕺 | 4 |
63
- | [mujoco_playground](https://github.com/google-deepmind/mujoco_playground) | 🕺 | 54 |
64
- | | |
65
- | Total | 🕺 / 👯 | 193 / 1 |
58
+
59
+
60
+ | 📦 | # 🤖 | # 🌍 |
61
+ | ------------------------------------------------------------------------- | ------- | ------- |
62
+ | [gymnax](https://github.com/RobertTLange/gymnax) | 🕺 | 24 |
63
+ | [brax](https://github.com/google/brax) | 🕺 | 12 |
64
+ | [jumanji](https://github.com/instadeepai/jumanji) | 🕺 / 👯 | 25 / 1 |
65
+ | [kinetix](https://github.com/flairox/kinetix) | 🕺 | 74 |
66
+ | [craftax](https://github.com/MichaelTMatthews/craftax) | 🕺 | 4 |
67
+ | [mujoco_playground](https://github.com/google-deepmind/mujoco_playground) | 🕺 | 54 |
68
+ | | | |
69
+ | Total | 🕺 / 👯 | 193 / 1 |
70
+
66
71
 
67
72
  ```python
68
73
  envelope.create("📦::🌍")
69
74
  ```
75
+
70
76
  let's you create environments from any of the above!
71
77
 
72
78
  ## 📝 Testing
73
- - **Default (no optional compat deps required)**: `uv run pytest -m "not compat"`
74
- - **Compat suite (requires full compat dependency group)**:
75
- - `uv sync --group compat`
76
- - `uv run pytest -m compat`
77
- - If any compat dependency is missing/broken, the run will fail fast with an error telling you what to install.
79
+
80
+ - **Default (no optional adapters deps required)**: `uv run pytest -m "not adapters"`
81
+ - **Adapters suite (requires full adapters dependency group)**:
82
+ - `uv sync --group adapters`
83
+ - `uv run pytest -m adapters`
84
+ - If any adapter dependency is missing/broken, the run will fail fast with an error telling you what to install.
78
85
 
79
86
  ## 🏗️ Installation
87
+
80
88
  ```bash
81
89
  pip install jax-envelope
82
90
  ```
83
91
 
84
92
  ## 💞 Related projects
85
- * [stoa](https://github.com/EdanToledo/Stoa) is a very similar project that provides adapters and wrappers for the jumanji-like interface.
86
- * Check out all the great suites we have adapters for! [gymnax](https://github.com/RobertTLange/gymnax), [brax](https://github.com/google/brax), [jumanji](https://github.com/instadeepai/jumanji), [kinetix](https://github.com/flairox/kinetix), [craftax](https://github.com/MichaelTMatthews/craftax), [mujoco_playground](https://github.com/google-deepmind/mujoco_playground).
87
- * We will be adding support for [jaxmarl](https://github.com/flairox/jaxmarl) and [pgx](https://github.com/sotetsuk/pgx) in the future, as soon as we figured out the best ever MARL interface for JAX!
93
+
94
+ - [stoa](https://github.com/EdanToledo/Stoa) is a very similar project that provides adapters and wrappers for the jumanji-like interface.
95
+ - Check out all the great suites we have adapters for! [gymnax](https://github.com/RobertTLange/gymnax), [brax](https://github.com/google/brax), [jumanji](https://github.com/instadeepai/jumanji), [kinetix](https://github.com/flairox/kinetix), [craftax](https://github.com/MichaelTMatthews/craftax), [mujoco_playground](https://github.com/google-deepmind/mujoco_playground).
96
+ - We will be adding support for [jaxmarl](https://github.com/flairox/jaxmarl) and [pgx](https://github.com/sotetsuk/pgx) in the future, as soon as we figured out the best ever MARL interface for JAX!
97
+
@@ -0,0 +1,32 @@
1
+ envelope/__init__.py,sha256=3IaAPQ12MPozGX1a73c6_EGaI26-UdXYtu5rqJfPBYs,1331
2
+ envelope/environment.py,sha256=ZPTznHjMdpUkadDB0rpiGuG_E5ZhfgHt8SxfUhDgySU,2015
3
+ envelope/spaces.py,sha256=E05OsM999nJw03lyg-ALXbvuw8xQT4AJhLFyahDKFaA,6938
4
+ envelope/struct.py,sha256=tO1LLxk0bYc4kBCXraDv93T5d-xFRzpv4QCjASB0Lv4,5522
5
+ envelope/typing.py,sha256=2rWKiZnKXaUU6JIl51Wcj_txRwe8wQNRSLSiRuT3OP4,127
6
+ envelope/adapters/__init__.py,sha256=Xn2jNxQK1tdS1P1FEzCNTO8ulLuKfDaiozWkOqNRTPQ,3472
7
+ envelope/adapters/brax_envelope.py,sha256=xToV4fJYU9gPfF5SDs1Y3BaP96daz9iW8gKs2QYQDNA,3476
8
+ envelope/adapters/craftax_envelope.py,sha256=G9jLBJjV_YQh_dnnQlk0hm_hW5B2lyjz7GiRZnIo-FM,3961
9
+ envelope/adapters/gymnax_envelope.py,sha256=X0pfFKaFKzvBl2Zih4utiXSEI-lc5f9miJGMxeT-Xfw,4690
10
+ envelope/adapters/jumanji_envelope.py,sha256=zl6yE1yqdCxZmjVCexu5vD10f5a-M82EpjzyVlj13wM,4909
11
+ envelope/adapters/kinetix_envelope.py,sha256=Nh98HXzuXB_gCWKOdYWBiaxa2v1zyhmYzrUw2xmfwWo,6857
12
+ envelope/adapters/mujoco_playground_envelope.py,sha256=7iPaRuwXUHbP8hHNgSFipW8rJVruootMQPYBWNpPiHA,3705
13
+ envelope/adapters/navix_envelope.py,sha256=2n8vgr6gVeAYo12D1kDPxUrA9Hm-If3E7rV7TF_ksSA,3026
14
+ envelope/wrappers/__init__.py,sha256=GFs54kY1aWbf9xymzm22x6M2AhZCwFLx1f5LF23jX94,1409
15
+ envelope/wrappers/autoreset_wrapper.py,sha256=MIP9Svl1KyX0B7tfsSpcDJeEzPrAJQNWwpixrMZkF_s,3382
16
+ envelope/wrappers/clip_action_wrapper.py,sha256=YKJbKBUbf-LADxIfCE6mMOKo_NEa2xPfDncNKg9iEHM,955
17
+ envelope/wrappers/continuous_observation_wrapper.py,sha256=ud6MgMe7hY91q_26-A9t24RJSPrHI8wXue1LTdrm8O8,2041
18
+ envelope/wrappers/episode_statistics_wrapper.py,sha256=F5NM6U0Z5kImSr7uflPb6nnZbtDtyqewAlryzqiWnwU,1433
19
+ envelope/wrappers/flatten_action_wrapper.py,sha256=znR7WSebLy2wZnEWHi9SshxgV9b29khg-6ogQlmXq9U,2609
20
+ envelope/wrappers/flatten_observation_wrapper.py,sha256=-i7mH09MoTwlA9xT7768n2Pb0TrVYSTLiUdoPKR76Gc,2829
21
+ envelope/wrappers/normalization.py,sha256=YkNEZqvFSaPurBA-WfXDynV51552ucNkY0JqNSAKNUE,1819
22
+ envelope/wrappers/observation_normalization_wrapper.py,sha256=crCm4FREqUumFnNdr30-cORLp4rIOTiyKywMzvQz0Xk,4921
23
+ envelope/wrappers/pooled_init_vmap_wrapper.py,sha256=foqZBx-7u-n8VMidF_SoGFY0-Y-RFKB9v-Nj7ZTtmqk,4985
24
+ envelope/wrappers/state_injection_wrapper.py,sha256=2Lr_LG2-0EF1q2KP7yQ954xfrWTU9moSC4bB1VLK3Vs,3553
25
+ envelope/wrappers/truncation_wrapper.py,sha256=k3NZBMmnYiCN35jA1Yyi3rGdLcqM52CzuPabtvwGxe8,1345
26
+ envelope/wrappers/vmap_envs_wrapper.py,sha256=b-xtzhVcT7P0IV4dvVvQk-rX5zv9H4Msv3Ftf5KnZp4,2757
27
+ envelope/wrappers/vmap_wrapper.py,sha256=8Qxml26HZqjdKV7ey43wLszjMUxURQUFBHwJ1z5F9y8,1973
28
+ envelope/wrappers/wrapper.py,sha256=SSnKFs_9Rhf0zorMUTTSlzj81OwR7KWo4PgCgycROwo,1503
29
+ jax_envelope-0.4.0.dist-info/METADATA,sha256=gxqlnwoNe2wcVPfYMZQE9HRROD8q7CDkni1ND1akrKg,5056
30
+ jax_envelope-0.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
31
+ jax_envelope-0.4.0.dist-info/licenses/LICENSE,sha256=VyF-MK-gY2_fZlhf8uEnE2y8ziIXK-w55GM12eOgXrQ,1069
32
+ jax_envelope-0.4.0.dist-info/RECORD,,
@@ -1,32 +0,0 @@
1
- envelope/__init__.py,sha256=AAk0b9UHGUTC4EzERXYPOrHuECgO0fx_UN4_gpKvgT0,1329
2
- envelope/environment.py,sha256=OC15ooivLvD8EpmgRQGL-UjRpF0UXO0fOGu61I8HmRs,2015
3
- envelope/spaces.py,sha256=Nb8T11ilCtqKPSLZ98oqhnuw_8JNLOGCZdsy2X_pSwU,6944
4
- envelope/struct.py,sha256=tO1LLxk0bYc4kBCXraDv93T5d-xFRzpv4QCjASB0Lv4,5522
5
- envelope/typing.py,sha256=2rWKiZnKXaUU6JIl51Wcj_txRwe8wQNRSLSiRuT3OP4,127
6
- envelope/compat/__init__.py,sha256=6q2B2bfTIu7MlnIXc702ysqEnyURz-MEjTBMBCo8rdQ,3458
7
- envelope/compat/brax_envelope.py,sha256=XV07CYApXmf5g-Bin4pKWqExTZ7gd6AO2BCx3bk1Z1c,3482
8
- envelope/compat/craftax_envelope.py,sha256=ey_mhJh5KCDRvjpCAJbFexpWiFjzQr7BZuRYLM9g-yA,3959
9
- envelope/compat/gymnax_envelope.py,sha256=X0pfFKaFKzvBl2Zih4utiXSEI-lc5f9miJGMxeT-Xfw,4690
10
- envelope/compat/jumanji_envelope.py,sha256=WhU4WJLW2Fn4BG1_69EJO9ZiBDWEaY387dhkRZjqcGc,4916
11
- envelope/compat/kinetix_envelope.py,sha256=99RNN6nBN9CQEEe2fX2T1n8lL9ki0uKaz5BK052IU1o,6853
12
- envelope/compat/mujoco_playground_envelope.py,sha256=7iPaRuwXUHbP8hHNgSFipW8rJVruootMQPYBWNpPiHA,3705
13
- envelope/compat/navix_envelope.py,sha256=2n8vgr6gVeAYo12D1kDPxUrA9Hm-If3E7rV7TF_ksSA,3026
14
- envelope/wrappers/__init__.py,sha256=GFs54kY1aWbf9xymzm22x6M2AhZCwFLx1f5LF23jX94,1409
15
- envelope/wrappers/autoreset_wrapper.py,sha256=MJlbOj94_6I1KQ4oBNsKEB416YWp9ZMiqn85b1PdhKQ,3382
16
- envelope/wrappers/clip_action_wrapper.py,sha256=YKJbKBUbf-LADxIfCE6mMOKo_NEa2xPfDncNKg9iEHM,955
17
- envelope/wrappers/continuous_observation_wrapper.py,sha256=geDOeQDuASiRDEv5U3fF7GxblfB0Ku99GvTUGZSteJs,2041
18
- envelope/wrappers/episode_statistics_wrapper.py,sha256=-LzypX4aYS9mCkl33M2mf3aT3OhXWtDc_kVU6VlSJ0o,1433
19
- envelope/wrappers/flatten_action_wrapper.py,sha256=znR7WSebLy2wZnEWHi9SshxgV9b29khg-6ogQlmXq9U,2609
20
- envelope/wrappers/flatten_observation_wrapper.py,sha256=xgoToXdJ_6M7cwfL7pu91ZnnFh3LHPR_DSe7XJwkh2Q,2839
21
- envelope/wrappers/normalization.py,sha256=YkNEZqvFSaPurBA-WfXDynV51552ucNkY0JqNSAKNUE,1819
22
- envelope/wrappers/observation_normalization_wrapper.py,sha256=pAwyd01lAxvEu1Z4KZfhSRoVNeLpeohtw1t1wsE53w4,4921
23
- envelope/wrappers/pooled_init_vmap_wrapper.py,sha256=LAgrEjLXlq-gJiCGzI357a41tx1mVgb1ZQwtgN8mTuw,4985
24
- envelope/wrappers/state_injection_wrapper.py,sha256=L8f9yrlSigk0jNx6D6wSu0cR0bgssq2uJNd9OYPAY0E,3553
25
- envelope/wrappers/truncation_wrapper.py,sha256=-CSoXs-mDy1AiJVKLUc2zBHQxalhYjjyAAu1NmiIUsE,1345
26
- envelope/wrappers/vmap_envs_wrapper.py,sha256=gE_TAmfsu2PJlDGWTz5xENJBJ-6nSJpTUlKhOnmZPi0,2757
27
- envelope/wrappers/vmap_wrapper.py,sha256=B43lFNdAN_d8LgN6Rz6Ty85r9zcXLB3Wrn8R-Z2qHps,1973
28
- envelope/wrappers/wrapper.py,sha256=tWluyIg4eGl60R2v3qaahU5dlmKKkBsJiNj7GnWR06E,1517
29
- jax_envelope-0.2.0.dist-info/METADATA,sha256=jCWoAuoFsrLJakIhecRdz1zOIGCdrD4dBS0JSA2LVvw,4651
30
- jax_envelope-0.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
31
- jax_envelope-0.2.0.dist-info/licenses/LICENSE,sha256=VyF-MK-gY2_fZlhf8uEnE2y8ziIXK-w55GM12eOgXrQ,1069
32
- jax_envelope-0.2.0.dist-info/RECORD,,
File without changes
File without changes