rlmesh.jax

The JAX backend (experimental): the same clients, models, and sandbox sessions as the top-level package, with tensor leaves decoded to JAX arrays. JAX arrays are immutable, so there is no mutation hazard.

pip install "rlmesh[jax]"

Conversion semantics

• asarray() imports over DLPack. XLA shares RLMesh’s 64-byte-aligned buffers zero-copy and copies otherwise.
• from_array() moves the array to CPU if needed, blocks until ready, and copies it into a fresh RLMesh tensor.
• int64, uint64, and float64 need JAX 64-bit mode (jax.config.update("jax_enable_x64", True)); without it JAX demotes them.
• Requires jax >= 0.4.24, the first release with DLPack bool support.

Classes

• EnvFactoryExperimentalJAX-backed EnvFactory: served envs speak JAX arrays.
• ModelExperimentalExperimental JAX-backed model: predict works in JAX values.
• RemoteEnvExperimentalExperimental JAX-backed remote client for one environment.
• RemoteModelExperimentalExperimental JAX-backed handle to a served model (policy).
• RemoteVectorEnvExperimentalExperimental JAX-backed remote client for vectorized environments.
• SandboxBuildExperimentalBuild-from-source infrastructure for a sandbox image.
• SandboxEnvExperimentalExperimental JAX-backed owned sandbox session for one environment.
• SandboxInfoExperimentalInformation about a running RLMesh sandbox container.
• SandboxModelExperimentalA model served from an isolated container.
• SandboxRuntimeExperimentalContainer run-time settings for a sandbox -- docker run flags.
• SandboxVectorEnvExperimentalExperimental JAX-backed owned sandbox session for vectorized environments.

Functions

• asarray()ExperimentalReturn an immutable JAX array for an RLMesh tensor.
• ensure_available()ExperimentalRaise if JAX is not installed or is older than the supported floor.
• from_array()ExperimentalEncode a JAX array as an RLMesh value.
• space_from_spec()ExperimentalCreate a JAX-adapted space wrapper for a native space spec.

Type aliases

• JaxValueExperimentalJaxValue = PrimitiveValue | JaxArray | list['JaxValue'] | tuple['JaxValue', ...] | dict[str, 'JaxValue']