import numpy as np
from rlbench.environment import Environment
from rlbench.action_modes.action_mode import MoveArmThenGripper
from rlbench.action_modes.arm_action_modes import EndEffectorPoseViaPlanning
from rlbench.action_modes.gripper_action_modes import Discrete
from rlbench.tasks import PutRubbishInBin

def setup_environment():
    np.random.seed(42)
    action_mode = MoveArmThenGripper(
        arm_action_mode=EndEffectorPoseViaPlanning(),
        gripper_action_mode=Discrete()
    )
    env = Environment(action_mode=action_mode, headless=False)
    env.launch()
    task = env.get_task(PutRubbishInBin)
    task.set_variation(0)
    return env, task

def shutdown_environment(env):
    env.shutdown()