Metadata-Version: 2.5
Name: forkscope
Version: 0.1.1
Summary: A local workbench for evaluating test-time planning rollouts.
Author: Forkscope contributors
License-Expression: MIT
License-File: LICENSE
Requires-Python: >=3.10
Requires-Dist: numpy>=1.26
Requires-Dist: pydantic>=2.8
Requires-Dist: typer>=0.12
Provides-Extra: dev
Requires-Dist: build>=1.2; extra == 'dev'
Requires-Dist: mypy>=1.11; extra == 'dev'
Requires-Dist: pytest>=8.3; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Provides-Extra: gymnasium
Requires-Dist: gymnasium>=1.0; extra == 'gymnasium'
Provides-Extra: jax
Requires-Dist: jax>=0.4.30; extra == 'jax'
Provides-Extra: torch
Requires-Dist: torch>=2.4; extra == 'torch'
Description-Content-Type: text/markdown

# ForkScope

a tool for finding the optimum hyperparameters for world model rollouts at inference.

works for setups that use test-time, action-conditioned planning rollouts (planner generates a bunch
of actions, world model scores those actions based on outcome, planner evaluates and selects the best
one to execute).

## Runnable example

The [point-mass example](examples/point_mass/README.md) includes a small environment, a planner,
zero-argument factories for the CLI, a reference YAML config, and a programmatic example. Install
ForkScope in editable mode from the repository root with `python -m pip install -e .`, then run:

```sh
PYTHONPATH=. forkscope run \
  --env-factory examples.point_mass.example:make_environment \
  --planner-factory examples.point_mass.example:make_planner \
  --seeds 0,1,2 \
  --horizons 1,2,4 \
  --max-steps 10 \
  --output outputs/point_mass/results.jsonl
```

The same factories can be passed to `RolloutBenchmarker` from Python. See the example README for
that usage and details about the JSON Lines output. The YAML file documents the matching settings;
the CLI currently accepts options directly and does not load YAML configs.

## Array backends

The Python API uses NumPy arrays by default. To use JAX or PyTorch, install the corresponding extra
and select the backend when creating a benchmark:

```sh
python -m pip install -e '.[jax]'
python -m pip install -e '.[torch]'
```

```python
from forkscope import RolloutBenchmarker

benchmarker = RolloutBenchmarker(
    env_factory,
    planner_factory,
    seeds=[0, 1],
    horizons=[1, 2],
    max_steps=10,
    backend="jax",  # or "torch"; omit for NumPy
)
records = benchmarker.run()
benchmarker.write_results("results.jsonl")
```

The environment and planner must produce arrays from the selected framework. For example, a JAX
environment can return `{"position": jax.numpy.asarray([position])}` from `reset` and `step`,
and its planner can return a `jax.numpy.asarray(action)` action. The equivalent PyTorch values are
`{"position": torch.tensor([position])}` and `torch.tensor(action)`. Python scalar actions remain
supported for environments that use them. ForkScope rejects arrays from a different backend at
the environment and planner boundaries; it does not convert them during a run.

`run_episode`, `run_sweep`, and the tree utilities also accept `backend="jax"` or
`backend="torch"`. Tree operations retain the framework's array type and device. JSON Lines output
converts arrays to host values when writing; PyTorch tensors are detached first. The CLI currently
uses the default NumPy backend.
