Skip to content

Vectorized Environments

JoltGym provides JoltVectorEnv, a high-performance vectorized environment that steps N independent physics worlds in parallel using native C++ threads.

Unlike Python-based vectorization (e.g., SubprocVecEnv), JoltVectorEnv wraps the C++ WorldPool class which:

  • Maintains N independent PhysicsSystem instances
  • Steps all environments in parallel via C++ ParallelFor
  • Releases the GIL during the entire hot loop
  • Returns batched NumPy arrays with zero-copy where possible

Quick Usage

import numpy as np
from joltgym.vector import JoltVectorEnv

envs = JoltVectorEnv(
    num_envs=64,
    model_path="python/joltgym/assets/half_cheetah.xml",
)

obs, infos = envs.reset(seed=42)

actions = np.random.uniform(-1, 1, (64, 6)).astype(np.float32)
obs, rewards, dones, truncs, infos = envs.step(actions)

Threading Architecture

Python thread (GIL released)
  +-- C++ WorldPool::StepAll()
       +-- ParallelFor across min(hardware_concurrency, 16) OS threads
            |-- thread 0: step envs [0, chunk)
            |-- thread 1: step envs [chunk, 2*chunk)
            +-- thread N: step envs [last_chunk, num_envs)

When to Use

Scenario Recommended Approach
Single environment joltgym.make(...)
2--8 parallel envs (SB3) SubprocVecEnv with joltgym.make
16+ parallel envs JoltVectorEnv (C++ WorldPool)
Maximum throughput JoltVectorEnv with 64--256 envs

API Reference

JoltVectorEnv(num_envs, model_path, **kwargs)

N parallel HalfCheetah environments stepped in C++ threads.

Wraps the C++ WorldPool class for maximum throughput. All N PhysicsSystem instances are stepped in parallel via native OS threads with the GIL released — the entire hot loop (action apply, physics step, observation extraction, reward computation) runs in C++.

Achieves ~73K env-steps/sec at 256 environments on Apple Silicon.

Attributes:

Name Type Description
num_envs

Number of parallel environments.
observation_space

Batched observation space (num_envs, obs_dim).
action_space

Batched action space (num_envs, act_dim).
single_observation_space

Single-env observation space (obs_dim,).
single_action_space

Single-env action space (act_dim,).

Initialize the vectorized environment pool.

Parameters:

Name Type Description Default
num_envs

Number of parallel environments to create.

 required
model_path

Path to the MJCF XML model file.

 required
**kwargs

Forwarded to WorldPool (e.g. forward_reward_weight, ctrl_cost_weight).

 {}

step(actions)

Step all environments in parallel.

The GIL is released for the entire duration of the C++ computation. Environments that reach a terminal state are auto-reset.

Parameters:

Name Type Description Default
actions

Array of shape (num_envs, act_dim), dtype float32.

 required

Returns:

Name Type Description
obs

Observations, shape (num_envs, obs_dim).
rewards

Rewards, shape (num_envs,).
dones

Terminal flags, shape (num_envs,). True indicates the environment was auto-reset.
truncs

Truncation flags, shape (num_envs,) (always False).
infos

List of empty dicts.

reset(*, seed=None, options=None)

Reset all environments in parallel.

Parameters:

Name Type Description Default
seed

Optional base seed. Environment i receives seed + i.

 None
options

Unused, present for compatibility.

 None

Returns:

Name Type Description
obs

Initial observations, shape (num_envs, obs_dim).
infos

List of empty dicts.

close()

Clean up resources (no-op, pool is managed by C++).