code/utils/generate_manipspace.py

'''
Adapted from OGBench
https://github.com/seohongpark/ogbench/blob/master/data_gen_scripts/generate_manipspace.py
'''
import pathlib
import os
import sys
import pickle

os.environ["MUJOCO_GL"] = "egl"

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import gymnasium
import numpy as np
from absl import app, flags
from tqdm import trange
import torch
from gymnasium.envs.registration import registry, register
import ogbench.manipspace  # noqa
from ogbench.manipspace.oracles.markov.button_markov import ButtonMarkovOracle
from ogbench.manipspace.oracles.markov.cube_markov import CubeMarkovOracle
from ogbench.manipspace.oracles.markov.drawer_markov import DrawerMarkovOracle
from ogbench.manipspace.oracles.markov.window_markov import WindowMarkovOracle
from ogbench.manipspace.oracles.plan.button_plan import ButtonPlanOracle
from ogbench.manipspace.oracles.plan.cube_plan import CubePlanOracle
from ogbench.manipspace.oracles.plan.drawer_plan import DrawerPlanOracle
from ogbench.manipspace.oracles.plan.window_plan import WindowPlanOracle
from envs.cube.cube_env import CubeEnv

FLAGS = flags.FLAGS

flags.DEFINE_string('env_name', 'cube-single-v1', 'Environment name.')
flags.DEFINE_string('dataset_type', 'play', 'Dataset type.')
flags.DEFINE_string('save_path', None, 'Directory to save outputs (will be created).')
flags.DEFINE_float('noise', 0.1, 'Action noise level.')
flags.DEFINE_float('noise_smoothing', 0.5, 'Action noise smoothing level for PlanOracle.')
flags.DEFINE_float('min_norm', 0.4, 'Minimum action norm for MarkovOracle.')
flags.DEFINE_float('p_random_action', 0, 'Probability of selecting a random action.')
flags.DEFINE_integer('num_episodes', 1000, 'Number of train episodes.')
flags.DEFINE_integer('max_episode_steps', None, 'Max steps per episode.')
flags.DEFINE_integer('init_cube', None, 'Force initial target cube index.')


def main(_):
    assert FLAGS.dataset_type in ['play', 'noisy']
    if FLAGS.save_path is None:
        raise ValueError("Please provide --save_path (directory).")

    save_folder = pathlib.Path(FLAGS.save_path)
    save_folder.mkdir(parents=True, exist_ok=True)
    obses_folder = save_folder / "obses"
    obses_folder.mkdir(parents=True, exist_ok=True)

    # Environment
    env_kwargs = {
        "terminate_at_goal": True,
        "mode": "task",
        "reward_task_id": 5,  # lock to task 5
    }

    if FLAGS.max_episode_steps:
        env_kwargs["max_episode_steps"] = FLAGS.max_episode_steps

    env = gymnasium.make(
        FLAGS.env_name,
        **env_kwargs
    )

    # Oracles initialization
    oracle_type = 'plan' if FLAGS.dataset_type == 'play' else 'markov'
    has_button_states = hasattr(env.unwrapped, '_cur_button_states')
    if 'cube' in FLAGS.env_name:
        if oracle_type == 'markov':
            agents = {'cube': CubeMarkovOracle(env=env, min_norm=FLAGS.min_norm)}
        else:
            agents = {'cube': CubePlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing)}
    elif 'scene' in FLAGS.env_name:
        if oracle_type == 'markov':
            agents = {
                'cube': CubeMarkovOracle(env=env, min_norm=FLAGS.min_norm, max_step=100),
                'button': ButtonMarkovOracle(env=env, min_norm=FLAGS.min_norm),
                'drawer': DrawerMarkovOracle(env=env, min_norm=FLAGS.min_norm),
                'window': WindowMarkovOracle(env=env, min_norm=FLAGS.min_norm),
            }
        else:
            agents = {
                'cube': CubePlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
                'button': ButtonPlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
                'drawer': DrawerPlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
                'window': WindowPlanOracle(env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing),
            }
    elif 'puzzle' in FLAGS.env_name:
        if oracle_type == 'markov':
            agents = {'button': ButtonMarkovOracle(env=env, min_norm=FLAGS.min_norm, gripper_always_closed=True)}
        else:
            agents = {
                'button': ButtonPlanOracle(
                    env=env, noise=FLAGS.noise, noise_smoothing=FLAGS.noise_smoothing, gripper_always_closed=True
                )
            }
    else:
        raise ValueError(f"Unknown env_name: {FLAGS.env_name}")

    # Counters and accumulators
    successful_episodes = 0
    total_steps = 0
    total_train_steps = 0
    num_train_episodes = FLAGS.num_episodes
    num_val_episodes = FLAGS.num_episodes // 10
    target_episodes = num_train_episodes + num_val_episodes

    seq_lengths = []
    episode_latents = []   # list of (Ti, latent_dim)
    episode_actions = []   # list of (Ti, action_dim)
    episode_index = 0

    # Collection loop with retry logic (keeps original behavior)
    while successful_episodes < target_episodes:
        while True:  # retry loop
            if FLAGS.init_cube is not None and 'cube' in FLAGS.env_name:
                env.unwrapped._target_block = FLAGS.init_cube
            ob, info = env.reset(options={'task_id':5})
            if oracle_type == 'markov':
                xi = np.random.uniform(0, FLAGS.noise)
            agent = agents[info['privileged/target_task']]
            agent.reset(ob, info)

            done = False
            step = 0
            ep_qpos = []

            # Per-episode buffers
            ep_pixels = []   # store only pixels per timestep
            ep_latents = []  # store latents per timestep
            ep_actions = []

            while not done:
                if np.random.rand() < FLAGS.p_random_action:
                    # Sample a random action.
                    action = env.action_space.sample()
                else:
                    # Get an action from the oracle.
                    action = agent.select_action(ob, info)
                    action = np.array(action)
                    if oracle_type == 'markov':
                        # Add Gaussian noise to the action.
                        action += np.random.normal(0, [xi, xi, xi, xi*3, xi*10], action.shape)
                action = np.clip(action, -1, 1)

                next_ob, reward, terminated, truncated, info = env.step(action)
                done = terminated or truncated

                # Validate observation format
                if not (isinstance(ob, dict) and 'pixels' in ob and 'latent' in ob):
                    raise KeyError("Expected observation dict with keys 'pixels' and 'latent'")

                # Append only pixels to the episode file buffer
                pixels = ob['pixels']
                ep_pixels.append(pixels)

                # Append latents and actions to global lists
                lat = ob['latent']

                ep_latents.append(np.asarray(lat))
                ep_actions.append(np.asarray(action.copy()))

                ep_qpos.append(info['prev_qpos'])

                # Handle subtask completion
                # Handle subtask completion: only switch if the episode hasn't terminated
                # AND there are still unsatisfied blocks to switch to.
                if agent.done and not done:
                    successes = env.unwrapped._compute_successes()
                    num_satisfied = int(sum(bool(s) for s in successes))

                    if num_satisfied < int(env.unwrapped._num_cubes):
                        agent_ob, agent_info = env.unwrapped.switch_target_block()
                        agent = agents[agent_info['privileged/target_task']]
                        agent.reset(agent_ob, agent_info)

                ob = next_ob
                step += 1

            # Scene health check
            if 'scene' in FLAGS.env_name:
                ep_qpos_arr = np.array(ep_qpos)
                block_xyzs = ep_qpos_arr[:, 14:17]
                is_healthy = not ((block_xyzs[:,1] >= 0.29).any() or
                                ((block_xyzs[:,1] <= -0.3) & ((block_xyzs[:,2]<0.06) | (block_xyzs[:,2]>0.08))).any())
                if not is_healthy:
                    print("Unhealthy episode — retrying (no data saved).", flush=True)
                    continue
            # If episode ended because of time limit (truncated) AND the final reward is negative,
            # drop this episode (do not save it and do not increment successful_episodes).
            # `reward` holds the last reward from the loop above.
            if truncated and (reward < 0):
                print(f"Episode ended by time limit with reward {reward:.3f} < 0 — dropping episode and retrying.", flush=True)
                # Do NOT increment successful_episodes or episode_index; restart retry loop
                continue


            # Episode successful -> save per-episode pixels and accumulate latents/actions
            successful_episodes += 1
            total_steps += step
            if successful_episodes <= num_train_episodes:
                total_train_steps += step
                print(f"Train Episode {successful_episodes}/{target_episodes} Steps: {step}", flush=True)
            else:
                print(f"Val Episode {successful_episodes}/{target_episodes} Steps: {step}", flush=True)

            # Save episode pixels only
            ep_file = obses_folder / f"episode_{episode_index:05d}.pth"
            # Convert lists to arrays when sensible (keep dtype uint8 for images if available)
            try:
                # try stack if all same shape
                ep_pixels_arr = np.stack(ep_pixels, axis=0)
            except Exception:
                # fallback to list (torch.save will serialize)
                ep_pixels_arr = ep_pixels
            torch.save(ep_pixels_arr, str(ep_file))

            # Record sequence length
            seq_lengths.append(len(ep_pixels))

            # Append latents and actions to global lists (maintain time order)
            episode_latents.append(
                np.stack(ep_latents, axis=0)
            )
            episode_actions.append(
                np.stack(ep_actions, axis=0)
            )
            episode_index += 1
            break  # exit retry loop and move to next episode

    print("Collected total steps:", total_steps)

    # Concatenate latents and actions into tensors and save
    num_eps = len(seq_lengths)
    max_len = max(seq_lengths)

    latent_dim = episode_latents[0].shape[1]
    action_dim = episode_actions[0].shape[1]

    latents_ep = torch.zeros((num_eps, max_len, latent_dim), dtype=torch.float32)
    actions_ep = torch.zeros((num_eps, max_len, action_dim), dtype=torch.float32)

    for i, (lat_ep, act_ep) in enumerate(zip(episode_latents, episode_actions)):
        L = lat_ep.shape[0]
        latents_ep[i, :L] = torch.from_numpy(lat_ep)
        actions_ep[i, :L] = torch.from_numpy(act_ep)

    torch.save(latents_ep, save_folder / "latents.pth")
    torch.save(actions_ep, save_folder / "actions.pth")

    with open(save_folder / "seq_lengths.pkl", "wb") as f:
        pickle.dump(seq_lengths, f)

    print("Saved:")
    print(" latents:", latents_ep.shape)
    print(" actions:", actions_ep.shape)



if __name__ == "__main__":
    app.run(main)