Add files using upload-large-folder tool

a3f6fc6 verified about 1 month ago

11.3 kB

'''
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)

	'''
	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, xi3, xi10], 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)