feat: Timestep now supports indexing

helx/_version.py

@@ -1,4 +1,4 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-__version__ = "1.1.4"
+__version__ = "1.1.5"
 __version_info__ = tuple(int(i) for i in __version__.split(".") if i.isdigit())

helx/base/mdp.py

@@ -19,6 +19,7 @@
 
 from jax import Array
 import jax.numpy as jnp
+import jax.tree_util as jtu
 from flax import struct
 
 
@@ -42,3 +43,9 @@ class Timestep(struct.PyTreeNode):
     """The true state of the MDP, $s_t$ before taking action `action`"""
     info: Dict[str, Any] = struct.field(default_factory=dict)
     """Additional information about the environment. Useful for accumulations (e.g. returns)"""
+
+    def __getitem__(self, key: Any) -> Timestep:
+        return jtu.tree_map(lambda x: x[key], self)
+
+    def __setitem__(self, key: Any, value: Any) -> Timestep:
+        return jtu.tree_map(lambda x: x.at[key].set(value), self)