62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
"""Debug script to check what actions agent takes."""
|
|
|
|
import sys
|
|
sys.path.insert(0, '/home/frank14f/Frank_LBM/scripts')
|
|
sys.path.insert(0, '/home/frank14f/Frank_LBM')
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
from disco_cartpole_env import DiscoCartPoleEnv
|
|
import disco_rl.agent as disco_agent
|
|
import disco_rl.types as types
|
|
|
|
# Create environment
|
|
env = DiscoCartPoleEnv(batch_size=1, max_steps=500)
|
|
|
|
# Create agent
|
|
agent_settings = disco_agent.get_settings_disco()
|
|
agent = disco_agent.Agent(
|
|
single_observation_spec=env.single_observation_spec(),
|
|
single_action_spec=env.single_action_spec(),
|
|
agent_settings=agent_settings,
|
|
batch_axis_name=None,
|
|
)
|
|
|
|
# Initialize state
|
|
rng = jax.random.PRNGKey(42)
|
|
rng, subkey = jax.random.split(rng)
|
|
learner_state = agent.initial_learner_state(subkey)
|
|
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_state = agent.initial_actor_state(subkey)
|
|
|
|
# Reset environment
|
|
_, env_t = env.reset()
|
|
|
|
print("Manual trajectory collection:")
|
|
print(f"Initial env_t.step_type: {env_t.step_type}")
|
|
print(f"Initial env_t.reward: {env_t.reward}")
|
|
|
|
for step in range(20):
|
|
rng, subkey = jax.random.split(rng)
|
|
|
|
# Get action from agent
|
|
actor_timestep, actor_state = agent.actor_step(
|
|
learner_state.params,
|
|
subkey,
|
|
env_t,
|
|
actor_state,
|
|
)
|
|
|
|
action = actor_timestep.actions[0]
|
|
|
|
# Step environment
|
|
rng, subkey = jax.random.split(rng)
|
|
_, env_t = env.step(None, actor_timestep.actions)
|
|
|
|
print(f"Step {step+1:2d}: action={int(action)}, reward={env_t.reward[0]:.1f}, step_type={env_t.step_type[0]} (0=FIRST, 1=MID, 2=LAST), done={env._episode_done[0]}")
|
|
|
|
if env._episode_done[0]:
|
|
print(f" -> Episode done at step {step+1}!")
|