117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
"""Example: run DiscoRL agent inference on a Gym env (PoC).
|
|
|
|
This script demonstrates how to:
|
|
- discretize a continuous Gym action space (quick PoC),
|
|
- adapt the Gym env to DiscoRL's Environment interface using
|
|
`disco_gym_adapter.GymToDiscoEnv`, and
|
|
- run a few actor steps with the DiscoRL Agent (inference only).
|
|
|
|
Important constraints and notes:
|
|
- The provided DiscoRL code expects a scalar discrete action space by
|
|
design (see Agent.__init__). To avoid large code changes we discretize
|
|
your continuous action space into a small action set.
|
|
- This example performs inference only using the freshly initialised
|
|
agent parameters (no meta-training). If you want to run meta-training or
|
|
reproduce the paper, keep using the JAX harness and the original
|
|
training notebooks / scripts.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
# ============================================================================
|
|
# JAX/CUDA Configuration: Prevent GPU conflicts between JAX and LBM simulator.
|
|
# If you encounter segmentation faults, try setting JAX_PLATFORM_NAME=cpu
|
|
# in your shell environment before running this script.
|
|
# ============================================================================
|
|
# Uncomment below to force JAX to CPU-only mode (safest for PoC):
|
|
os.environ['JAX_PLATFORMS'] = 'cpu'
|
|
|
|
import time
|
|
import numpy as np
|
|
import jax
|
|
|
|
# Make sure DiscoRL package is importable. If you installed it with
|
|
# `pip install -e ./disco_rl` then the plain import below will work. If not,
|
|
# try adding the repo-local path to sys.path.
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
repo_root = os.path.abspath(os.path.join(current_dir, os.pardir))
|
|
sys.path.append(os.path.join(repo_root, "disco_rl"))
|
|
|
|
from disco_gym_adapter import GymToDiscoEnv
|
|
from gym_env_250326_erase import CustomEnv
|
|
|
|
from disco_rl import agent as disco_agent
|
|
|
|
|
|
def main():
|
|
# Choose a small discrete action set (example: grid over each continuous dim).
|
|
# Here the original env has action shape (3,) with each element in [-1,1].
|
|
# We build a simple 5-action set: no-op + ±1 on each axis (sparse).
|
|
discrete_actions = np.array([
|
|
[0.0, 0.0, 0.0],
|
|
[1.0, 0.0, 0.0],
|
|
[-1.0, 0.0, 0.0],
|
|
[0.0, 1.0, 0.0],
|
|
[0.0, -1.0, 0.0],
|
|
], dtype=np.float32)
|
|
|
|
# Build the Disco adapter with discrete actions (batch_size=1).
|
|
base_kwargs = dict(device_id=2)
|
|
disco_env = GymToDiscoEnv(
|
|
gym_env_cls=lambda **kw: CustomEnv(**base_kwargs),
|
|
batch_size=1,
|
|
env_settings=None,
|
|
discrete_actions=discrete_actions,
|
|
)
|
|
|
|
# Create agent and settings.
|
|
settings = disco_agent.get_settings_disco()
|
|
|
|
# Get action/observation specs from the adapted env.
|
|
single_obs_spec = disco_env.single_observation_spec()
|
|
single_act_spec = disco_env.single_action_spec()
|
|
|
|
print('Action spec:', single_act_spec)
|
|
print('Observation spec:', single_obs_spec)
|
|
|
|
# Create the DiscoRL agent.
|
|
agent = disco_agent.Agent(
|
|
single_observation_spec=single_obs_spec,
|
|
single_action_spec=single_act_spec,
|
|
agent_settings=settings,
|
|
batch_axis_name=None,
|
|
)
|
|
|
|
rng = jax.random.PRNGKey(0)
|
|
learner_state = agent.initial_learner_state(rng)
|
|
actor_state = agent.initial_actor_state(rng)
|
|
|
|
# Reset environment and run a small rollout.
|
|
_, env_t = disco_env.reset()
|
|
print('\nInitial observation shape:', env_t.observation['observation'].shape)
|
|
print('Initial observation:', env_t.observation)
|
|
|
|
print('\nRunning 10 inference steps...')
|
|
for step in range(10):
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_timestep, actor_state = agent.actor_step(
|
|
learner_state.params, subkey, env_t, actor_state
|
|
)
|
|
# actor_timestep.actions is a JAX array (discrete indices). Convert to numpy.
|
|
action = np.asarray(actor_timestep.actions)[0]
|
|
# Step the Gym env via the adapter: pass the discrete action index.
|
|
_, env_t = disco_env.step(None, np.array([action]))
|
|
reward = np.asarray(env_t.reward)
|
|
print(f' step {step:2d} action={int(action)} reward={reward[0] if hasattr(reward, "__len__") else reward:.4f}')
|
|
|
|
print('\nPoC completed successfully!')
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
main()
|
|
except Exception as e:
|
|
print(f'Error during execution: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|