Skip to content

Struct

Overview

Similarly to flax.struct, the struct submodule implements datastructures that are automatically registered as JAX pytrees with jax.tree_util, allowing us to pass them to JAX transformations. The two datastructures defined fullfil two main purposes in envelope:

  • FrozenPyTreeNode creates a dataclass with fixed fields that you can mark as either static or dynamic using struct.field.
    For example, environments are instances of FrozenPyTreeNode, allowing them to have both static (structural, e.g. world size) and dynamic params (not affecting array shapes and control flow, e.g. wind speed).
  • Container creates a dataclass with a core set of fields that will always be present. Additionally, instances of Container can be updated to hold any additional fields at runtime (and even within traced methods).
    For example, the step function of an environment emits an Info object that is a Container, and holds observation, reward and terminated/truncated flags. Wrappers can add information to this, such as current episode statistics.

API Reference

envelope.struct.FrozenPyTreeNode

Frozen dataclass base that is a JAX pytree node. Fields can be declared as either dynamic (pytree nodes) or static (not pytree nodes) using the field and static_field helpers.

Usage 
class Foo(FrozenPyTreeNode):
    a: Any                      # pytree leaf
    b: int = static_field()     # static, not a leaf

x = Foo(a={"w": 1.0}, b=0)
y = x.replace(b=1)

Functions

replace(**changes)

Shorthand for dataclasses.replace(self, **changes).

envelope.struct.field(*, pytree_node=True, **kwargs)

Dataclass field helper. See typing.FrozenPyTreeNode for more details. Set pytree_node=False for static (non-transformed) fields.

envelope.struct.static_field(**kwargs)

Shorthand for field(pytree_node=False, ...).

envelope.struct.Container dataclass

This class implements a container for arbitrary PyTree-valued fields. While subclasses can define arbitrary core fields, instances of this class can be updated to hold any additional extras fields.

Usage example 
class Foo(Container):
    bar: int

foo = Foo(bar=1)
foo = foo.update(bar=2, baz=3.0)
print(foo)  # Foo(bar=2, baz=3.0)

Functions

update(**changes)

Update the container with new values. The changes overwrite fields in the container, both for core fields and extras.

Parameters:

Name Type Description Default
**changes dict[str, PyTree]

A dictionary of field names and values to update.

 {}

Returns:

Type Description
Self

A new instance of the container with the updated values.